fix(runtime): replace TLS with destructors in allocator hot path

The PianoAllocator accessed STACK (RefCell<Vec<StackEntry>>) from
track_alloc/track_dealloc, which has a destructor. On Rust < 1.93.1
this triggers "fatal runtime error: the global allocator may not use
TLS with destructors, aborting". Rust 1.93.1 relaxed this check
(rust-lang/rust#144465), masking the bug in our test suite.

Replace STACK access with a destructor-free Cell<AllocSnapshot> where
AllocSnapshot is Copy (no destructor). enter() saves and zeroes the
Cell; Guard::drop() reads and restores, correctly handling nesting.
Benchmark shows ~0ns overhead per allocation vs ~25ns before.

Also fix build_instrumented to remove RUSTUP_TOOLCHAIN so target
projects' rust-toolchain.toml is respected (previously, nested cargo
inherited the parent toolchain, silently masking version-specific bugs).

Update integration test to pin Rust 1.91.0 with rayon, representative
of real programs like chainsaw that triggered the crash.

Closes #27

Author: rocketman-code

piano-runtime/src/alloc.rs

@@ -1,18 +1,41 @@
 use std::alloc::{GlobalAlloc, Layout};
-use std::time::Instant;
-
-use crate::collector::STACK;
+use std::cell::Cell;
+
+/// Allocation counters accumulated in the allocator hot path.
+/// All fields are plain integers → `Copy` → `Cell<AllocSnapshot>` has no
+/// destructor → safe for use in global allocator TLS on all Rust versions.
+#[derive(Clone, Copy, Default)]
+pub(crate) struct AllocSnapshot {
+    pub(crate) alloc_count: u32,
+    pub(crate) alloc_bytes: u64,
+    pub(crate) free_count: u32,
+    pub(crate) free_bytes: u64,
+}
 
-// Re-entrancy guard: when tracking code itself allocates, we skip tracking.
 thread_local! {
-    static IN_ALLOC_TRACKING: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
+    /// Destructor-free counters that the allocator hot path increments.
+    /// `enter()` saves and zeroes this; `Guard::drop()` reads and restores.
+    pub(crate) static ALLOC_COUNTERS: Cell<AllocSnapshot> = const { Cell::new(AllocSnapshot::new()) };
+}
+
+impl AllocSnapshot {
+    pub(crate) const fn new() -> Self {
+        Self {
+            alloc_count: 0,
+            alloc_bytes: 0,
+            free_count: 0,
+            free_bytes: 0,
+        }
+    }
 }
 
 /// A global allocator wrapper that tracks allocation counts and bytes
 /// per instrumented function scope, with zero timing distortion.
 ///
-/// Wraps any inner `GlobalAlloc`. Measures its own bookkeeping overhead
-/// and subtracts it from the current function's timing via `overhead_ns`.
+/// Wraps any inner `GlobalAlloc`. Uses a destructor-free `Cell<AllocSnapshot>`
+/// for thread-local bookkeeping, which is safe on all Rust versions
+/// (including < 1.93.1 where TLS with destructors is forbidden for
+/// global allocators).
 pub struct PianoAllocator<A: GlobalAlloc> {
     inner: A,
 }
@@ -52,56 +75,26 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for PianoAllocator<A> {
     }
 }
 
