Add a test of memory use when logging a lot of big images (#1372)

* Add a test of memory use when logging a lot of big images

* copy-pasta

Co-authored-by: Clement Rey <[email protected]>

* copy-pasta

Co-authored-by: Clement Rey <[email protected]>

* Add store archiving to the stack shortening

* fix feature name

---------

Co-authored-by: Clement Rey <[email protected]>

Cargo.toml

@@ -1,6 +1,12 @@
 [workspace]
 resolver = "2"
-members = ["crates/*", "examples/rust/*", "rerun_py", "run_wasm"]
+members = [
+  "crates/*",
+  "examples/rust/*",
+  "rerun_py",
+  "run_wasm",
+  "tests/rust/*",
+]
 
 [workspace.package]
 authors = ["rerun.io <[email protected]>"]
@@ -55,6 +61,7 @@ half = "2.0"
 image = "0.24"
 lazy_static = "1.4"
 macaw = "0.18"
+mimalloc = "0.1.29"
 ndarray = "0.15"
 polars-core = "0.27.1"
 polars-lazy = "0.27.1"

crates/re_arrow_store/Cargo.toml

@@ -71,7 +71,7 @@ polars-ops = { workspace = true, optional = true, features = [
 
 [dev-dependencies]
 criterion = "0.4"
-mimalloc = "0.1"
+mimalloc.workspace = true
 polars-core = { workspace = true, features = [
   "diagonal_concat",
   "dtype-date",

crates/re_data_store/Cargo.toml

@@ -46,7 +46,7 @@ puffin.workspace = true
 
 [dev-dependencies]
 criterion = "0.4"
-mimalloc = "0.1"
+mimalloc.workspace = true
 rand = "0.8"
 
 [lib]

crates/re_int_histogram/Cargo.toml

@@ -24,7 +24,7 @@ static_assertions = "1.1"
 [dev-dependencies]
 criterion = "0.4"
 insta = "1.23"
-mimalloc = "0.1"
+mimalloc.workspace = true
 
 
 [lib]

crates/re_log_types/Cargo.toml

@@ -104,7 +104,7 @@ ruzstd = { version = "0.3.0", optional = true } # works on wasm
 
 [dev-dependencies]
 criterion = "0.4"
-mimalloc = "0.1"
+mimalloc.workspace = true
 serde_test = { version = "1" }
 arrow2 = { workspace = true, features = [
   "io_ipc",

crates/re_query/Cargo.toml

@@ -55,7 +55,7 @@ polars-core = { workspace = true, optional = true, features = [
 [dev-dependencies]
 criterion = "0.4"
 itertools = "0.10"
-mimalloc = "0.1"
+mimalloc.workspace = true
 polars-core = { workspace = true, features = [
   "dtype-date",
   "dtype-time",

crates/re_sdk/src/lib.rs

@@ -63,6 +63,9 @@ pub mod external {
     pub use re_memory;
     pub use re_sdk_comms;
 
+    #[cfg(feature = "re_viewer")]
+    pub use re_viewer;
+
     #[cfg(feature = "glam")]
     pub use re_log_types::external::glam;
 

crates/re_viewer/src/ui/memory_panel.rs

@@ -357,11 +357,16 @@ impl MemoryPanel {
 
 fn summarize_callstack(callstack: &str) -> String {
     let patterns = [
+        ("MsgSender", "MsgSender"),
+        ("App::receive_messages", "App::receive_messages"),
+        ("w_store::store::ComponentBucket>::archive", "archive"),
+        ("DataStore>::insert", "DataStore"),
         ("LogDb", "LogDb"),
         ("EntityDb", "EntityDb"),
         ("EntityTree", "EntityTree"),
         ("::LogMsg>::deserialize", "LogMsg"),
         ("::TimePoint>::deserialize", "TimePoint"),
+        ("ImageCache", "ImageCache"),
         ("gltf", "gltf"),
         ("image::image", "image"),
         // -----

crates/rerun/Cargo.toml

@@ -78,7 +78,7 @@ webbrowser = { version = "0.8", optional = true }
 # Native dependencies:
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 clap = { workspace = true, features = ["derive"] }
-mimalloc = "0.1.29"
+mimalloc.workspace = true
 puffin_http = "0.11"
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 

crates/rerun/src/run.rs

@@ -5,6 +5,8 @@ use re_smart_channel::Receiver;
 use anyhow::Context as _;
 use clap::Subcommand;
 
+// Note the extra blank lines between the point-lists below: it is required by `clap`.
+
 /// The Rerun Viewer and Server
 ///
 /// Features:

examples/rust/raw_mesh/Cargo.toml

@@ -13,5 +13,5 @@ anyhow.workspace = true
 bytes = "1.3"
 clap = { workspace = true, features = ["derive"] }
 gltf.workspace = true
-mimalloc = "0.1"
+mimalloc.workspace = true
 reqwest = { workspace = true, features = ["blocking", "rustls-tls"] }

rerun_py/Cargo.toml

@@ -59,7 +59,7 @@ image = { version = "0.24", default-features = false, features = [
 ] }
 itertools = "0.10"
 macaw.workspace = true
-mimalloc = { version = "0.1.29", features = ["local_dynamic_tls"] }
+mimalloc = { workspace = true, features = ["local_dynamic_tls"] }
 numpy = { version = "0.18.0", features = ["half"] }
 pyo3 = { version = "0.18.0", features = ["abi3-py38"] }
 rand = { version = "0.8", features = ["std_rng"] }

tests/README.md

@@ -0,0 +1 @@
+Uses of the rerun SDK designed to test different things.

tests/rust/test_image_memory/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "test_image_memory"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+publish = false
+
+[dependencies]
+rerun.workspace = true
+re_format.workspace = true
+
+mimalloc.workspace = true

tests/rust/test_image_memory/src/main.rs

@@ -0,0 +1,56 @@
+//! Logs a bunch of big images to test Rerun memory usage.
+
+use mimalloc::MiMalloc;
+use re_memory::AccountingAllocator;
+use rerun::external::{image, re_memory, re_viewer};
+
+#[global_allocator]
+static GLOBAL: AccountingAllocator<MiMalloc> = AccountingAllocator::new(MiMalloc);
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    re_memory::accounting_allocator::turn_on_tracking_if_env_var(
+        re_viewer::env_vars::RERUN_TRACK_ALLOCATIONS,
+    );
+
+    let session = rerun::Session::init("test_image_memory_rs", true);
+
+    session.spawn(|mut session| {
+        log_images(&mut session).unwrap();
+    })?;
+
+    Ok(())
+}
+
+fn log_images(session: &mut rerun::Session) -> Result<(), Box<dyn std::error::Error>> {
+    let (w, h) = (2048, 1024);
+    let n = 100;
+
+    let image = image::RgbaImage::from_fn(w, h, |x, y| {
+        if (x + y) % 2 == 0 {
+            image::Rgba([0, 0, 0, 255])
+        } else {
+            image::Rgba([255, 255, 255, 255])
+        }
+    });
+    let tensor = rerun::components::Tensor::from_image(image)?;
+
+    for _ in 0..n {
+        rerun::MsgSender::new("image")
+            .with_component(&[tensor.clone()])?
+            .send(session)?;
+    }
+
+    eprintln!(
+        "Logged {n} {w}x{h} RGBA images = {}",
+        re_format::format_bytes((n * w * h * 4) as _)
+    );
+
+    // Give viewer time to load it:
+    std::thread::sleep(std::time::Duration::from_secs(2));
+
+    if let Some(allocs) = re_memory::accounting_allocator::global_allocs() {
+        eprintln!("{} RAM used", re_format::format_bytes(allocs.size as _));
+    }
+
+    Ok(())
+}