feat(trace_exporter)!: add restart_after_fork configuration option [APMSP-3843]

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

@@ -8,6 +8,7 @@ use super::config::OtlpMetricsConfig;
 use super::exporter::{send_otlp_http, OTLP_MAX_RETRIES, OTLP_SHUTDOWN_MAX_RETRIES};
 use async_trait::async_trait;
 use libdd_capabilities::{HttpClientCapability, MaybeSend, SleepCapability};
+use libdd_common::MutexExt;
 use libdd_ddsketch::DDSketch;
 use libdd_shared_runtime::Worker;
 use libdd_trace_protobuf::pb;
@@ -257,8 +258,7 @@ impl<C: HttpClientCapability + SleepCapability> OtlpStatsExporter<C> {
     /// Flush the concentrator and export stats; returns `Ok(true)` if anything was sent.
     async fn send(&self, force_flush: bool, max_retries: u32) -> anyhow::Result<bool> {
         let buckets = {
-            #[allow(clippy::unwrap_used)]
-            let mut c = self.concentrator.lock().unwrap();
+            let mut c = self.concentrator.lock_or_panic();
             c.flush_with_otlp_exact(SystemTime::now(), force_flush)
         };
         if buckets.is_empty() {
@@ -301,6 +301,13 @@ impl<C: HttpClientCapability + SleepCapability + MaybeSend + Sync + 'static> Wor
         }
     }
 
+    fn reset(&mut self) {
+        let _ = self
+            .concentrator
+            .lock_or_panic()
+            .flush(SystemTime::now(), true);
+    }
+
     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 {

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

@@ -101,6 +101,9 @@ pub struct TraceExporterBuilder<R: SharedRuntime> {
     output_to_log: bool,
     /// Optional override for the maximum size of a single emitted log line.
     log_max_line_size: Option<usize>,
+    /// Whether background workers spawned should be
+    /// restarted in the child after a `fork()`. Defaults to `true`.
+    restart_after_fork: bool,
 }
 
 /// Default is impl'd for `R = ForkSafeRuntime` only so that bare
@@ -167,6 +170,7 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
             agentless_timeout: None,
             output_to_log: false,
             log_max_line_size: None,
+            restart_after_fork: true,
         }
     }
 }
@@ -344,6 +348,18 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
         self
     }
 
+    /// Sets whether background workers spawned by this exporter  are restarted in the child after
+    /// the process `fork()`s. Defaults to `true`.
+    ///
+    /// Set to `false` if the trace exporter is recreated after the fork to avoid restarting workers
+    /// which are going to be discarded anyway.
+    ///
+    /// TODO(APMSP-3846): Remove once python no longer recreates the exporter on forks.
+    pub fn set_restart_after_fork(&mut self, restart_after_fork: bool) -> &mut Self {
+        self.restart_after_fork = restart_after_fork;
+        self
+    }
+
     /// Enable client-side stats obfuscation. Disabled by default.
     ///
     /// Final activation also requires the agent to advertise a supported
