First pass

gyuheon0h · gyuheon0h · commit a530d9951cea · 2025-10-03T16:48:22.000Z
diff --git a/datadog-crashtracker-ffi/src/runtime_callback.rs b/datadog-crashtracker-ffi/src/runtime_callback.rs
@@ -15,9 +15,7 @@ use datadog_crashtracker::{
 pub use datadog_crashtracker::CallbackType as ddog_CallbackType;
 pub use datadog_crashtracker::RuntimeType as ddog_RuntimeType;
 
-#[cfg(test)]
-use datadog_crashtracker::clear_runtime_callback;
-use std::ffi::c_char;
+pub use datadog_crashtracker::RuntimeStackFrame as ddog_RuntimeStackFrame;
 
 /// Result type for runtime callback registration
 #[repr(C)]
@@ -60,7 +58,7 @@ impl From<CallbackError> for CallbackResult {
 /// static void my_runtime_callback(
 ///     void (*emit_frame)(const ddog_RuntimeStackFrame*),
 ///     void (*emit_stacktrace_string)(const char*),
-///     void* context
+///     void* writer_ctx
 /// ) {
 ///     // Collect runtime frames and call emit_frame for each one
 ///     ddog_RuntimeStackFrame frame = {
@@ -71,7 +69,7 @@ impl From<CallbackError> for CallbackResult {
 ///         .class_name = "MyClass",
 ///         .module_name = NULL
 ///     };
-///     emit_frame(&frame);
+///     emit_frame(writer_ctx, &frame);
 /// }
 ///
 ///
@@ -131,7 +129,7 @@ pub extern "C" fn ddog_crasht_is_runtime_callback_registered() -> bool {
 /// - The caller should not free the returned pointer
 /// - The returned string should be copied if it needs to persist beyond callback lifetime
 #[no_mangle]
-pub unsafe extern "C" fn ddog_crasht_get_registered_runtime_type() -> *const c_char {
+pub unsafe extern "C" fn ddog_crasht_get_registered_runtime_type() -> *const std::ffi::c_char {
     get_registered_runtime_type_ptr()
 }
 
@@ -145,18 +143,15 @@ pub unsafe extern "C" fn ddog_crasht_get_registered_runtime_type() -> *const c_c
 /// - The caller should not free the returned pointer
 /// - The returned string should be copied if it needs to persist beyond callback lifetime
 #[no_mangle]
-pub unsafe extern "C" fn ddog_crasht_get_registered_callback_type() -> *const c_char {
+pub unsafe extern "C" fn ddog_crasht_get_registered_callback_type() -> *const std::ffi::c_char {
     get_registered_callback_type_ptr()
 }
 
-// Re-export the core RuntimeStackFrame for C interop; consumers need to create RuntimeStackFrame
-pub use datadog_crashtracker::RuntimeStackFrame as ddog_RuntimeStackFrame;
-
 #[cfg(test)]
 mod tests {
     use super::*;
-    use datadog_crashtracker::RuntimeStackFrame;
-    use std::ffi::CString;
+    use datadog_crashtracker::{clear_runtime_callback, RuntimeStackFrame};
+    use std::ffi::{c_char, c_void, CString};
     use std::ptr;
     use std::sync::Mutex;
 
@@ -165,8 +160,9 @@ mod tests {
     static TEST_MUTEX: Mutex<()> = Mutex::new(());
 
     unsafe extern "C" fn test_runtime_callback(
-        emit_frame: unsafe extern "C" fn(*const RuntimeStackFrame),
-        _emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
+        emit_frame: unsafe extern "C" fn(*mut c_void, *const RuntimeStackFrame),
+        _emit_stacktrace_string: unsafe extern "C" fn(*mut c_void, *const c_char),
+        writer_ctx: *mut c_void,
     ) {
         let function_name = CString::new("test_function").unwrap();
         let file_name = CString::new("test.rb").unwrap();
@@ -183,7 +179,7 @@ mod tests {
             module_name: ptr::null(),
         };
 
-        emit_frame(&frame);
+        emit_frame(writer_ctx, &frame);
     }
 
     #[test]
diff --git a/datadog-crashtracker/src/runtime_callback.rs b/datadog-crashtracker/src/runtime_callback.rs
@@ -114,11 +114,6 @@ impl std::str::FromStr for CallbackType {
 static RUNTIME_CALLBACK: AtomicPtr<(RuntimeStackCallback, RuntimeType, CallbackType)> =
     AtomicPtr::new(ptr::null_mut());
 
-thread_local! {
-    static FRAME_WRITER: std::cell::RefCell<Option<*mut dyn std::io::Write>> = const { std::cell::RefCell::new(None) };
-    static FRAME_COUNT: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
-}
-
 #[repr(C)]
 #[derive(Debug, Clone)]
 pub struct RuntimeStackFrame {
@@ -145,17 +140,21 @@ pub struct RuntimeStackFrame {
 /// - Only async-signal-safe functions
 ///
 /// # Parameters
-/// - `emit_frame`: Function to call for each runtime frame (marked unsafe because it takes raw pointers)
-/// - `emit_stacktrace_string`: Function to call for complete stacktrace string (marked unsafe because it takes raw C string pointers)
+/// - `emit_frame`: Function to call for each runtime frame (takes writer context and frame pointer)
+/// - `emit_stacktrace_string`: Function to call for complete stacktrace string (takes writer
+///   context and C string)
+/// - `writer_ctx`: Opaque pointer to writer context that should be passed to emit functions
 ///
 /// # Safety
 /// The callback function is marked unsafe because:
 /// - It receives function pointers that take raw pointers as parameters
 /// - The callback must ensure any pointers it passes to these functions are valid
 /// - All C strings passed must be null-terminated and remain valid for the call duration
+/// - The writer_ctx must be passed unchanged to the emit functions
 pub type RuntimeStackCallback = unsafe extern "C" fn(
-    emit_frame: unsafe extern "C" fn(*const RuntimeStackFrame),
-    emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
+    emit_frame: unsafe extern "C" fn(*mut std::ffi::c_void, *const RuntimeStackFrame),
+    emit_stacktrace_string: unsafe extern "C" fn(*mut std::ffi::c_void, *const c_char),
+    writer_ctx: *mut std::ffi::c_void,
 );
 
 /// Runtime stack representation for JSON serialization
@@ -367,78 +366,69 @@ pub(crate) unsafe fn invoke_runtime_callback_with_writer<W: std::io::Write>(
     // to a properly aligned, initialized tuple.
     let (callback_fn, _, _) = &*callback_ptr;
 
-    // Reset frame counter
-    FRAME_COUNT.with(|count| count.set(0));
-
-    // Define the emit_frame function that writes directly to the pipe
-    // Safety: This function is called by the user's runtime callback. It's marked unsafe
-    // because it dereferences raw pointers, but the safety is guaranteed by the contract
-    // with the runtime callback system.
-    unsafe extern "C" fn emit_frame_collector(frame: *const RuntimeStackFrame) {
-        // We need access to the writer, so we'll use thread-local storage for it
-        FRAME_WRITER.with(|writer_cell| {
-            if let Some(writer_ptr) = *writer_cell.borrow() {
-                // Safety: writer_ptr was stored in thread-local storage before
-                // calling the user callback, so it's guaranteed to be valid and point
-                // to the same writer object we passed in.
-                let writer = &mut *writer_ptr;
-
-                FRAME_COUNT.with(|count| {
-                    let current_count = count.get();
-
-                    // Add comma separator for frames after the first
-                    if current_count > 0 {
-                        let _ = write!(writer, ", ");
-                    }
-
-                    // Write the frame as JSON
-                    // Safety: frame pointer is passed from the runtime callback, writer is valid
-                    let _ = emit_frame_as_json(writer, frame);
-                    let _ = writer.flush();
-
-                    count.set(current_count + 1);
-                });
-            }
-        });
+    let mut frame_count = 0usize;
+
+    // Define the emit_frame function that writes directly to the writer
+    // Safety: This function receives writer context and frame pointer from the runtime callback.
+    // The writer context is guaranteed to be valid for the duration of the callback.
+    unsafe extern "C" fn emit_frame_collector(
+        writer_ctx: *mut std::ffi::c_void,
+        frame: *const RuntimeStackFrame,
+    ) {
+        if writer_ctx.is_null() || frame.is_null() {
+            return;
+        }
+
+        // Safety: writer_ctx was created from a valid writer reference and frame_count pointer
+        let (writer, frame_count) =
+            &mut *(writer_ctx as *mut (&mut dyn std::io::Write, &mut usize));
+
+        // Add comma separator for frames after the first
+        if **frame_count > 0 {
+            let _ = write!(writer, ", ");
+        }
+
+        // Write the frame as JSON
+        // Safety: frame pointer is passed from the runtime callback
+        let _ = emit_frame_as_json(*writer, frame);
+        let _ = writer.flush();
+
+        **frame_count += 1;
     }
 
-    // Safety: This function is called by the user's runtime callback with a C string pointer.
-    // The safety requirements are documented in the RuntimeStackCallback type.
-    unsafe extern "C" fn emit_stacktrace_string_collector(stacktrace_string: *const c_char) {
-        if stacktrace_string.is_null() {
+    // Safety: This function receives writer context and C string pointer from the runtime callback.
+    unsafe extern "C" fn emit_stacktrace_string_collector(
+        writer_ctx: *mut std::ffi::c_void,
+        stacktrace_string: *const c_char,
+    ) {
+        if writer_ctx.is_null() || stacktrace_string.is_null() {
             return;
         }
 
+        // Safety: writer_ctx was created from a valid writer reference
+        let (writer, _) = &mut *(writer_ctx as *mut (&mut dyn std::io::Write, &mut usize));
+
         // Safety: stacktrace_string is guaranteed by the runtime callback contract to be
         // a valid, null-terminated C string that remains valid for the duration of this call.
         let cstr = std::ffi::CStr::from_ptr(stacktrace_string);
         let bytes = cstr.to_bytes();
 
-        FRAME_WRITER.with(|writer_cell| {
-            if let Some(writer_ptr) = *writer_cell.borrow() {
-                // Safety: writer_ptr was stored in thread-local storage just before
-                // calling the user callback, so it's guaranteed to be valid.
-                let writer = &mut *writer_ptr;
-                let _ = writer.write_all(bytes);
-                let _ = writer.flush();
-            }
-        });
+        let _ = writer.write_all(bytes);
+        let _ = writer.flush();
     }
 
-    // Store the writer in thread-local storage so the emit functions can access it
-    FRAME_WRITER.with(|writer_cell| {
-        *writer_cell.borrow_mut() = Some(writer as *mut W as *mut dyn std::io::Write);
-    });
+    // Create writer context that bundles writer and frame counter
+    let mut writer_context = (writer as &mut dyn std::io::Write, &mut frame_count);
+    let writer_ctx = &mut writer_context as *mut _ as *mut std::ffi::c_void;
 
-    // Invoke the user callback
+    // Invoke the user callback with the writer context
     // Safety: callback_fn was verified to be non-null during registration, and the
-    // emit functions we're passing are compatible with the expected function signatures.
-    callback_fn(emit_frame_collector, emit_stacktrace_string_collector);
-
-    // Clear the writer reference
-    FRAME_WRITER.with(|writer_cell| {
-        *writer_cell.borrow_mut() = None;
-    });
+    // emit functions and writer context are valid for the duration of this call.
+    callback_fn(
+        emit_frame_collector,
+        emit_stacktrace_string_collector,
+        writer_ctx,
+    );
 
     Ok(())
 }
@@ -555,8 +545,9 @@ mod tests {
     static TEST_MUTEX: Mutex<()> = Mutex::new(());
 
     unsafe extern "C" fn test_emit_frame_callback(
-        emit_frame: unsafe extern "C" fn(*const RuntimeStackFrame),
-        _emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
+        emit_frame: unsafe extern "C" fn(*mut std::ffi::c_void, *const RuntimeStackFrame),
+        _emit_stacktrace_string: unsafe extern "C" fn(*mut std::ffi::c_void, *const c_char),
+        writer_ctx: *mut std::ffi::c_void,
     ) {
         let function_name = CString::new("test_function").unwrap();
         let file_name = CString::new("test.rb").unwrap();
@@ -571,16 +562,19 @@ mod tests {
             module_name: ptr::null(),
         };
 
-        emit_frame(&frame);
+        // Safety: frame is a valid RuntimeStackFrame with valid C string pointers
+        emit_frame(writer_ctx, &frame);
     }
 
     unsafe extern "C" fn test_emit_stacktrace_string_callback(
-        _emit_frame: unsafe extern "C" fn(*const RuntimeStackFrame),
-        emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
+        _emit_frame: unsafe extern "C" fn(*mut std::ffi::c_void, *const RuntimeStackFrame),
+        emit_stacktrace_string: unsafe extern "C" fn(*mut std::ffi::c_void, *const c_char),
+        writer_ctx: *mut std::ffi::c_void,
     ) {
         let stacktrace_string = CString::new("test_stacktrace_string").unwrap();
 
-        emit_stacktrace_string(stacktrace_string.as_ptr());
+        // Safety: stacktrace_string.as_ptr() returns a valid null-terminated C string
+        emit_stacktrace_string(writer_ctx, stacktrace_string.as_ptr());
     }
 
     fn ensure_callback_cleared() {
