Add unit tests for thread context buffer

Author: gyuheon0h

libdd-crashtracker/src/collector/emitters.rs

@@ -454,7 +454,6 @@ fn emit_thread_block(
 
     // Header JSON line
     write!(w, "{{\"tid\": {tid}, \"crashed\": {crashed}")?;
-    // Escape name for JSON: replace control chars and quotes
     let safe_name = name
         .replace('\\', "\\\\")
         .replace('"', "\\\"")
@@ -497,7 +496,7 @@ fn read_thread_name(pid: libc::pid_t, tid: libc::pid_t) -> Option<String> {
 }
 
 /// Read a thread's scheduler state letter from /proc/<pid>/task/<tid>/status.
-/// Returns the single-letter state (e.g. "S", "R", "D") or None on failure.
+/// Returns the single-letter state ("S", "R", "D") or None on failure.
 #[cfg(target_os = "linux")]
 fn read_thread_state(pid: libc::pid_t, tid: libc::pid_t) -> Option<String> {
     let path = format!("/proc/{pid}/task/{tid}/status");
@@ -535,7 +534,7 @@ fn enumerate_task_tids(pid: libc::pid_t) -> Vec<libc::pid_t> {
 /// Called from the collector child (after fork), so std::fs is safe to use.
 /// For each thread:
 ///   - Reads name and state from /proc/<ppid>/task/<tid>/
-///   - If a ucontext was captured (Phase 2 / collect_all_threads), emits the stack trace.
+///   - If a ucontext was captured (collect_all_threads), emits the stack trace.
 ///   - Otherwise emits an empty (incomplete) stack trace.
 #[cfg(target_os = "linux")]
 fn emit_all_threads(
@@ -546,7 +545,7 @@ fn emit_all_threads(
 ) -> Result<(), EmitterError> {
     use crate::collector::thread_context_buffer::iter_collected_contexts;
 
-    // Build a map from TID to captured ucontext pointer (may be empty if Phase 2 buffer
+    // Build a map from TID to captured ucontext pointer (may be empty if buffer
     // was not initialised or the thread did not respond in time).
     let contexts: std::collections::HashMap<libc::pid_t, *const ucontext_t> =
         iter_collected_contexts()

libdd-crashtracker/src/collector/thread_context_buffer.rs

@@ -99,7 +99,7 @@ static BUFFER: AtomicPtr<ThreadContextBuffer> = AtomicPtr::new(ptr::null_mut());
 
 /// Allocate the thread context buffer.
 ///
-/// Must be called once from a non-signal context (typically inside `init()`).
+/// Must be called once from a non-signal context (inside `init()`).
 /// Subsequent calls are a no-op.
 pub fn init_thread_context_buffer(capacity: usize) {
     let capacity = capacity.clamp(1, MAX_TRACKED_THREADS);
@@ -164,7 +164,7 @@ pub(crate) extern "C" fn handle_collect_context_signal(
         if slot.tid.load(Ordering::SeqCst) == tid {
             if !ucontext.is_null() {
                 // SAFETY: ucontext is the kernel-saved register state, valid for the
-                // duration of the signal handler.  The slot ctx field is not accessed
+                // duration of the signal handler. The slot ctx field is not accessed
                 // by anyone else until ready is observed as true.
                 unsafe {
                     ptr::copy_nonoverlapping(
@@ -279,7 +279,7 @@ fn parse_tid(bytes: &[u8]) -> Option<libc::pid_t> {
 /// then wait for each contacted thread to write its ucontext into its slot.
 ///
 /// Called from the primary crash signal handler before fork()-ing the collector
-/// child.  Must be async-signal-safe.
+/// child. Must be async-signal-safe.
 ///
 /// - `crashing_tid`  TID of the thread that received the fatal signal (skipped).
 /// - `max_threads`   cap on how many threads to sample.
@@ -380,6 +380,37 @@ pub fn iter_collected_contexts() -> impl Iterator<Item = CollectedContext> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    /// Claim the first free slot for `tid`.  Returns the slot index on success.
+    fn claim_slot(buf: &ThreadContextBuffer, tid: i32) -> Option<usize> {
+        buf.slots.iter().enumerate().find_map(|(i, slot)| {
+            slot.tid
+                .compare_exchange(0, tid, Ordering::SeqCst, Ordering::Relaxed)
+                .ok()
+                .map(|_| i)
+        })
+    }
+
+    /// Install `handle_collect_context_signal` as the SIGUSR2 handler and return
+    /// the previous `sigaction` so the caller can restore it.
+    fn install_context_handler() -> libc::sigaction {
+        let mut sa: libc::sigaction = unsafe { std::mem::zeroed() };
+        sa.sa_sigaction = handle_collect_context_signal as *const () as usize;
+        sa.sa_flags = libc::SA_SIGINFO | libc::SA_RESTART | libc::SA_NODEFER;
+        let mut old: libc::sigaction = unsafe { std::mem::zeroed() };
+        unsafe { libc::sigaction(libc::SIGUSR2, &sa, &mut old) };
+        old
+    }
+
+    fn restore_handler(old: &libc::sigaction) {
+        unsafe { libc::sigaction(libc::SIGUSR2, old, std::ptr::null_mut()) };
+    }
+
+    fn current_tid() -> i32 {
+        unsafe { libc::syscall(libc::SYS_gettid) as i32 }
+    }
 
     #[test]
     fn test_parse_tid_valid() {
@@ -419,4 +450,186 @@ mod tests {
         assert_eq!(buf.slots[0].tid.load(Ordering::SeqCst), 0);
         assert!(!buf.slots[0].ready.load(Ordering::SeqCst));
     }
+
+    #[test]
+    fn test_signal_handler_writes_context_and_sets_ready() {
+        init_thread_context_buffer(16);
+        reset_thread_context_buffer();
+
+        let tid = current_tid();
+        let buf = get_buffer().unwrap();
+        let idx = claim_slot(buf, tid).expect("should find a free slot");
+
+        let mut fake_ctx: libc::ucontext_t = unsafe { std::mem::zeroed() };
+        fake_ctx.uc_flags = 0x1234_5678;
+
+        handle_collect_context_signal(
+            libc::SIGUSR2,
+            std::ptr::null_mut(),
+            &mut fake_ctx as *mut libc::ucontext_t as *mut libc::c_void,
+        );
+
+        let slot = &buf.slots[idx];
+        assert!(
+            slot.ready.load(Ordering::Acquire),
+            "slot should be marked ready after handler runs"
+        );
+        // Verify the context bytes were actually copied.
+        let copied_flags = unsafe { (*(*slot.ctx.get()).as_ptr()).uc_flags };
+        assert_eq!(
+            copied_flags, 0x1234_5678,
+            "uc_flags sentinel should survive the copy"
+        );
+
+        reset_thread_context_buffer();
+    }
+
+    /// Even with a null ucontext (kernel gave us nothing), the slot must be
+    /// marked ready so the spin-wait in collect_thread_contexts can complete.
+    #[test]
+    fn test_signal_handler_null_ucontext_still_sets_ready() {
+        init_thread_context_buffer(16);
+        reset_thread_context_buffer();
+
+        let tid = current_tid();
+        let buf = get_buffer().unwrap();
+        let idx = claim_slot(buf, tid).expect("should find a free slot");
+
+        handle_collect_context_signal(
+            libc::SIGUSR2,
+            std::ptr::null_mut(),
+            std::ptr::null_mut(), // null ucontext
+        );
+
+        assert!(
+            buf.slots[idx].ready.load(Ordering::Acquire),
+            "null ucontext should still set ready so the caller doesn't stall"
+        );
+
+        reset_thread_context_buffer();
+    }
+
+    /// The handler must silently do nothing when the current TID has no slot
+    /// claimed for it
+    #[test]
+    fn test_signal_handler_ignores_unclaimed_tid() {
+        init_thread_context_buffer(16);
+        reset_thread_context_buffer(); // all slots tid=0, ready=false
+
+        let mut fake_ctx: libc::ucontext_t = unsafe { std::mem::zeroed() };
+        // Call with no slot pre-claimed for the current TID.
+        handle_collect_context_signal(
+            libc::SIGUSR2,
+            std::ptr::null_mut(),
+            &mut fake_ctx as *mut libc::ucontext_t as *mut libc::c_void,
+        );
+
+        let buf = get_buffer().unwrap();
+        let tid = current_tid();
+        for slot in &buf.slots {
+            assert!(
+                !slot.ready.load(Ordering::Acquire),
+                "no slot should be ready when TID was not pre-claimed"
+            );
+            assert_ne!(
+                slot.tid.load(Ordering::Acquire),
+                tid,
+                "handler must not claim a slot for an unclaimed TID"
+            );
+        }
+
+        reset_thread_context_buffer();
+    }
+
+    /// The crashing thread shouldn't appear as one of the captured threads
+    #[test]
+    #[cfg_attr(miri, ignore)] // sigaction is not supported by Miri
+    fn test_collect_never_captures_crashing_tid() {
+        init_thread_context_buffer(16);
+        reset_thread_context_buffer();
+
+        let old_sa = install_context_handler();
+        let crashing_tid = unsafe { libc::syscall(libc::SYS_gettid) as libc::pid_t };
+
+        collect_thread_contexts(crashing_tid, 16, 300);
+
+        let buf = get_buffer().unwrap();
+        for slot in &buf.slots {
+            assert_ne!(
+                slot.tid.load(Ordering::Acquire),
+                crashing_tid as i32,
+                "crashing TID ({crashing_tid}) must not appear in any slot"
+            );
+        }
+
+        restore_handler(&old_sa);
+        reset_thread_context_buffer();
+    }
+
+    /// Two background threads sleeping in a syscall should both respond to
+    /// SIGUSR2 and have their contexts captured before the timeout fires.
+    ///
+    /// We record each worker's TID before they block and verify that exactly
+    /// those TIDs appear in the buffer afterwards.
+    #[test]
+    #[cfg_attr(miri, ignore)] // sigaction is not supported by Miri
+    fn test_collect_captures_background_threads() {
+        init_thread_context_buffer(16);
+        reset_thread_context_buffer();
+
+        let old_sa = install_context_handler();
+
+        let (tx, rx) = std::sync::mpsc::channel::<i32>();
+        let keep = Arc::new(std::sync::atomic::AtomicBool::new(true));
+
+        let handles: Vec<_> = (0..2)
+            .map(|_| {
+                let tx = tx.clone();
+                let keep = Arc::clone(&keep);
+                std::thread::spawn(move || {
+                    tx.send(current_tid()).unwrap();
+                    while keep.load(Ordering::Relaxed) {
+                        std::thread::sleep(Duration::from_millis(20));
+                    }
+                })
+            })
+            .collect();
+
+        // Receive exactly one TID per worker
+        drop(tx);
+        let worker_tids: Vec<i32> = (0..2).map(|_| rx.recv().unwrap()).collect();
+        std::thread::sleep(Duration::from_millis(50));
+
+        let crashing_tid = unsafe { libc::syscall(libc::SYS_gettid) as libc::pid_t };
+        collect_thread_contexts(crashing_tid, 16, 500);
+
+        let buf = get_buffer().unwrap();
+
+        // Both workers must have responded and their contexts must be in the buffer.
+        for &wtid in &worker_tids {
+            let found = buf
+                .slots
+                .iter()
+                .any(|s| s.tid.load(Ordering::Acquire) == wtid && s.ready.load(Ordering::Acquire));
+            assert!(
+                found,
+                "worker TID {wtid} should have its context captured in the buffer"
+            );
+        }
+
+        // The crashing TID must still be absent.
+        assert!(
+            buf.slots
+                .iter()
+                .all(|s| s.tid.load(Ordering::Acquire) != crashing_tid as i32),
+            "crashing TID ({crashing_tid}) must not be in the buffer"
+        );
+
+        keep.store(false, Ordering::Relaxed);
+        for h in handles {
+            let _ = h.join();
+        }
+        restore_handler(&old_sa);
+        reset_thread_context_buffer();
+    }
 }

libdd-crashtracker/src/receiver/receive_report.rs

@@ -661,8 +661,8 @@ fn enrich_thread_name(builder: &mut CrashInfoBuilder) -> anyhow::Result<()> {
     // the collector child may have been unable to read it in time.
     //
     // Note: builder.error.threads holds ThreadData values only after they have been
-    // parsed from the wire protocol. We iterate them and patch names that look like
-    // bare integers (i.e. the TID fallback).
+    // parsed from the unix socket. We iterate them and patch names that look like
+    // bare integers (the TID fallback).
     if let Some(pid_info) = builder.proc_info.as_ref() {
         let pid = pid_info.pid;
         if let Some(threads) = builder.error.threads.as_mut() {
@@ -671,7 +671,7 @@ fn enrich_thread_name(builder: &mut CrashInfoBuilder) -> anyhow::Result<()> {
                 if thread.name.parse::<i32>().is_err() {
                     continue;
                 }
-                // Name is a bare TID integer — try to read from /proc.
+                // Name is a bare TID integer; try to read from /proc.
                 let Ok(tid) = thread.name.parse::<i32>() else {
                     continue;
                 };