fix(data-pipeline): Apply suggestions

Author: VianneyRuhlmann

data-pipeline/src/agent_info/fetcher.rs

@@ -6,8 +6,7 @@
 use super::{schema::AgentInfo, AgentInfoArc};
 use anyhow::{anyhow, Result};
 use arc_swap::ArcSwapOption;
-use ddcommon::hyper_migration;
-use ddcommon::Endpoint;
+use ddcommon::{hyper_migration, worker::Worker, Endpoint};
 use http_body_util::BodyExt;
 use hyper::{self, body::Buf, header::HeaderName};
 use log::{error, info};
@@ -140,11 +139,20 @@ impl AgentInfoFetcher {
         }
     }
 
+    /// Return an AgentInfoArc storing the info received by the agent.
+    ///
+    /// When the fetcher is running it updates the AgentInfoArc when the agent's info changes.
+    pub fn get_info(&self) -> AgentInfoArc {
+        self.info.clone()
+    }
+}
+
+impl Worker for AgentInfoFetcher {
     /// Start fetching the info endpoint with the given interval.
     ///
     /// # Warning
     /// This method does not return and should be called within a dedicated task.
-    pub async fn run(&self) {
+    async fn run(&mut self) {
         loop {
             let current_info = self.info.load();
             let current_hash = current_info.as_ref().map(|info| info.state_hash.as_str());
@@ -164,13 +172,6 @@ impl AgentInfoFetcher {
             sleep(self.refresh_interval).await;
         }
     }
-
-    /// Return an AgentInfoArc storing the info received by the agent.
-    ///
-    /// When the fetcher is running it updates the AgentInfoArc when the agent's info changes.
-    pub fn get_info(&self) -> AgentInfoArc {
-        self.info.clone()
-    }
 }
 
 #[cfg(test)]
@@ -329,7 +330,7 @@ mod tests {
             })
             .await;
         let endpoint = Endpoint::from_url(server.url("/info").parse().unwrap());
-        let fetcher = AgentInfoFetcher::new(endpoint.clone(), Duration::from_millis(100));
+        let mut fetcher = AgentInfoFetcher::new(endpoint.clone(), Duration::from_millis(100));
         let info = fetcher.get_info();
         assert!(info.load().is_none());
         tokio::spawn(async move {

data-pipeline/src/pausable_worker.rs

@@ -2,42 +2,29 @@
 // SPDX-License-Identifier: Apache-2.0
 
 //! Defines a pausable worker to be able to stop background processes before forks
-use anyhow::{anyhow, Result};
-use ddtelemetry::worker::TelemetryWorker;
-use tokio::{runtime::Runtime, select, task::JoinHandle};
-use tokio_util::sync::CancellationToken;
-
-use crate::{agent_info::AgentInfoFetcher, stats_exporter::StatsExporter};
-
-/// Trait representing a worker which can be wrapped by `PausableWorker`
-pub trait Worker {
-    /// Main worker loop
-    fn run(&mut self) -> impl std::future::Future<Output = ()> + Send;
-}
-
-impl Worker for StatsExporter {
-    async fn run(&mut self) {
-        Self::run(self).await
-    }
-}
 
-impl Worker for AgentInfoFetcher {
-    async fn run(&mut self) {
-        Self::run(self).await
-    }
-}
-
-impl Worker for TelemetryWorker {
-    async fn run(&mut self) {
-        Self::run(self).await
-    }
-}
+use ddcommon::worker::Worker;
+use std::fmt::Display;
+use tokio::{
+    runtime::Runtime,
+    select,
+    task::{JoinError, JoinHandle},
+};
+use tokio_util::sync::CancellationToken;
 
-/// A pausable worker which can be paused and restarded on forks.
+/// A pausable worker which can be paused and restarted on forks.
+///
+/// Used to allow a [`ddcommon::worker::Worker`] to be paused while saving its state when dropping
+/// a tokio runtime to be able to restart with the same state on a new runtime. This is used to
+/// stop all threads before a fork to avoid deadlocks in child.
 ///
-/// # Requirements
-/// When paused the worker will exit on the next awaited call. To be able to safely restart the
-/// worker must be in a valid state on every call to `.await`.
+/// # Time-to-pause
+/// This loop should yield regularly to reduce time-to-pause. See [`tokio::task::yield_now`].
+///
+/// # Cancellation safety
+/// The main loop can be interrupted at any yield point (`.await`ed call). The state of the worker
+/// at this point will be saved and used to restart the worker. To be able to safely restart, the
+/// worker must be in a valid state on every call to `.await`. See [`tokio::select#cancellation-safety`] for more details.
 #[derive(Debug)]
 pub enum PausableWorker<T: Worker + Send + Sync + 'static> {
     Running {
@@ -50,6 +37,25 @@ pub enum PausableWorker<T: Worker + Send + Sync + 'static> {
     InvalidState,
 }
 
+#[derive(Debug)]
+pub enum PausableWorkerError {
+    InvalidState,
+    TaskAborted,
+}
+
+impl Display for PausableWorkerError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PausableWorkerError::InvalidState => {
+                write!(f, "Worker is in an invalid state and must be recreated.")
+            }
+            PausableWorkerError::TaskAborted => {
+                write!(f, "Worker task has been aborted and state has been lost.")
+            }
+        }
+    }
+}
+
 impl<T: Worker + Send + Sync + 'static> PausableWorker<T> {
     /// Create a new pausable worker from the given worker.
     pub fn new(worker: T) -> Self {
@@ -58,13 +64,15 @@ impl<T: Worker + Send + Sync + 'static> PausableWorker<T> {
 
     /// Start the worker on the given runtime.
     ///
+    /// The worker's main loop will be run on the runtime.
+    ///
     /// # Errors
     /// Fails if the worker is in an invalid state.
-    pub fn start(&mut self, rt: &Runtime) -> Result<()> {
+    pub fn start(&mut self, rt: &Runtime) -> Result<(), PausableWorkerError> {
         if let Self::Running { .. } = self {
             Ok(())
         } else if let Self::Paused { mut worker } = std::mem::replace(self, Self::InvalidState) {
-            // Worker is temporarly in an invalid state, but since this block is failsafe it will
+            // Worker is temporarily in an invalid state, but since this block is failsafe it will
             // be replaced by a valid state.
             let stop_token = CancellationToken::new();
             let cloned_token = stop_token.clone();
@@ -78,34 +86,34 @@ impl<T: Worker + Send + Sync + 'static> PausableWorker<T> {
             *self = PausableWorker::Running { handle, stop_token };
             Ok(())
         } else {
-            Err(anyhow!("Failed to start service"))
+            Err(PausableWorkerError::InvalidState)
         }
     }
 
     /// Pause the worker saving it's state to be restarted.
     ///
     /// # Errors
     /// Fails if the worker handle has been aborted preventing the worker from being retrieved.
-    pub async fn pause(&mut self) -> Result<()> {
+    pub async fn pause(&mut self) -> Result<(), PausableWorkerError> {
         match self {
             PausableWorker::Running { handle, stop_token } => {
                 stop_token.cancel();
                 if let Ok(worker) = handle.await {
                     *self = PausableWorker::Paused { worker };
                     Ok(())
                 } else {
-                    // Worker isn't retrieved and can't be restarted.
+                    // The task has been aborted and the worker can't be retrieved.
                     *self = PausableWorker::InvalidState;
-                    Err(anyhow!("Failed to stop worker. Worker must be recreated"))
+                    Err(PausableWorkerError::TaskAborted)
                 }
             }
             PausableWorker::Paused { .. } => Ok(()),
-            PausableWorker::InvalidState => Err(anyhow!("Worker is in invalid state")),
+            PausableWorker::InvalidState => Err(PausableWorkerError::InvalidState),
         }
     }
 
     /// Wait for the run method of the worker to exit.
-    pub async fn join(self) -> Result<()> {
+    pub async fn join(self) -> Result<(), JoinError> {
         if let PausableWorker::Running { handle, .. } = self {
             handle.await?;
         }

data-pipeline/src/stats_exporter.rs

@@ -13,7 +13,7 @@ use std::{
 
 use datadog_trace_protobuf::pb;
 use datadog_trace_utils::send_with_retry::{send_with_retry, RetryStrategy};
-use ddcommon::Endpoint;
+use ddcommon::{worker::Worker, Endpoint};
 use hyper;
 use log::error;
 use tokio::select;
@@ -128,13 +128,15 @@ impl StatsExporter {
                 .flush(time::SystemTime::now(), force_flush),
         )
     }
+}
 
+impl Worker for StatsExporter {
     /// Run loop of the stats exporter
     ///
     /// Once started, the stats exporter will flush and send stats on every `self.flush_interval`.
     /// If the `self.cancellation_token` is cancelled, the exporter will force flush all stats and
     /// return.
-    pub async fn run(&self) {
+    async fn run(&mut self) {
         loop {
             select! {
                 _ = self.cancellation_token.cancelled() => {
@@ -316,7 +318,7 @@ mod tests {
             })
             .await;
 
-        let stats_exporter = StatsExporter::new(
+        let mut stats_exporter = StatsExporter::new(
             BUCKETS_DURATION,
             Arc::new(Mutex::new(get_test_concentrator())),
             get_test_metadata(),
@@ -356,7 +358,7 @@ mod tests {
         let buckets_duration = Duration::from_secs(10);
         let cancellation_token = CancellationToken::new();
 
-        let stats_exporter = StatsExporter::new(
+        let mut stats_exporter = StatsExporter::new(
             buckets_duration,
             Arc::new(Mutex::new(get_test_concentrator())),
             get_test_metadata(),

data-pipeline/src/telemetry/mod.rs

@@ -266,7 +266,7 @@ impl TelemetryClient {
 
 #[cfg(test)]
 mod tests {
-    use ddcommon::hyper_migration;
+    use ddcommon::{hyper_migration, worker::Worker};
     use httpmock::Method::POST;
     use httpmock::MockServer;
     use hyper::{Response, StatusCode};

data-pipeline/src/trace_exporter/mod.rs

@@ -395,7 +395,9 @@ impl TraceExporter {
             );
             let mut stats_worker = PausableWorker::new(stats_exporter);
             let runtime = self.runtime()?;
-            stats_worker.start(&runtime)?;
+            stats_worker.start(&runtime).map_err(|e| {
+                TraceExporterError::Internal(InternalErrorKind::InvalidWorkerState(e.to_string()))
+            })?;
 
             self.workers.lock_or_panic().stats = Some(stats_worker);
 

data-pipeline/tests/test_fetch_info.rs

@@ -6,7 +6,7 @@ mod tracing_integration_tests {
     use arc_swap::access::Access;
     use data_pipeline::agent_info::{fetch_info, AgentInfoFetcher};
     use datadog_trace_utils::test_utils::datadog_test_agent::DatadogTestAgent;
-    use ddcommon::Endpoint;
+    use ddcommon::{worker::Worker, Endpoint};
     use std::time::Duration;
 
     #[cfg_attr(miri, ignore)]
@@ -30,7 +30,7 @@ mod tracing_integration_tests {
     async fn test_agent_info_fetcher_with_test_agent() {
         let test_agent = DatadogTestAgent::new(None, None, &[]).await;
         let endpoint = Endpoint::from_url(test_agent.get_uri_for_endpoint("info", None).await);
-        let fetcher = AgentInfoFetcher::new(endpoint, Duration::from_secs(1));
+        let mut fetcher = AgentInfoFetcher::new(endpoint, Duration::from_secs(1));
         let info_arc = fetcher.get_info();
         tokio::spawn(async move { fetcher.run().await });
         let info_received = async {

ddcommon/src/lib.rs

@@ -26,6 +26,7 @@ pub mod rate_limiter;
 pub mod tag;
 pub mod tracer_metadata;
 pub mod unix_utils;
+pub mod worker;
 
 /// Extension trait for `Mutex` to provide a method that acquires a lock, panicking if the lock is
 /// poisoned.

ddcommon/src/worker.rs

@@ -0,0 +1,12 @@
+// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+/// Trait representing a generic worker.
+///
+/// The worker runs an async looping function running periodic tasks.
+///
+/// This trait can be used to provide wrapper around a worker.
+pub trait Worker {
+    /// Main worker loop
+    fn run(&mut self) -> impl std::future::Future<Output = ()> + Send;
+}