+#[inline(always)]
 fn track_alloc(bytes: u64) {
-    let ok = IN_ALLOC_TRACKING.try_with(|flag| {
-        if flag.get() {
-            return false;
-        }
-        flag.set(true);
-        true
+    // Cell::get/set are plain memory reads/writes — no allocation, no
+    // re-entrancy risk, no destructor. Safe from the global allocator.
+    let _ = ALLOC_COUNTERS.try_with(|cell| {
+        let mut snap = cell.get();
+        snap.alloc_count += 1;
+        snap.alloc_bytes += bytes;
+        cell.set(snap);
     });
-    if ok != Ok(true) {
-        return;
-    }
-
-    let t0 = Instant::now();
-    let _ = STACK.try_with(|stack| {
-        if let Ok(mut s) = stack.try_borrow_mut()
-            && let Some(top) = s.last_mut()
-        {
-            top.alloc_count += 1;
-            top.alloc_bytes += bytes;
-            top.overhead_ns += t0.elapsed().as_nanos();
-        }
-    });
-
-    let _ = IN_ALLOC_TRACKING.try_with(|flag| flag.set(false));
 }
 
+#[inline(always)]
 fn track_dealloc(bytes: u64) {
-    let ok = IN_ALLOC_TRACKING.try_with(|flag| {
-        if flag.get() {
-            return false;
-        }
-        flag.set(true);
-        true
+    let _ = ALLOC_COUNTERS.try_with(|cell| {
+        let mut snap = cell.get();
+        snap.free_count += 1;
+        snap.free_bytes += bytes;
+        cell.set(snap);
     });
-    if ok != Ok(true) {
-        return;
-    }
-
-    let t0 = Instant::now();
-    let _ = STACK.try_with(|stack| {
-        if let Ok(mut s) = stack.try_borrow_mut()
-            && let Some(top) = s.last_mut()
-        {
-            top.free_count += 1;
-            top.free_bytes += bytes;
-            top.overhead_ns += t0.elapsed().as_nanos();
-        }
-    });
-
-    let _ = IN_ALLOC_TRACKING.try_with(|flag| flag.set(false));
 }
 
 #[cfg(test)]
@@ -126,6 +119,37 @@ mod tests {
         assert_eq!(rec.free_bytes, 256);
     }
 
+    #[test]
+    fn alloc_tracking_nested_scopes() {
+        reset();
+        {
+            let _outer = enter("outer_alloc");
+            track_alloc(100);
+            {
+                let _inner = enter("inner_alloc");
+                track_alloc(200);
+                track_dealloc(50);
+            }
+            track_alloc(300);
+            track_dealloc(75);
+        }
+        let invocations = collect_invocations();
+        let outer = invocations.iter().find(|r| r.name == "outer_alloc").unwrap();
+        let inner = invocations.iter().find(|r| r.name == "inner_alloc").unwrap();
+
+        // Inner scope should only see its own allocations
+        assert_eq!(inner.alloc_count, 1, "inner alloc_count");
+        assert_eq!(inner.alloc_bytes, 200, "inner alloc_bytes");
+        assert_eq!(inner.free_count, 1, "inner free_count");
+        assert_eq!(inner.free_bytes, 50, "inner free_bytes");
+
+        // Outer scope should see its own allocations (before + after inner)
+        assert_eq!(outer.alloc_count, 2, "outer alloc_count");
+        assert_eq!(outer.alloc_bytes, 400, "outer alloc_bytes");
+        assert_eq!(outer.free_count, 1, "outer free_count");
+        assert_eq!(outer.free_bytes, 75, "outer free_bytes");
+    }
+
     #[test]
     fn alloc_tracking_does_not_distort_timing() {
         reset();

piano-runtime/src/collector.rs

@@ -76,11 +76,8 @@ pub(crate) struct StackEntry {
     pub(crate) name: &'static str,
     pub(crate) start: Instant,
     pub(crate) children_ms: f64,
-    pub(crate) overhead_ns: u128,
-    pub(crate) alloc_count: u32,
-    pub(crate) alloc_bytes: u64,
-    pub(crate) free_count: u32,
-    pub(crate) free_bytes: u64,
+    /// Saved ALLOC_COUNTERS from before this scope, restored on Guard::drop.
+    pub(crate) saved_alloc: crate::alloc::AllocSnapshot,
     pub(crate) depth: u16,
 }
 
@@ -124,18 +121,31 @@ pub struct Guard {
 
 impl Drop for Guard {
     fn drop(&mut self) {
+        // Read this scope's alloc counters and restore the parent's saved state.
+        let scope_alloc = crate::alloc::ALLOC_COUNTERS
+            .try_with(|cell| {
+                let snap = cell.get();
+                // Will be overwritten below after we pop the stack entry.
+                snap
+            })
+            .unwrap_or_default();
+
         STACK.with(|stack| {
             let entry = stack.borrow_mut().pop();
             let Some(entry) = entry else {
                 eprintln!("piano-runtime: guard dropped without matching stack entry (bug)");
                 return;
             };
+
+            // Restore parent's saved alloc counters.
+            let _ = crate::alloc::ALLOC_COUNTERS.try_with(|cell| {
+                cell.set(entry.saved_alloc);
+            });
+
             let elapsed_ns = entry.start.elapsed().as_nanos() as u64;
             let elapsed_ms = elapsed_ns as f64 / 1_000_000.0;
             let children_ns = (entry.children_ms * 1_000_000.0) as u64;
-            let self_ns = elapsed_ns
-                .saturating_sub(children_ns)
-                .saturating_sub(entry.overhead_ns as u64);
+            let self_ns = elapsed_ns.saturating_sub(children_ns);
             let start_ns = entry.start.duration_since(*EPOCH).as_nanos() as u64;
             let children_ms = entry.children_ms;
 
@@ -160,10 +170,10 @@ impl Drop for Guard {
                 start_ns,
                 elapsed_ns,
                 self_ns,
-                alloc_count: entry.alloc_count,
-                alloc_bytes: entry.alloc_bytes,
-                free_count: entry.free_count,
-                free_bytes: entry.free_bytes,
+                alloc_count: scope_alloc.alloc_count,
+                alloc_bytes: scope_alloc.alloc_bytes,
+                free_count: scope_alloc.free_count,
+                free_bytes: scope_alloc.free_bytes,
                 depth: entry.depth,
             };
 
@@ -193,17 +203,23 @@ impl Drop for Guard {
 pub fn enter(name: &'static str) -> Guard {
     // Touch EPOCH so relative timestamps are anchored to process start.
     let _ = *EPOCH;
+
+    // Save current alloc counters and zero them for this scope.
+    let saved_alloc = crate::alloc::ALLOC_COUNTERS
+        .try_with(|cell| {
+            let snap = cell.get();
+            cell.set(crate::alloc::AllocSnapshot::new());
+            snap
+        })
+        .unwrap_or_default();
+
     STACK.with(|stack| {
         let depth = stack.borrow().len() as u16;
         stack.borrow_mut().push(StackEntry {
             name,
             start: Instant::now(),
             children_ms: 0.0,
-            overhead_ns: 0,
-            alloc_count: 0,
-            alloc_bytes: 0,
-            free_count: 0,
-            free_bytes: 0,
+            saved_alloc,
             depth,
         });
     });
@@ -633,17 +649,22 @@ pub fn fork() -> Option<SpanContext> {
 /// thread correctly attributes children time. Returns an `AdoptGuard` that
 /// propagates elapsed time back to the parent on drop.
 pub fn adopt(ctx: &SpanContext) -> AdoptGuard {
+    // Save current alloc counters and zero them, same as enter().
+    let saved_alloc = crate::alloc::ALLOC_COUNTERS
+        .try_with(|cell| {
+            let snap = cell.get();
+            cell.set(crate::alloc::AllocSnapshot::new());
+            snap
+        })
+        .unwrap_or_default();
+
     STACK.with(|stack| {
         let depth = stack.borrow().len() as u16;
         stack.borrow_mut().push(StackEntry {
             name: ctx.parent_name,
             start: Instant::now(),
             children_ms: 0.0,
-            overhead_ns: 0,
-            alloc_count: 0,
-            alloc_bytes: 0,
-            free_count: 0,
-            free_bytes: 0,
+            saved_alloc,
             depth,
         });
     });

src/build.rs

@@ -111,10 +111,14 @@ fn extract_rendered_errors(json_output: &str) -> Vec<String> {
 /// Build the instrumented binary using `cargo build --message-format=json`.
 /// Returns the path to the compiled executable.
 pub fn build_instrumented(staging_dir: &Path, target_dir: &Path) -> Result<PathBuf, Error> {
+    // Remove RUSTUP_TOOLCHAIN so the target project's rust-toolchain.toml
+    // is respected. Without this, nested cargo invocations inherit the
+    // parent's toolchain, ignoring the project's pinned version.
     let output = Command::new("cargo")
         .arg("build")
         .arg("--message-format=json")
         .env("CARGO_TARGET_DIR", target_dir)
+        .env_remove("RUSTUP_TOOLCHAIN")
         .current_dir(staging_dir)
         .output()?;