feat(data-pipeline)!: add fork safety hooks and cancellation token for trace exporter FFI (#2051)

# What does this PR do?

Add fork safety and cooperative cancellation support to the trace exporter FFI:

1. **Fork hooks**: Expose `ddog_trace_exporter_before_fork`, `_after_fork_in_parent`,
   and `_after_fork_in_child` that delegate to the underlying `SharedRuntime`. These
   allow C callers to coordinate the tokio runtime lifecycle around process forks.

2. **Cancellation token**: Introduce `ddog_trace_exporter_cancel_token_new`, `_cancel`,
   and `_drop` for managing cancellation tokens. Thread the token into
   `ddog_trace_exporter_send_trace_chunks` via `tokio::select!` so callers can
   cooperatively abort in-flight HTTP requests.

# Motivation

FUP to #1952.

Ruby application servers (Puma, Unicorn, Passenger) fork worker processes. The tokio
runtime does not survive `fork()` — its threads are not carried over. Without the fork
hooks, the trace exporter is dead in child processes.

The existing `RUBY_UBF_IO` unblock function on the dd-trace-rb side sends a signal
that cannot actually cancel an in-flight Rust HTTP request. The cancellation token
enables cooperative abort: when Ruby interrupts a thread (shutdown, `Thread#kill`),
the UBF cancels the token, which causes `tokio::select!` to abort the send and return
promptly.

# Additional Notes

- Added a `pub fn shared_runtime()` accessor to `TraceExporter` since the field was private.
- The cancellation token uses `Handle<CancellationToken>` from `libdd-common-ffi`, same
  pattern as profiling-ffi, with a `ddog_TraceExporterCancelToken` cbindgen rename to
  avoid symbol conflicts.
- The `cancel` parameter on `send_trace_chunks` is nullable; passing `NULL` preserves
  existing behavior.

AI was used to accelerate implementation; all code was reviewed and understood.

# How to test the change?

```
cargo test -p libdd-data-pipeline-ffi --lib
```

46 tests pass (40 existing + 6 new for fork hooks and cancellation token).

Co-authored-by: dd-octo-sts[bot] <200755185+dd-octo-sts[bot]@users.noreply.github.com>
Co-authored-by: VianneyRuhlmann <vianney.ruhlmann@datadoghq.com>

Author: 3 people

Cargo.lock

@@ -19,6 +19,7 @@ default = [
   "log",
   "ddsketch",
   "ffe",
+  "shared-runtime",
 ]
 crashtracker = []
 profiling = []
@@ -29,6 +30,7 @@ library-config = []
 log = []
 ddsketch = []
 ffe = []
+shared-runtime = []
 regex-lite = ["libdd-common/regex-lite"]
 
 [lib]

builder/Cargo.toml

@@ -74,6 +74,8 @@ pub fn main() {
         f.push("ddsketch-ffi".to_string());
         #[cfg(feature = "ffe")]
         f.push("datadog-ffe-ffi".to_string());
+        #[cfg(feature = "shared-runtime")]
+        f.push("shared-runtime".to_string());
         f
     };
 

builder/src/bin/release.rs

@@ -59,6 +59,8 @@ impl Profiling {
         headers.push("ddsketch.h");
         #[cfg(feature = "ffe")]
         headers.push("ffe.h");
+        #[cfg(feature = "shared-runtime")]
+        headers.push("shared-runtime.h");
 
         let mut origin_path: PathBuf = [&self.source_include, "dummy.h"].iter().collect();
         let mut target_path: PathBuf = [&self.target_include, "dummy.h"].iter().collect();

builder/src/profiling.rs

@@ -36,4 +36,5 @@ libdd-shared-runtime = { version = "1.0.0", path = "../libdd-shared-runtime" }
 libdd-common-ffi = { path = "../libdd-common-ffi", default-features = false }
 libdd-tinybytes = { path = "../libdd-tinybytes" }
 libdd-trace-utils = { path = "../libdd-trace-utils" }
+tokio-util = "0.7.11"
 tracing = { version = "0.1", default-features = false }

libdd-data-pipeline-ffi/Cargo.toml

@@ -5,3 +5,7 @@ C FFI bindings for the libdd-data-pipeline library.
 ## Overview
 
 `libdd-data-pipeline-ffi` provides C-compatible FFI bindings for `libdd-data-pipeline`, enabling high-performance trace processing from C, C++, PHP, Ruby, Python, and other languages.
+
+## Dependencies
+
+This crate depends on `tokio-util` for its `CancellationToken` type. The cancellation token created by `ddog_trace_exporter_cancel_token_new` and passed to `ddog_trace_exporter_send_trace_chunks` is a `tokio_util::sync::CancellationToken`, which the data pipeline uses to cooperatively abort an in-flight send. The token is exposed opaquely to C, so callers never need to depend on `tokio-util` themselves.

libdd-data-pipeline-ffi/README.md

@@ -13,12 +13,13 @@ after_includes = """
 typedef struct ddog_TraceExporter ddog_TraceExporter;
 typedef struct ddog_TracerSpan ddog_TracerSpan;
 typedef struct ddog_TracerTraceChunks ddog_TracerTraceChunks;
+typedef struct ddog_TraceExporterCancelToken ddog_TraceExporterCancelToken;
 """
 
 [export]
 prefix = "ddog_"
 renaming_overrides_prefixing = true
-exclude = ["TraceExporter", "TracerSpan", "TracerTraceChunks"]
+exclude = ["TraceExporter", "TracerSpan", "TracerTraceChunks", "TokioCancellationToken"]
 
 [export.rename]
 "ByteSlice" = "ddog_ByteSlice"
@@ -32,6 +33,7 @@ exclude = ["TraceExporter", "TracerSpan", "TracerTraceChunks"]
 "TracerSpan" = "ddog_TracerSpan"
 "TracerSpanFields" = "ddog_TracerSpanFields"
 "TracerTraceChunks" = "ddog_TracerTraceChunks"
+"TokioCancellationToken" = "ddog_TraceExporterCancelToken"
 
 [export.mangle]
 rename_types = "PascalCase"

libdd-data-pipeline-ffi/cbindgen.toml

@@ -19,6 +19,8 @@ use libdd_tinybytes::BytesString;
 use libdd_trace_utils::span::v04::SpanBytes;
 use std::ptr::NonNull;
 
+type TokioCancellationToken = tokio_util::sync::CancellationToken;
+
 // ---------------------------------------------------------------------------
 // Helper
 // ---------------------------------------------------------------------------
@@ -295,6 +297,49 @@ pub unsafe extern "C" fn ddog_tracer_trace_chunks_push_span(
     )
 }
 
+// ---------------------------------------------------------------------------
+// Cancellation token
+// ---------------------------------------------------------------------------
+
+/// Create a new cancellation token.
+///
+/// The returned token must be freed with
+/// [`ddog_trace_exporter_cancel_token_drop`].
+#[no_mangle]
+pub extern "C" fn ddog_trace_exporter_cancel_token_new() -> Box<TokioCancellationToken> {
+    Box::new(TokioCancellationToken::new())
+}
+
+/// Cancel a cancellation token.
+///
+/// All clones of the same token observe the cancellation. If a
+/// [`ddog_trace_exporter_send_trace_chunks`] call is using this token at the
+/// time of cancellation, that send stops waiting for the agent at its next
+/// await point and returns an error; the trace chunks it was sending may be
+/// lost.
+///
+/// Cancellation only affects a send that is in progress. If no send is using
+/// the token, cancelling it has no immediate effect: a send started afterwards
+/// with an already-cancelled token returns an error without contacting the
+/// agent, and a token cancelled after its send has already finished does
+/// nothing.
+#[no_mangle]
+pub extern "C" fn ddog_trace_exporter_cancel_token_cancel(token: Option<&TokioCancellationToken>) {
+    if let Some(token) = token {
+        token.cancel();
+    }
+}
+
+/// Free a cancellation token.
+///
+/// After this call the token is invalid and must not be reused.
+#[no_mangle]
+pub extern "C" fn ddog_trace_exporter_cancel_token_drop(
+    token: Option<Box<TokioCancellationToken>>,
+) {
+    drop(token);
+}
+
 // ---------------------------------------------------------------------------
 // Send trace chunks
 // ---------------------------------------------------------------------------
@@ -305,13 +350,22 @@ pub unsafe extern "C" fn ddog_tracer_trace_chunks_push_span(
 /// serializes in the configured output format, and sends to the agent
 /// with retry logic.
 ///
+/// When `cancel` is non-null, cancelling that token (via
+/// [`ddog_trace_exporter_cancel_token_cancel`]) while the send is in progress
+/// aborts the in-flight request and returns an error with code
+/// [`ExporterErrorCode::IoError`]. Cancellation is cooperative: it only takes
+/// effect while a request is actually in flight. A token that is already
+/// cancelled when the send starts makes this function return that error
+/// immediately, and cancelling after the send has finished has no effect.
+/// Cancelling an in-flight send may cause the trace chunks being sent to be
+/// lost.
+///
 /// On success, if `response_out` is non-null, a heap-allocated
 /// [`ExporterResponse`] is written there.  The caller owns it and must
 /// free it with `ddog_trace_exporter_response_free`.
 ///
 /// # Safety
 ///
-/// * `exporter` must be a valid `TraceExporter` pointer.
 /// * `chunks` is consumed and must not be used after this call.
 /// * If `response_out` is non-null it must point to valid writable memory for a
 ///   `Box<ExporterResponse>`.
@@ -320,6 +374,7 @@ pub unsafe extern "C" fn ddog_trace_exporter_send_trace_chunks(
     exporter: Option<&TraceExporter>,
     chunks: Option<Box<TracerTraceChunks>>,
     response_out: Option<NonNull<Box<ExporterResponse>>>,
+    cancel: Option<&TokioCancellationToken>,
 ) -> Option<Box<ExporterError>> {
     let Some(exporter) = exporter else {
         return gen_error!(ErrorCode::InvalidArgument);
@@ -329,7 +384,7 @@ pub unsafe extern "C" fn ddog_trace_exporter_send_trace_chunks(
     };
 
     catch_panic!(
-        match exporter.send_trace_chunks(chunks.0) {
+        match exporter.send_trace_chunks(chunks.0, cancel) {
             Ok(resp) => {
                 if let Some(out) = response_out {
                     out.as_ptr().write(Box::new(ExporterResponse::from(resp)));
@@ -651,7 +706,7 @@ mod tests {
     fn send_trace_chunks_null_exporter_returns_error() {
         unsafe {
             let chunks = make_chunks(0);
-            let err = ddog_trace_exporter_send_trace_chunks(None, Some(chunks), None);
+            let err = ddog_trace_exporter_send_trace_chunks(None, Some(chunks), None, None);
             assert!(err.is_some());
             assert_eq!(err.as_ref().unwrap().code, ErrorCode::InvalidArgument);
             ddog_trace_exporter_error_free(err);
@@ -687,4 +742,32 @@ mod tests {
             ddog_tracer_trace_chunks_free(chunks);
         }
     }
+
+    // -- Cancellation token -------------------------------------------------
+
+    #[test]
+    fn cancel_token_new_and_drop() {
+        let token = ddog_trace_exporter_cancel_token_new();
+        ddog_trace_exporter_cancel_token_drop(Some(token));
+    }
+
+    #[test]
+    fn cancel_token_cancel() {
+        let token = ddog_trace_exporter_cancel_token_new();
+        ddog_trace_exporter_cancel_token_cancel(Some(&token));
+        ddog_trace_exporter_cancel_token_drop(Some(token));
+    }
+
+    #[test]
+    fn send_trace_chunks_null_cancel_is_accepted() {
+        // Passing a null (None) cancel argument behaves like no cancellation.
+        unsafe {
+            let chunks = make_chunks(0);
+            let err = ddog_trace_exporter_send_trace_chunks(None, Some(chunks), None, None);
+            // exporter is None, so we get InvalidArgument, but no crash
+            // from the absent cancel argument.
+            assert!(err.is_some());
+            ddog_trace_exporter_error_free(err);
+        }
+    }
 }

libdd-data-pipeline-ffi/src/tracer.rs

@@ -52,6 +52,7 @@ use std::sync::{Arc, Once};
 use std::time::Duration;
 use std::{borrow::Borrow, str::FromStr};
 use tokio::task::JoinSet;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, warn};
 
 const INFO_ENDPOINT: &str = "/info";
@@ -477,14 +478,39 @@ impl<C: HttpClientCapability + SleepCapability + MaybeSend + Sync + 'static> Tra
     /// Send a list of trace chunks to the agent (or OTLP endpoint when configured).
     ///
     /// Sync facade over [`Self::send_trace_chunks_async`]; panics inside an existing
-    /// tokio context. Returns the agent response (or `Unchanged` for OTLP).
+    /// tokio context.
+    ///
+    /// # Arguments
+    /// * trace_chunks: A list of trace chunks. Each trace chunk is a list of spans.
+    /// * cancellation_token: When provided, cancelling the token aborts the send while it is in
+    ///   progress. The send only observes a token that is cancelled while the request is in-flight;
+    ///   a token cancelled before this call returns immediately, and a token cancelled after the
+    ///   send has already finished has no effect. Cancelling an in-flight send may cause the trace
+    ///   chunks being sent to be lost.
+    ///
+    /// # Returns
+    /// * Ok(AgentResponse): The response from the agent (or Unchanged for OTLP)
+    /// * Err(TraceExporterError): An error detailing what went wrong in the process
     #[cfg(not(target_arch = "wasm32"))]
     pub fn send_trace_chunks<T: TraceData>(
         &self,
         trace_chunks: Vec<Vec<Span<T>>>,
+        cancellation_token: Option<&CancellationToken>,
     ) -> Result<AgentResponse, TraceExporterError> {
-        self.shared_runtime
-            .block_on(self.send_trace_chunks_async(trace_chunks))?
+        self.shared_runtime.block_on(async {
+            match cancellation_token {
+                Some(token) => {
+                    tokio::select! {
+                        res = self.send_trace_chunks_async(trace_chunks) => res,
+                        _ = token.cancelled() => Err(TraceExporterError::Io(std::io::Error::new(
+                            std::io::ErrorKind::Interrupted,
+                            "send cancelled via cancellation token",
+                        ))),
+                    }
+                }
+                None => self.send_trace_chunks_async(trace_chunks).await,
+            }
+        })?
     }
 
     /// Send a list of trace chunks to the agent, asynchronously (or OTLP when configured).

libdd-data-pipeline/src/trace_exporter/mod.rs

@@ -18,7 +18,7 @@ name = "datadog_profiling_ffi"
 
 [features]
 default = ["ddcommon-ffi"]
-cbindgen = ["build_common/cbindgen", "libdd-common-ffi/cbindgen"]
+cbindgen = ["build_common/cbindgen", "libdd-common-ffi/cbindgen", "libdd-shared-runtime-ffi?/cbindgen"]
 ddtelemetry-ffi = ["dep:libdd-telemetry-ffi"]
 datadog-log-ffi = ["dep:libdd-log-ffi"]
 symbolizer = ["symbolizer-ffi"]
@@ -33,6 +33,7 @@ datadog-library-config-ffi = ["dep:libdd-library-config-ffi"]
 ddcommon-ffi = ["dep:libdd-common-ffi"]
 ddsketch-ffi = ["dep:libdd-ddsketch-ffi"]
 datadog-ffe-ffi = ["dep:datadog-ffe-ffi"]
+shared-runtime = ["dep:libdd-shared-runtime-ffi", "libdd-shared-runtime-ffi/catch_panic"]
 regex-lite = ["libdd-common/regex-lite"]
 
 [build-dependencies]
@@ -60,3 +61,4 @@ symbolizer-ffi = { path = "../symbolizer-ffi", optional = true, default-features
 thiserror = "2"
 tokio-util = "0.7.1"
 datadog-ffe-ffi = { path = "../datadog-ffe-ffi", default-features = false, optional = true }
+libdd-shared-runtime-ffi = { path = "../libdd-shared-runtime-ffi", default-features = false, optional = true }