Turn on allocation tracker at run-time and for web (#1591)

* Fix copy/paste in web viewer

* Memory tracker callstack: shorten file paths

* Fix typo

* clean up copy-pasta comment

* code refactor

* allow turning on allocation tracking at runtime from the memory panel

* nicer formatting in memory panel

* Notice text layout allocations in the summaries

* Notice failures to copy the backtrace

* Notice empty backtraces

* Fix stacktraces on web

* wasm_bindgen_check: build to target_wasm

* Always use target_wasm as build-dir when building wasm32

* Trim callstacks on web

* typo fixe

Co-authored-by: Andreas Reich <[email protected]>

* Use a checkbox for the detailed allocation tracking

* Make web callstack even nicer

* Even nicer callstacks

---------

Co-authored-by: Andreas Reich <[email protected]>

.github/workflows/rust.yml

@@ -218,13 +218,13 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: check
-          args: --locked --all-features --lib --target wasm32-unknown-unknown -p re_viewer
+          args: --locked --all-features --lib --target wasm32-unknown-unknown --target-dir target_wasm -p re_viewer
 
       - name: Check re_renderer examples wasm32
         uses: actions-rs/cargo@v1
         with:
           command: check
-          args: --locked --target wasm32-unknown-unknown -p re_renderer --examples
+          args: --locked --target wasm32-unknown-unknown --target-dir target_wasm -p re_renderer --examples
 
       - run: ./scripts/wasm_bindgen_check.sh --skip-setup
 

crates/re_build_web_viewer/src/lib.rs

@@ -56,13 +56,13 @@ pub fn build(release: bool) {
     let mut cmd = std::process::Command::new("cargo");
     cmd.args([
         "build",
-        "--target-dir",
-        target_wasm_dir.as_str(),
-        "-p",
+        "--package",
         crate_name,
         "--lib",
         "--target",
         "wasm32-unknown-unknown",
+        "--target-dir",
+        target_wasm_dir.as_str(),
     ]);
     if release {
         cmd.arg("--release");
@@ -78,7 +78,7 @@ pub fn build(release: bool) {
     // These pre-encoded flags are generally generated by Cargo itself when loading its
     // configuration from e.g. `$CARGO_HOME/config.toml`; which means they will contain
     // values that only make sense for the native target host, not for a wasm build.
-    cmd.env("CARGO_ENCODED_RUSTFLAGS", "");
+    cmd.env("CARGO_ENCODED_RUSTFLAGS", "--cfg=web_sys_unstable_apis");
 
     eprintln!("> {cmd:?}");
     let status = cmd

crates/re_memory/Cargo.toml

@@ -21,14 +21,19 @@ re_format.workspace = true
 re_log.workspace = true
 
 ahash.workspace = true
-backtrace = { version = "0.3" }
 emath.workspace = true
 instant = { version = "0.1", features = ["wasm-bindgen"] }
+itertools = "0.10"
 nohash-hasher = "0.2"
 once_cell = "1.16"
 parking_lot.workspace = true
 smallvec = "1.10"
 
 # native dependencies:
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+backtrace = "0.3"
 memory-stats = "1.0"
+
+# web dependencies:
+[target.'cfg(target_arch = "wasm32")'.dependencies]
+wasm-bindgen = "=0.2.84"

crates/re_memory/src/accounting_allocator.rs

@@ -31,7 +31,7 @@ thread_local! {
     ///
     /// Tracking an allocation (taking its backtrace etc) can itself create allocations.
     /// We don't want to track those allocations, or we will have infinite recursion.
-    static IS_TRHEAD_IN_ALLOCATION_TRACKER: std::cell::Cell<bool> = std::cell::Cell::new(false);
+    static IS_THREAD_IN_ALLOCATION_TRACKER: std::cell::Cell<bool> = std::cell::Cell::new(false);
 }
 
 // ----------------------------------------------------------------------------
@@ -188,7 +188,7 @@ pub fn tracking_stats() -> Option<TrackingStatistics> {
     }
 
     GLOBAL_STATS.track_callstacks.load(Relaxed).then(|| {
-        IS_TRHEAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
+        IS_THREAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
             // prevent double-lock of ALLOCATION_TRACKER:
             is_thread_in_allocation_tracker.set(true);
             let mut top_big_callstacks = tracker_stats(&BIG_ALLOCATION_TRACKER.lock());
@@ -303,7 +303,7 @@ fn note_alloc(ptr: *mut u8, size: usize) {
             // Big enough to track - but make sure we don't create a deadlock by trying to
             // track the allocations made by the allocation tracker:
 
-            IS_TRHEAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
+            IS_THREAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
                 if !is_thread_in_allocation_tracker.get() {
                     is_thread_in_allocation_tracker.set(true);
 
@@ -337,7 +337,7 @@ fn note_dealloc(ptr: *mut u8, size: usize) {
         } else {
             // Big enough to track - but make sure we don't create a deadlock by trying to
             // track the allocations made by the allocation tracker:
-            IS_TRHEAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
+            IS_THREAD_IN_ALLOCATION_TRACKER.with(|is_thread_in_allocation_tracker| {
                 if !is_thread_in_allocation_tracker.get() {
                     is_thread_in_allocation_tracker.set(true);
 

crates/re_memory/src/allocation_tracker.rs

@@ -1,6 +1,6 @@
-use std::{hash::Hash, sync::Arc};
+use std::sync::Arc;
 
-use backtrace::Backtrace;
+use crate::{Backtrace, BacktraceHash};
 
 use crate::CountAndSize;
 
@@ -22,25 +22,6 @@ impl PtrHash {
 
 // ----------------------------------------------------------------------------
 
-#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
-struct BacktraceHash(u64);
-
-impl nohash_hasher::IsEnabled for BacktraceHash {}
-
-fn hash_backtrace(backtrace: &Backtrace) -> BacktraceHash {
-    use std::hash::Hasher as _;
-    let mut hasher =
-        std::hash::BuildHasher::build_hasher(&ahash::RandomState::with_seeds(0, 1, 2, 3));
-
-    for frame in backtrace.frames() {
-        frame.ip().hash(&mut hasher);
-    }
-
-    BacktraceHash(hasher.finish())
-}
-
-// ----------------------------------------------------------------------------
-
 /// Formatted [`Backtrace`].
 ///
 /// Clones without allocating.
@@ -58,26 +39,10 @@ impl std::fmt::Display for ReadableBacktrace {
 
 impl ReadableBacktrace {
     fn new(mut backtrace: Backtrace) -> Self {
-        backtrace.resolve();
-        let readable = format_backtrace(&backtrace);
-        Self { readable }
-    }
-}
-
-fn format_backtrace(backtrace: &Backtrace) -> Arc<str> {
-    let stack = format!("{backtrace:?}");
-    let mut stack = stack.as_str();
-    let start_pattern = "re_memory::accounting_allocator::note_alloc\n";
-    if let Some(start_offset) = stack.find(start_pattern) {
-        stack = &stack[start_offset + start_pattern.len()..];
-    }
-
-    if let Some(end_offset) = stack.find("std::sys_common::backtrace::__rust_begin_short_backtrace")
-    {
-        stack = &stack[..end_offset];
+        Self {
+            readable: backtrace.format(),
+        }
     }
-
-    stack.into()
 }
 
 // ----------------------------------------------------------------------------
@@ -137,7 +102,7 @@ impl AllocationTracker {
         }
 
         let unresolved_backtrace = Backtrace::new_unresolved();
-        let hash = hash_backtrace(&unresolved_backtrace);
+        let hash = BacktraceHash::new(&unresolved_backtrace);
 
         self.readable_backtraces
             .entry(hash)

crates/re_memory/src/backtrace_native.rs

@@ -0,0 +1,120 @@
+use std::sync::Arc;
+
+pub(crate) struct Backtrace(backtrace::Backtrace);
+
+impl Backtrace {
+    pub fn new_unresolved() -> Self {
+        Self(backtrace::Backtrace::new_unresolved())
+    }
+
+    pub fn format(&mut self) -> Arc<str> {
+        self.0.resolve();
+        let stack = backtrace_to_string(&self.0);
+        trim_backtrace(&stack).into()
+    }
+}
+
+impl std::hash::Hash for Backtrace {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        for frame in self.0.frames() {
+            frame.ip().hash(state);
+        }
+    }
+}
+
+fn trim_backtrace(mut stack: &str) -> &str {
+    let start_pattern = "re_memory::accounting_allocator::note_alloc\n";
+    if let Some(start_offset) = stack.find(start_pattern) {
+        stack = &stack[start_offset + start_pattern.len()..];
+    }
+
+    let end_pattern = "std::sys_common::backtrace::__rust_begin_short_backtrace";
+    if let Some(end_offset) = stack.find(end_pattern) {
+        stack = &stack[..end_offset];
+    }
+
+    stack
+}
+
+fn backtrace_to_string(backtrace: &backtrace::Backtrace) -> String {
+    if backtrace.frames().is_empty() {
+        return "[empty backtrace]".to_owned();
+    }
+
+    // We need to get a `std::fmt::Formatter`, and there is no easy way to do that, so we do it the hard way:
+
+    struct AnonymizedBacktrace<'a>(&'a backtrace::Backtrace);
+
+    impl<'a> std::fmt::Display for AnonymizedBacktrace<'a> {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            format_backtrace_with_fmt(self.0, f)
+        }
+    }
+
+    AnonymizedBacktrace(backtrace).to_string()
+}
+
+fn format_backtrace_with_fmt(
+    backtrace: &backtrace::Backtrace,
+    fmt: &mut std::fmt::Formatter<'_>,
+) -> std::fmt::Result {
+    let mut print_path = |fmt: &mut std::fmt::Formatter<'_>,
+                          path: backtrace::BytesOrWideString<'_>| {
+        let path = path.into_path_buf();
+        let shortened = shorten_source_file_path(&path);
+        std::fmt::Display::fmt(&shortened, fmt)
+    };
+
+    let style = if fmt.alternate() {
+        backtrace::PrintFmt::Full
+    } else {
+        backtrace::PrintFmt::Short
+    };
+    let mut f = backtrace::BacktraceFmt::new(fmt, style, &mut print_path);
+    f.add_context()?;
+    for frame in backtrace.frames() {
+        f.frame().backtrace_frame(frame)?;
+    }
+    f.finish()?;
+    Ok(())
+}
+
+/// Anonymize a path to a Rust source file from a callstack.
+///
+/// Example input:
+/// * `/Users/emilk/.cargo/registry/src/jackfan.us.kg-1ecc6299db9ec823/tokio-1.24.1/src/runtime/runtime.rs`
+/// * `crates/rerun/src/main.rs`
+/// * `/rustc/d5a82bbd26e1ad8b7401f6a718a9c57c96905483/library/core/src/ops/function.rs`
+fn shorten_source_file_path(path: &std::path::Path) -> String {
+    // We must make sure we strip everything sensitive (especially user name).
+    // The easiest way is to look for `src` and strip everything up to it.
+
+    use itertools::Itertools as _;
+    let components = path.iter().map(|path| path.to_string_lossy()).collect_vec();
+
+    // Look for the last `src`:
+    if let Some((src_rev_idx, _)) = components.iter().rev().find_position(|&c| c == "src") {
+        let src_idx = components.len() - src_rev_idx - 1;
+        // Before `src` comes the name of the crate - let's include that:
+        let first_index = src_idx.saturating_sub(1);
+        components.iter().skip(first_index).format("/").to_string()
+    } else {
+        // No `src` directory found - weird!
+        path.display().to_string()
+    }
+}
+
+#[test]
+fn test_shorten_path() {
+    for (before, after) in [
+        ("/Users/emilk/.cargo/registry/src/jackfan.us.kg-1ecc6299db9ec823/tokio-1.24.1/src/runtime/runtime.rs", "tokio-1.24.1/src/runtime/runtime.rs"),
+        ("crates/rerun/src/main.rs", "rerun/src/main.rs"),
+        ("/rustc/d5a82bbd26e1ad8b7401f6a718a9c57c96905483/library/core/src/ops/function.rs", "core/src/ops/function.rs"),
+        ("/weird/path/file.rs", "/weird/path/file.rs"),
+        ]
+        {
+        use std::str::FromStr as _;
+        let before = std::path::PathBuf::from_str(before).unwrap();
+        assert_eq!(shorten_source_file_path(&before), after);
+    }
+}

crates/re_memory/src/backtrace_web.rs

@@ -0,0 +1,63 @@
+use wasm_bindgen::prelude::*;
+
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(js_namespace = console)]
+    fn error(msg: String);
+
+    type Error;
+
+    #[wasm_bindgen(constructor)]
+    fn new() -> Error;
+
+    #[wasm_bindgen(structural, method, getter)]
+    fn stack(error: &Error) -> String;
+}
+
+#[derive(Hash)]
+pub(crate) struct Backtrace(String);
+
+impl Backtrace {
+    pub fn new_unresolved() -> Self {
+        Self(Error::new().stack())
+    }
+
+    pub fn format(&mut self) -> std::sync::Arc<str> {
+        trim_backtrace(&self.0).into()
+    }
+}
+
+fn trim_backtrace(mut stack: &str) -> String {
+    let start_pattern = "__rust_alloc_zeroed";
+    if let Some(start_offset) = stack.find(start_pattern) {
+        if let Some(next_newline) = stack[start_offset..].find('\n') {
+            stack = &stack[start_offset + next_newline + 1..];
+        }
+    }
+
+    let end_pattern = "paint_and_schedule"; // normal eframe entry-point
+    if let Some(end_offset) = stack.find(end_pattern) {
+        if let Some(next_newline) = stack[end_offset..].find('\n') {
+            stack = &stack[..end_offset + next_newline];
+        }
+    }
+
+    stack.split('\n').map(trim_line).collect::<String>()
+}
+
+/// Example inputs:
+/// * `eframe::web::backend::AppRunner::paint::h584aff3234354fd5@http://127.0.0.1:9090/re_viewer.js line 366 > WebAssembly.instantiate:wasm-function[3352]:0x5d46b4`
+/// * `getImports/imports.wbg.__wbg_new_83e4891414f9e5c1/<@http://127.0.0.1:9090/re_viewer.js:453:21`
+/// * `__rg_realloc@http://127.0.0.1:9090/re_viewer.js line 366 > WebAssembly.instantiate:wasm-function[17996]:0x9b935f`
+fn trim_line(mut line: &str) -> String {
+    if let Some(index) = line.rfind("::") {
+        line = &line[..index];
+    }
+    if let Some(index) = line.find("/imports.wbg") {
+        line = &line[..index];
+    }
+    if let Some(index) = line.find("@http:") {
+        line = &line[..index];
+    }
+    format!("{line}\n")
+}