@@ -620,7 +636,7 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
         } else {
             Some(
                 shared_runtime
-                    .spawn_worker(info_fetcher, false)
+                    .spawn_worker(info_fetcher, self.restart_after_fork)
                     .map_err(|e| {
                         TraceExporterError::Builder(BuilderErrorKind::InvalidConfiguration(
                             e.to_string(),
@@ -675,11 +691,13 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
                 .transpose()?;
             match telemetry {
                 Some((client_tel, worker)) => {
-                    let handle = shared_runtime.spawn_worker(worker, false).map_err(|e| {
-                        TraceExporterError::Builder(BuilderErrorKind::InvalidConfiguration(
-                            e.to_string(),
-                        ))
-                    })?;
+                    let handle = shared_runtime
+                        .spawn_worker(worker, self.restart_after_fork)
+                        .map_err(|e| {
+                            TraceExporterError::Builder(BuilderErrorKind::InvalidConfiguration(
+                                e.to_string(),
+                            ))
+                        })?;
                     if let Err(e) = client_tel.start() {
                         tracing::warn!("Failed to start telemetry: {e}");
                     }
@@ -775,9 +793,13 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
                 test_token: self.test_session_token.clone(),
                 capabilities: capabilities.clone(),
             };
-            let worker_handle = shared_runtime.spawn_worker(worker, false).map_err(|e| {
-                TraceExporterError::Builder(BuilderErrorKind::InvalidConfiguration(e.to_string()))
-            })?;
+            let worker_handle = shared_runtime
+                .spawn_worker(worker, self.restart_after_fork)
+                .map_err(|e| {
+                    TraceExporterError::Builder(BuilderErrorKind::InvalidConfiguration(
+                        e.to_string(),
+                    ))
+                })?;
             stats = StatsComputationStatus::Enabled {
                 stats_concentrator: concentrator,
                 worker_handle,
@@ -853,6 +875,7 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
             trace_filterer: ArcSwap::from_pointee(TraceFilterer::with_empty_conf()),
             otlp_stats_enabled,
             log_output,
+            restart_after_fork: self.restart_after_fork,
         })
     }
 
@@ -993,7 +1016,8 @@ mod tests {
             .set_otlp_instrumentation_scope("dd-trace-js", "7.0.0-pre")
             .set_input_format(TraceExporterInputFormat::V04)
             .set_output_format(TraceExporterOutputFormat::V04)
-            .set_client_computed_stats();
+            .set_client_computed_stats()
+            .set_restart_after_fork(false);
         #[cfg(feature = "telemetry")]
         builder.enable_telemetry(TelemetryConfig {
             heartbeat: 1000,
@@ -1020,6 +1044,7 @@ mod tests {
         let otlp_config = exporter.otlp_config.as_ref().unwrap();
         assert_eq!(otlp_config.instrumentation_scope_name, "dd-trace-js");
         assert_eq!(otlp_config.instrumentation_scope_version, "7.0.0-pre");
+        assert!(!exporter.restart_after_fork);
         #[cfg(feature = "telemetry")]
         assert!(exporter.telemetry.is_some());
     }
@@ -1043,6 +1068,7 @@ mod tests {
         assert_eq!(exporter.metadata.language_version, "");
         assert_eq!(exporter.metadata.language_interpreter, "");
         assert!(!exporter.metadata.client_computed_stats);
+        assert!(exporter.restart_after_fork);
         #[cfg(feature = "telemetry")]
         assert!(exporter.telemetry.is_none());
     }

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

@@ -279,6 +279,8 @@ pub struct TraceExporter<
     /// path) instead of being sent to an agent. Used in serverless environments
     /// where no agent is reachable.
     log_output: Option<usize>,
+    /// Whether background workers should be restarted in the child after a `fork()`.
+    restart_after_fork: bool,
 }
 
 impl<
@@ -460,6 +462,7 @@ impl<
                     endpoint_url: &self.endpoint.url,
                     shared_runtime: &*self.shared_runtime,
                     stats_cardinality_limit: self.client_side_stats.stats_cardinality_limit,
+                    restart_after_fork: self.restart_after_fork,
                     dogstatsd: if self.health_metrics_enabled {
                         self.dogstatsd.clone()
                     } else {

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

@@ -48,6 +48,8 @@ pub(crate) struct StatsContext<
     pub endpoint_url: &'a http::Uri,
     pub shared_runtime: &'a R,
     pub stats_cardinality_limit: Option<usize>,
+    /// Configuration option to pass to [SharedRuntime::spawn_worker]
+    pub restart_after_fork: bool,
     /// Optional DogStatsD client forwarded to the [`StatsExporter`].
     pub dogstatsd: Option<libdd_dogstatsd_client::DogStatsDClient>,
     /// Optional telemetry handle forwarded to the [`StatsExporter`].
@@ -172,7 +174,7 @@ fn create_and_start_stats_worker<
     );
     let worker_handle = ctx
         .shared_runtime
-        .spawn_worker(stats_exporter, false)
+        .spawn_worker(stats_exporter, ctx.restart_after_fork)
         .map_err(|e| anyhow::anyhow!(e))?;
 
     // Update the stats computation state with the new worker components.

libdd-trace-stats/src/stats_exporter.rs

@@ -14,7 +14,7 @@ use async_trait::async_trait;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt as _;
 use libdd_capabilities::{HttpClientCapability, MaybeSend, SleepCapability};
-use libdd_common::Endpoint;
+use libdd_common::{Endpoint, MutexExt};
 use libdd_shared_runtime::Worker;
 use libdd_trace_protobuf::pb;
 use libdd_trace_utils::send_with_retry::{send_with_retry, RetryBackoffType, RetryStrategy};
@@ -178,8 +178,7 @@ impl<
     /// Returns `Ok(true)` if stats were sent, `Ok(false)` if the concentrator had nothing to send.
     pub async fn send(&self, force_flush: bool) -> anyhow::Result<bool> {
         let flush = {
-            #[allow(clippy::unwrap_used)]
-            let mut concentrator = self.concentrator.lock().unwrap();
+            let mut concentrator = self.concentrator.lock_or_panic();
             concentrator.flush_buckets(force_flush)
         };
 
@@ -280,6 +279,11 @@ impl<
         let _ = self.send(false).await; // bool return ignored by Worker
     }
 
+    fn reset(&mut self) {
+        let _ = self.concentrator.lock_or_panic().flush_buckets(true);
+        self.sequence_id.store(0, Ordering::Relaxed);
+    }
+
     async fn shutdown(&mut self) {
         let _ = self.send(true).await;
     }