remove the allocation in signal handle of framehop

Signed-off-by: Yang Keao <[email protected]>

Author: YangKeao

.github/workflows/rust.yml

@@ -159,3 +159,9 @@ jobs:
         with:
           command: test
           args: --features flamegraph,protobuf-codec --target ${{ matrix.target }}
+
+      - name: Run cargo test framehop
+        uses: actions-rs/[email protected]
+        with:
+          command: test
+          args: --features flamegraph,protobuf-codec,framehop-unwinder --target ${{ matrix.target }}

src/backtrace/framehop_unwinder.rs

@@ -3,6 +3,7 @@ use framehop::{
 };
 use libc::{c_void, ucontext_t};
 use once_cell::sync::Lazy;
+use spin::RwLock;
 mod shlib;
 
 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
@@ -123,7 +124,8 @@ fn read_stack(addr: u64) -> Result<u64, ()> {
     }
 }
 
-static mut UNWINDER: Lazy<FramehopUnwinder> = Lazy::new(|| FramehopUnwinder::new());
+static UNWINDER: Lazy<RwLock<FramehopUnwinder>> =
+    Lazy::new(|| RwLock::new(FramehopUnwinder::new()));
 #[derive(Clone, Debug)]
 pub struct Frame {
     pub ip: usize,
@@ -158,10 +160,22 @@ pub struct Trace;
 impl super::Trace for Trace {
     type Frame = Frame;
 
+    fn init() {
+        let _ = UNWINDER.read();
+    }
+
     fn trace<F: FnMut(&Self::Frame) -> bool>(ctx: *mut c_void, cb: F)
     where
         Self: Sized,
     {
-        unsafe { UNWINDER.iter_frames(ctx, cb) };
+        // For Linux, this `try_write` should always succeed, because `SIGPROF` will never be delivered to
+        // another thread while the signal handler is running. However, I'm not sure about other OSes, so
+        // we use `try_write` to be safe instead of using `static mut` and `unsafe` directly.
+        match UNWINDER.try_write() {
+            None => return,
+            Some(mut unwinder) => {
+                unwinder.iter_frames(ctx, cb);
+            }
+        }
     }
 }

src/backtrace/mod.rs

@@ -43,6 +43,9 @@ pub trait Frame: Sized + Clone {
 pub trait Trace {
     type Frame;
 
+    // init will be called before running the first trace in signal handler
+    fn init() {}
+
     fn trace<F: FnMut(&Self::Frame) -> bool>(_: *mut libc::c_void, cb: F)
     where
         Self: Sized;

src/profiler.rs

@@ -188,6 +188,7 @@ pub struct ProfilerGuard<'a> {
 fn trigger_lazy() {
     let _ = backtrace::Backtrace::new();
     let _profiler = PROFILER.read();
+    TraceImpl::init();
 }
 
 impl ProfilerGuard<'_> {
@@ -524,3 +525,80 @@ impl Profiler {
         if let Ok(()) = self.data.add(frames, 1) {}
     }
 }
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    struct AllocDetector {
+        should_count_alloc: std::sync::atomic::AtomicBool,
+        alloc_count: std::sync::atomic::AtomicUsize,
+    }
+
+    unsafe impl std::alloc::GlobalAlloc for AllocDetector {
+        unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
+            if self
+                .should_count_alloc
+                .load(std::sync::atomic::Ordering::SeqCst)
+            {
+                if self
+                    .alloc_count
+                    .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
+                    == 0
+                {
+                    // it's not safe to print and unwind here, but it's fine for test purpose
+                    println!(
+                        "allocation happened during unwinding! {:?}",
+                        backtrace::Backtrace::new()
+                    );
+                }
+            }
+
+            unsafe { std::alloc::System.alloc(layout) }
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
+            unsafe { std::alloc::System.dealloc(ptr, layout) }
+        }
+    }
+    impl AllocDetector {
+        fn enable_count_alloc(&self) {
+            self.should_count_alloc
+                .store(true, std::sync::atomic::Ordering::SeqCst);
+        }
+
+        fn disable_count_alloc(&self) {
+            self.should_count_alloc
+                .store(false, std::sync::atomic::Ordering::SeqCst);
+        }
+
+        fn alloc_count(&self) -> usize {
+            self.alloc_count.load(std::sync::atomic::Ordering::SeqCst)
+        }
+    }
+
+    #[global_allocator]
+    static ALLOC: AllocDetector = AllocDetector {
+        should_count_alloc: std::sync::atomic::AtomicBool::new(false),
+        alloc_count: std::sync::atomic::AtomicUsize::new(0),
+    };
+
+    #[test]
+    fn test_no_alloc_during_unwind() {
+        trigger_lazy();
+        PROFILER.write().as_mut().unwrap().start().unwrap();
+        let timer = Timer::new(999);
+        let start = std::time::Instant::now();
+        ALLOC.enable_count_alloc();
+
+        // busy loop for a while to trigger some samples
+        while start.elapsed().as_millis() < 500 {
+            let _ = std::hint::black_box(());
+        }
+        assert_eq!(ALLOC.alloc_count(), 0);
+
+        ALLOC.disable_count_alloc();
+        drop(timer);
+        PROFILER.write().as_mut().unwrap().stop().unwrap();
+    }
+}