Move PausableWorker to a separate file

Author: VianneyRuhlmann

data-pipeline/src/lib.rs

@@ -12,6 +12,7 @@
 
 pub mod agent_info;
 mod health_metrics;
+pub mod pausable_worker;
 #[allow(missing_docs)]
 pub mod span_concentrator;
 #[allow(missing_docs)]

data-pipeline/src/pausable_worker.rs

@@ -0,0 +1,89 @@
+//! Defines a pausable worker to be able to stop background processes before forks
+use ddtelemetry::worker::TelemetryWorker;
+use tokio::{runtime::Runtime, select, task::JoinHandle};
+use tokio_util::sync::CancellationToken;
+
+use crate::{agent_info::AgentInfoFetcher, stats_exporter::StatsExporter};
+
+pub trait Worker {
+    /// Main worker loop
+    fn run(&mut self) -> impl std::future::Future<Output = ()> + Send;
+    /// Hook called on the worker when pausing.
+    /// The worker can be paused on any await call. This hook can be use to clean the worker
+    /// state after pausing.
+    fn on_pause(&mut self);
+}
+
+impl Worker for StatsExporter {
+    async fn run(&mut self) {
+        Self::run(self).await
+    }
+    fn on_pause(&mut self) {}
+}
+
+impl Worker for AgentInfoFetcher {
+    async fn run(&mut self) {
+        Self::run(self).await
+    }
+    fn on_pause(&mut self) {}
+}
+
+impl Worker for TelemetryWorker {
+    async fn run(&mut self) {
+        Self::run(self).await
+    }
+    fn on_pause(&mut self) {
+        self.cleanup();
+    }
+}
+
+/// A pausable worker which can be paused and restarded on forks.
+#[derive(Debug)]
+pub enum PausableWorker<T: Worker + Send + Sync + 'static> {
+    Running {
+        handle: JoinHandle<T>,
+        stop_token: CancellationToken,
+    },
+    Paused {
+        worker: T,
+    },
+    InvalidState,
+}
+
+impl<T: Worker + Send + Sync + 'static> PausableWorker<T> {
+    pub fn new(worker: T) -> Self {
+        Self::Paused { worker }
+    }
+
+    pub fn start(&mut self, rt: &Runtime) {
+        if let Self::Paused { mut worker } = std::mem::replace(self, Self::InvalidState) {
+            let stop_token = CancellationToken::new();
+            let cloned_token = stop_token.clone();
+            let handle = rt.spawn(async move {
+                select! {
+                    _ = worker.run() => {worker}
+                    _ = cloned_token.cancelled() => {worker}
+                }
+            });
+
+            *self = PausableWorker::Running { handle, stop_token };
+        }
+    }
+
+    pub async fn stop(&mut self) {
+        if let PausableWorker::Running { handle, stop_token } = self {
+            stop_token.cancel();
+            let worker = handle.await.unwrap();
+            worker.on_pause();
+            *self = PausableWorker::Paused { worker };
+        }
+    }
+
+    /// Wait for the run method of the worker to exit.
+    pub async fn join(&mut self) {
+        if let PausableWorker::Running { handle, stop_token } = self {
+            let worker = handle.await.unwrap();
+            *self = PausableWorker::Paused { worker };
+        }
+    }
+}

data-pipeline/src/trace_exporter/mod.rs

@@ -3,6 +3,7 @@
 pub mod agent_response;
 pub mod error;
 use crate::agent_info::{AgentInfoArc, AgentInfoFetcher};
+use crate::pausable_worker::PausableWorker;
 use crate::stats_exporter::StatsExporter;
 use crate::telemetry::{self, SendPayloadTelemetry, TelemetryClient, TelemetryClientBuilder};
 use crate::trace_exporter::error::{RequestError, TraceExporterError};
@@ -68,91 +69,6 @@ pub enum TraceExporterOutputFormat {
     V05,
 }
 
-mod pausable_worker {
-    use ddtelemetry::worker::TelemetryWorker;
-    use tokio::{runtime::Runtime, select, task::JoinHandle};
-    use tokio_util::sync::CancellationToken;
-
-    use crate::{agent_info::AgentInfoFetcher, stats_exporter::StatsExporter};
-
-    pub trait Worker {
-        fn run(&mut self) -> impl std::future::Future<Output = ()> + Send;
-        fn on_pause(&mut self);
-    }
-
-    impl Worker for StatsExporter {
-        async fn run(&mut self) {
-            Self::run(self).await
-        }
-        fn on_pause(&mut self) {}
-    }
-
-    impl Worker for AgentInfoFetcher {
-        async fn run(&mut self) {
-            Self::run(self).await
-        }
-        fn on_pause(&mut self) {}
-    }
-
-    impl Worker for TelemetryWorker {
-        async fn run(&mut self) {
-            Self::run(self).await
-        }
-        fn on_pause(&mut self) {
-            self.cleanup();
-        }
-    }
-
-    #[derive(Debug)]
-    pub enum PausableWorker<T: Worker + Send + Sync + 'static> {
-        Running {
-            handle: JoinHandle<T>,
-            stop_token: CancellationToken,
-        },
-        Paused {
-            worker: T,
-        },
-        InvalidState,
-    }
-
-    impl<T: Worker + Send + Sync + 'static> PausableWorker<T> {
-        pub fn new(worker: T) -> Self {
-            Self::Paused { worker }
-        }
-
-        pub fn start(&mut self, rt: &Runtime) {
-            if let Self::Paused { mut worker } = std::mem::replace(self, Self::InvalidState) {
-                let stop_token = CancellationToken::new();
-                let cloned_token = stop_token.clone();
-                let handle = rt.spawn(async move {
-                    select! {
-                        _ = worker.run() => {worker}
-                        _ = cloned_token.cancelled() => {worker}
-                    }
-                });
-
-                *self = PausableWorker::Running { handle, stop_token };
-            }
-        }
-
-        pub async fn stop(&mut self) {
-            if let PausableWorker::Running { handle, stop_token } = self {
-                stop_token.cancel();
-                let worker = handle.await.unwrap();
-                *self = PausableWorker::Paused { worker };
-            }
-        }
-
-        /// Wait for the run method of the worker to exit.
-        pub async fn join(&mut self) {
-            if let PausableWorker::Running { handle, stop_token } = self {
-                let worker = handle.await.unwrap();
-                *self = PausableWorker::Paused { worker };
-            }
-        }
-    }
-}
-
 impl TraceExporterOutputFormat {
     /// Add the agent trace endpoint path to the URL.
     fn add_path(&self, url: &Uri) -> Uri {
@@ -248,8 +164,6 @@ enum StatsComputationStatus {
     },
 }
 
-use pausable_worker::PausableWorker;
-
 #[derive(Debug)]
 struct TraceExporterWorkers {
     pub info: PausableWorker<AgentInfoFetcher>,