From 0b0ffc3a4cafb0e2df8b57552ac09503ced64c27 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Mon, 22 Apr 2024 18:03:04 +0200 Subject: [PATCH 1/3] Add memory overhead test for the datastore Part of https://github.com/rerun-io/rerun/issues/6066 --- Cargo.lock | 1 + crates/re_data_store/Cargo.toml | 2 + crates/re_data_store/tests/memory_test.rs | 107 ++++++++++++++++++ .../memory_test__scalars_on_one_timeline.snap | 10 ++ 4 files changed, 120 insertions(+) create mode 100644 crates/re_data_store/tests/memory_test.rs create mode 100644 crates/re_data_store/tests/snapshots/memory_test__scalars_on_one_timeline.snap diff --git a/Cargo.lock b/Cargo.lock index 134c857a11a1..413fa0bb1c7b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4209,6 +4209,7 @@ dependencies = [ "criterion", "document-features", "indent", + "insta", "itertools 0.12.0", "mimalloc", "nohash-hasher", diff --git a/crates/re_data_store/Cargo.toml b/crates/re_data_store/Cargo.toml index 5356836263dc..0674815bed3b 100644 --- a/crates/re_data_store/Cargo.toml +++ b/crates/re_data_store/Cargo.toml @@ -51,9 +51,11 @@ web-time.workspace = true [dev-dependencies] re_types = { workspace = true, features = ["datagen", "testing"] } +re_format.workspace = true anyhow.workspace = true criterion.workspace = true +insta.workspace = true mimalloc.workspace = true rand.workspace = true similar-asserts.workspace = true diff --git a/crates/re_data_store/tests/memory_test.rs b/crates/re_data_store/tests/memory_test.rs new file mode 100644 index 000000000000..2b9b1f5c409f --- /dev/null +++ b/crates/re_data_store/tests/memory_test.rs @@ -0,0 +1,107 @@ +//! Measures the memory overhead of the data store. + +use std::sync::atomic::{AtomicUsize, Ordering::Relaxed}; + +thread_local! { + static LIVE_BYTES_IN_THREAD: AtomicUsize = AtomicUsize::new(0); +} + +pub struct TrackingAllocator { + allocator: std::alloc::System, +} + +#[global_allocator] +pub static GLOBAL_ALLOCATOR: TrackingAllocator = TrackingAllocator { + allocator: std::alloc::System, +}; + +#[allow(unsafe_code)] +// SAFETY: +// We just do book-keeping and then let another allocator do all the actual work. +unsafe impl std::alloc::GlobalAlloc for TrackingAllocator { + #[allow(clippy::let_and_return)] + unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { + LIVE_BYTES_IN_THREAD.with(|bytes| bytes.fetch_add(layout.size(), Relaxed)); + + // SAFETY: + // Just deferring + unsafe { self.allocator.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) { + LIVE_BYTES_IN_THREAD.with(|bytes| bytes.fetch_sub(layout.size(), Relaxed)); + + // SAFETY: + // Just deferring + unsafe { self.allocator.dealloc(ptr, layout) }; + } +} + +fn live_bytes() -> usize { + LIVE_BYTES_IN_THREAD.with(|bytes| bytes.load(Relaxed)) +} + +/// Assumes all allocations are on the calling thread. +/// +/// The reason we use thread-local counting is so that +/// the counting won't be confused by any other running threads (e.g. other tests). +fn memory_use(run: impl Fn() -> R) -> usize { + let used_bytes_start = live_bytes(); + let ret = run(); + let bytes_used = live_bytes() - used_bytes_start; + drop(ret); + bytes_used +} + +// ---------------------------------------------------------------------------- + +use re_data_store::{DataStore, DataStoreConfig}; +use re_log_types::{DataRow, RowId, TimePoint, TimeType, Timeline}; +use re_types::components::{InstanceKey, Scalar}; +use re_types_core::Loggable as _; + +/// The memory overhead of storing many scalars in the store. +#[test] +fn scalar_memory_overhead() { + re_log::setup_logging(); + + const NUM_SCALARS: usize = 1024 * 1024; + + let total_mem_use = memory_use(|| { + let mut store = DataStore::new( + re_log_types::StoreId::random(re_log_types::StoreKind::Recording), + InstanceKey::name(), + DataStoreConfig::default(), + ); + + for i in 0..NUM_SCALARS { + let entity_path = re_log_types::entity_path!("scalar"); + let timepoint = + TimePoint::default().with(Timeline::new("log_time", TimeType::Time), i as i64); + let num_instances = 1; + let row = DataRow::from_cells1_sized( + RowId::new(), + entity_path, + timepoint, + num_instances, + vec![Scalar(i as f64)], + ) + .unwrap(); + store.insert_row(&row).unwrap(); + } + + store + }); + + insta::assert_debug_snapshot!( + "scalars_on_one_timeline", + [ + format!("{NUM_SCALARS} scalars"), + format!("{} in total", re_format::format_bytes(total_mem_use as _)), + format!( + "{} per row", + re_format::format_bytes(total_mem_use as f64 / NUM_SCALARS as f64) + ), + ] + ); +} diff --git a/crates/re_data_store/tests/snapshots/memory_test__scalars_on_one_timeline.snap b/crates/re_data_store/tests/snapshots/memory_test__scalars_on_one_timeline.snap new file mode 100644 index 000000000000..6ba35e5cf106 --- /dev/null +++ b/crates/re_data_store/tests/snapshots/memory_test__scalars_on_one_timeline.snap @@ -0,0 +1,10 @@ +--- +source: crates/re_data_store/tests/memory_test.rs +assertion_line: 96 +expression: "[format!(\"{NUM_SCALARS} scalars\"),\n format!(\"{} in total\", re_format::format_bytes(total_mem_use as _)),\n format!(\"{} per row\",\n re_format::format_bytes(total_mem_use as f64 / NUM_SCALARS as\n f64))]" +--- +[ + "1048576 scalars", + "936 MiB in total", + "936 B per row", +] From b51a5b6d030c94966f2d65c7a6b983c46aa7c37f Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Mon, 22 Apr 2024 18:27:08 +0200 Subject: [PATCH 2/3] Fix `pixi run rs-update-insta-tests` --- pixi.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pixi.toml b/pixi.toml index be6b61507499..bc76285745df 100644 --- a/pixi.toml +++ b/pixi.toml @@ -186,7 +186,7 @@ meilisearch = "meilisearch --db-path=./meilisearch/data.ms --dump-dir=./meilisea download-design-tokens = "curl --fail https://rerun-docs.netlify.app/api/tokens | jq > crates/re_ui/data/design_tokens.json" # Update the results of `insta` snapshot regression tests -rs-update-insta-tests = "cargo test && cargo insta review" +rs-update-insta-tests = "cargo test ; cargo insta review" # Upload image to gcloud storage. upload-image = "python scripts/upload_image.py" From 6afadd71d22da2bbda41600091638d77aa1618ce Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 23 Apr 2024 09:32:27 +0200 Subject: [PATCH 3/3] sort deps --- crates/re_data_store/Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/re_data_store/Cargo.toml b/crates/re_data_store/Cargo.toml index 0674815bed3b..a81e47ecf21d 100644 --- a/crates/re_data_store/Cargo.toml +++ b/crates/re_data_store/Cargo.toml @@ -30,8 +30,8 @@ deadlock_detection = ["parking_lot/deadlock_detection"] # Rerun dependencies: re_format.workspace = true re_format_arrow.workspace = true -re_log_types.workspace = true re_log = { workspace = true, features = ["setup"] } +re_log_types.workspace = true re_tracing.workspace = true re_types_core.workspace = true @@ -50,8 +50,8 @@ web-time.workspace = true [dev-dependencies] -re_types = { workspace = true, features = ["datagen", "testing"] } re_format.workspace = true +re_types = { workspace = true, features = ["datagen", "testing"] } anyhow.workspace = true criterion.workspace = true