fix(shared-runtime): prevent stale MicroVM runtime data

Author: litianningdatadog

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

@@ -901,6 +901,42 @@ pub unsafe extern "C" fn ddog_trace_exporter_free(handle: Box<TraceExporter>) {
     let _ = catch_panic!(handle.shutdown(None), Ok(()));
 }
 
+/// Discards buffered data without sending it, then frees the TraceExporter instance.
+///
+/// This is intended for runtime identity refreshes. Unlike
+/// [`ddog_trace_exporter_free`], it does not execute the normal flush-capable shutdown path.
+///
+/// Returns `None` when every worker has been discarded and sets `*handle` to null. On an ordinary
+/// error, the exporter remains owned by `*handle` so the caller can retry after the fork lifecycle
+/// has completed. A caught panic sets `*handle` to null.
+///
+/// # Arguments
+///
+/// * handle - A non-null pointer to the TraceExporter handle. The pointee must be non-null.
+#[no_mangle]
+pub unsafe extern "C" fn ddog_trace_exporter_free_without_flush(
+    handle: &mut *mut TraceExporter,
+) -> Option<Box<ExporterError>> {
+    catch_panic!(
+        {
+            let Some(exporter) = NonNull::new(*handle) else {
+                return gen_error!(ErrorCode::InvalidArgument);
+            };
+            let mut exporter = Box::from_raw(exporter.as_ptr());
+            *handle = std::ptr::null_mut();
+
+            match exporter.shutdown_without_flush() {
+                Ok(()) => None,
+                Err(err) => {
+                    *handle = Box::into_raw(exporter);
+                    Some(Box::new(ExporterError::from(err)))
+                }
+            }
+        },
+        gen_error!(ErrorCode::Panic)
+    )
+}
+
 /// Send traces to the Datadog Agent.
 ///
 /// # Arguments
@@ -944,6 +980,7 @@ mod tests {
     use crate::error::ddog_trace_exporter_error_free;
     use httpmock::prelude::*;
     use httpmock::MockServer;
+    use libdd_shared_runtime::SharedRuntime;
     use libdd_trace_utils::span::v04::SpanSlice;
     use std::{borrow::Borrow, mem::MaybeUninit};
 
@@ -1351,6 +1388,46 @@ mod tests {
         }
     }
 
