fix(shared-runtime): discard inherited MicroVM runtime data

.gitignore

@@ -15,7 +15,7 @@ release
 target
 tmp
 __fuzz__
-
+/graphify-out
 #==============================================================================#
 # Files to ignore
 #==============================================================================#

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

@@ -901,6 +901,44 @@ pub unsafe extern "C" fn ddog_trace_exporter_free(handle: Box<TraceExporter>) {
     let _ = catch_panic!(handle.shutdown(None), Ok(()));
 }
 
+/// Discards runtime-identity-bound state without sending buffered data, then frees the exporter.
+///
+/// Why this exists: SDKs that restore from a process snapshot, such as dd-trace-py in a MicroVM
+/// `/run` hook, can inherit buffered stats, telemetry state, and cached agent `/info` responses
+/// from the snapshotted runtime. The normal [`ddog_trace_exporter_free`] path flushes before
+/// freeing, which is correct for ordinary shutdown but wrong when the SDK is discarding inherited
+/// state and immediately constructing a fresh exporter for the restored runtime identity.
+///
+/// This function gives those snapshot-aware callers an explicit no-flush teardown API. Ordinary
+/// non-snapshot shutdown callers should continue using [`ddog_trace_exporter_free`].
+///
+/// Returns `None` when the inherited state was discarded. Once a non-null handle is accepted,
+/// `*handle` is always set to null, including when an ordinary error is returned.
+///
+/// # 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 exporter = Box::from_raw(exporter.as_ptr());
+            *handle = std::ptr::null_mut();
+
+            match exporter.shutdown_without_flush() {
+                Ok(()) => None,
+                Err(err) => Some(Box::new(ExporterError::from(err))),
+            }
+        },
+        gen_error!(ErrorCode::Panic)
+    )
+}
+
 /// Send traces to the Datadog Agent.
 ///
 /// # Arguments
@@ -1351,6 +1389,28 @@ 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_send_test_arguments_test() {
         unsafe {

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

@@ -1554,8 +1554,8 @@ mod tests {
             assert!(event.0.attributes.is_empty());
 
             // Misaligned pointer
-            let buf = [0u8; 32];
-            let misaligned = buf.as_ptr().add(1).cast::<i64>();
+            let data = [0i64; 3];
+            let misaligned = data.as_ptr().cast::<u8>().add(1).cast::<i64>();
             assert!(!misaligned.is_aligned());
             let err = ddog_tracer_span_event_set_int_array(
                 Some(&mut event),

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

@@ -29,12 +29,17 @@ pub fn get_agent_info() -> Option<Arc<schema::AgentInfo>> {
     AGENT_INFO_CACHE.load_full()
 }
 
+/// Clear the global agent info cache after restoring from an inherited runtime identity.
+pub(crate) fn clear_cache_for_runtime_identity_refresh() {
+    AGENT_INFO_CACHE.store(None);
+}
+
 pub use fetcher::{
     fetch_info, fetch_info_with_state, AgentInfoFetcher, FetchInfoStatus, ResponseObserver,
 };
 
 #[cfg(test)]
 /// Clear the global agent info cache for test isolation in envs where nextest isn't used.
 pub fn clear_cache_for_test() {
-    AGENT_INFO_CACHE.store(None);
+    clear_cache_for_runtime_identity_refresh();
 }

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

@@ -634,6 +634,20 @@ impl<T> Receiver<T> {
         Ok(())
     }
 
+    /// Permanently close the channel, dropping any buffered chunks without exporting them.
+    ///
+    /// Unlike [`Self::reset`], which reopens the channel for a worker that keeps running (e.g.
+    /// after a fork), this is for a worker that is never run again. Leaving the channel
+    /// `Running` would let a surviving [`Sender`] keep queuing chunks that no receiver will ever
+    /// take, and would leave callers blocked forever in [`Sender::wait_close_done`] or
+    /// [`Sender::wait_flush_done`].
+    fn discard(&self) -> Result<(), MutexPoisonedError> {
+        let mut state = self.lock_state()?;
+        state.batch.reset();
+        self.waiter.mark_stopped(state);
+        Ok(())
+    }
+
     async fn receive(&self, timeout: Duration) -> Result<Vec<TraceChunk<T>>, MutexPoisonedError> {
         loop {
             // Enable the notify future BEFORE acquiring the lock to avoid lost wakeups:
@@ -950,6 +964,13 @@ impl<T: Send + Debug + 'static> Worker for TraceExporterWorker<T> {
     fn reset(&mut self) {
         let _ = self.rx.reset();
     }
+
+    // Override the default (`Worker::discard`'s `reset()` forwarding): `reset()` reopens the
+    // channel for continued use, which is wrong once this worker is permanently discarded (see
+    // `Receiver::discard`).
+    fn discard(&mut self) {
+        let _ = self.rx.discard();
+    }
 }
 
 #[cfg(test)]
@@ -958,7 +979,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 +1571,46 @@ mod tests {
         assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 2);
         rt.shutdown(None).unwrap();
     }
+
+    #[test]
+    fn test_worker_reset_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.reset();
+        assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 0);
+    }
+
+    #[test]
+    fn test_worker_discard_closes_channel_without_exporting() {
+        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();
+
+        worker.discard();
+
+        // Unlike `reset`, which reopens the channel for a worker that keeps running (e.g. after
+        // a fork), `discard` is for a worker that is never run again. A channel left `Running`
+        // would let this sender keep queuing chunks that no receiver will ever take.
+        assert!(matches!(
+            sender.send_chunk(vec![()]),
+            Err(TraceBufferError::AlreadyClosed)
+        ));
+    }
 }