+    #[cfg_attr(miri, ignore)]
+    #[test]
+    fn exporter_free_without_flush_test() {
+        unsafe {
+            let mut config: MaybeUninit<Box<TraceExporterConfig>> = MaybeUninit::uninit();
+            ddog_trace_exporter_config_new(NonNull::new_unchecked(&mut config).cast());
+            let cfg = config.assume_init();
+
+            let mut exporter: MaybeUninit<Box<TraceExporter>> = MaybeUninit::uninit();
+            let error = ddog_trace_exporter_new(
+                NonNull::new_unchecked(&mut exporter).cast(),
+                Some(cfg.borrow()),
+            );
+            assert!(error.is_none());
+
+            let mut exporter = Box::into_raw(exporter.assume_init());
+            assert!(ddog_trace_exporter_free_without_flush(&mut exporter).is_none());
+            assert!(exporter.is_null());
+            ddog_trace_exporter_config_free(cfg);
+        }
+    }
+
+    #[test]
+    fn exporter_free_without_flush_preserves_handle_during_fork() {
+        unsafe {
+            let runtime = Arc::new(ForkSafeRuntime::new().unwrap());
+            let mut builder = TraceExporter::builder();
+            builder.set_shared_runtime(runtime.clone());
+            let mut exporter = Box::into_raw(Box::new(builder.build().unwrap()));
+
+            runtime.before_fork();
+            assert!(ddog_trace_exporter_free_without_flush(&mut exporter).is_some());
+            assert!(!exporter.is_null());
+
+            runtime.after_fork_parent().unwrap();
+            assert!(ddog_trace_exporter_free_without_flush(&mut exporter).is_none());
+            assert!(exporter.is_null());
+        }
+    }
+
     #[test]
     fn exporter_send_test_arguments_test() {
         unsafe {

libdd-data-pipeline/src/otlp/metrics.rs

@@ -348,6 +348,10 @@ impl<C: HttpClientCapability + SleepCapability + MaybeSend + Sync + 'static> Wor
             .flush(SystemTime::now(), true);
     }
 
+    fn discard(&mut self) {
+        self.reset();
+    }
+
     async fn shutdown(&mut self) {
         // Single attempt: a long backoff could miss the bounded shutdown window.
         if let Err(e) = self.send(true, OTLP_SHUTDOWN_MAX_RETRIES).await {
@@ -487,6 +491,55 @@ mod tests {
         str_at(p["attributes"].as_array().unwrap(), "status.code") == Some(STATUS_CODE_ERROR)
     }
 
+    #[test]
+    fn discard_clears_buffered_stats_without_export() {
+        use crate::otlp::config::OtlpProtocol;
+        use http::HeaderMap;
+        use libdd_capabilities_impl::NativeCapabilities;
+        use libdd_trace_utils::span::{trace_utils, v04::SpanSlice};
+        use std::borrow::Cow;
+
+        let concentrator = Arc::new(Mutex::new(SpanConcentrator::new(
+            Duration::from_secs(10),
+            SystemTime::now(),
+            vec![],
+            vec![],
+            None,
+            vec![],
+            #[cfg(feature = "stats-obfuscation")]
+            None,
+        )));
+        let mut spans = vec![SpanSlice {
+            service: Cow::Borrowed("svc"),
+            duration: 1,
+            ..Default::default()
+        }];
+        trace_utils::compute_top_level_span(&mut spans);
+        concentrator.lock_or_panic().add_span(&spans[0]);
+
+        let mut exporter = OtlpStatsExporter {
+            flush_interval: Duration::from_secs(10),
+            concentrator: concentrator.clone(),
+            config: OtlpMetricsConfig {
+                endpoint_url: "http://127.0.0.1:1/v1/metrics".to_string(),
+                headers: HeaderMap::new(),
+                timeout: Duration::from_millis(1),
+                protocol: OtlpProtocol::HttpJson,
+                otel_trace_semantics_enabled: false,
+            },
+            resource: OtlpResourceInfo::default(),
+            test_token: None,
+            capabilities: NativeCapabilities::new_client(),
+        };
+
+        exporter.discard();
+
+        assert!(concentrator
+            .lock_or_panic()
+            .flush_with_otlp_exact(SystemTime::now(), true)
+            .is_empty());
+    }
+
     #[test]
     fn metric_shape_and_resource_attributes() {
         assert!(map_stats_to_otlp_metrics(&[], &resource()).is_none());

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

@@ -958,7 +958,7 @@ mod tests {
     use std::sync::Arc;
     use std::time::{Duration, Instant};
 
-    use libdd_shared_runtime::{BlockingRuntime, ForkSafeRuntime, SharedRuntime};
+    use libdd_shared_runtime::{BlockingRuntime, ForkSafeRuntime, SharedRuntime, Worker};
 
     use crate::trace_buffer::{BufferSize, Export, TraceBuffer, TraceBufferConfig};
     use crate::trace_exporter::agent_response::AgentResponse;
@@ -1550,4 +1550,22 @@ mod tests {
         assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 2);
         rt.shutdown(None).unwrap();
     }
+
+    #[test]
+    fn test_worker_discard_drops_buffered_chunk() {
+        let (sender, mut worker) = TraceBuffer::new(
+            TraceBufferConfig::default().flush_threshold_bytes(2),
+            Box::new(|_| {}),
+            Box::new(AssertExporter(
+                Box::new(|_| panic!("discard must not export buffered chunks")),
+                Arc::new(tokio::sync::Semaphore::new(0)),
+            )),
+        );
+
+        sender.send_chunk(vec![()]).unwrap();
+        assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 1);
+
+        worker.discard();
+        assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 0);
+    }
 }