Add pausable worker for telemetry

Author: VianneyRuhlmann

data-pipeline/src/telemetry/mod.rs

@@ -12,11 +12,11 @@ use datadog_trace_utils::{
 };
 use ddcommon::tag::Tag;
 use ddtelemetry::worker::{
-    LifecycleAction, TelemetryActions, TelemetryWorkerBuilder, TelemetryWorkerFlavor,
-    TelemetryWorkerHandle,
+    LifecycleAction, TelemetryActions, TelemetryWorker, TelemetryWorkerBuilder,
+    TelemetryWorkerFlavor, TelemetryWorkerHandle,
 };
 use std::{collections::HashMap, time::Duration};
-use tokio::task::JoinHandle;
+use tokio::runtime::Handle;
 
 /// Structure to build a Telemetry client.
 ///
@@ -86,7 +86,10 @@ impl TelemetryClientBuilder {
     }
 
     /// Builds the telemetry client.
-    pub async fn build(self) -> Result<TelemetryClient, TelemetryError> {
+    pub fn build(
+        self,
+        runtime: Handle,
+    ) -> Result<(TelemetryClient, TelemetryWorker), TelemetryError> {
         #[allow(clippy::unwrap_used)]
         let mut builder = TelemetryWorkerBuilder::new_fetch_host(
             self.service_name.unwrap(),
@@ -102,16 +105,17 @@ impl TelemetryClientBuilder {
             builder.runtime_id = Some(id);
         }
 
-        let (worker, handle) = builder
-            .spawn()
-            .await
+        let (worker_handle, worker) = builder
+            .build_worker(runtime)
             .map_err(|e| TelemetryError::Builder(e.to_string()))?;
 
-        Ok(TelemetryClient {
-            handle,
-            metrics: Metrics::new(&worker),
+        Ok((
+            TelemetryClient {
+                metrics: Metrics::new(&worker_handle),
+                worker: worker_handle,
+            },
             worker,
-        })
+        ))
     }
 }
 
@@ -120,7 +124,6 @@ impl TelemetryClientBuilder {
 pub struct TelemetryClient {
     metrics: Metrics,
     worker: TelemetryWorkerHandle,
-    handle: JoinHandle<()>,
 }
 
 /// Telemetry describing the sending of a trace payload
@@ -246,26 +249,18 @@ impl TelemetryClient {
 
     /// Starts the client
     pub async fn start(&self) {
-        if let Err(_e) = self
+        _ = self
             .worker
             .send_msg(TelemetryActions::Lifecycle(LifecycleAction::Start))
-            .await
-        {
-            self.handle.abort();
-        }
+            .await;
     }
 
     /// Shutdowns the telemetry client.
     pub async fn shutdown(self) {
-        if let Err(_e) = self
+        _ = self
             .worker
             .send_msg(TelemetryActions::Lifecycle(LifecycleAction::Stop))
-            .await
-        {
-            self.handle.abort();
-        }
-
-        let _ = self.handle.await;
+            .await;
     }
 }
 
@@ -276,21 +271,23 @@ mod tests {
     use httpmock::MockServer;
     use hyper::{Response, StatusCode};
     use regex::Regex;
+    use tokio::time::sleep;
 
     use super::*;
 
     async fn get_test_client(url: &str) -> TelemetryClient {
-        TelemetryClientBuilder::default()
+        let (client, mut worker) = TelemetryClientBuilder::default()
             .set_service_name("test_service")
             .set_language("test_language")
             .set_language_version("test_language_version")
             .set_tracer_version("test_tracer_version")
             .set_url(url)
             .set_heartbeat(100)
             .set_debug_enabled(true)
-            .build()
-            .await
-            .unwrap()
+            .build(Handle::current())
+            .unwrap();
+        tokio::spawn(async move { worker.run().await });
+        client
     }
 
     #[test]
@@ -320,15 +317,14 @@ mod tests {
     }
 
     #[cfg_attr(miri, ignore)]
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn spawn_test() {
         let client = TelemetryClientBuilder::default()
             .set_service_name("test_service")
             .set_language("test_language")
             .set_language_version("test_language_version")
             .set_tracer_version("test_tracer_version")
-            .build()
-            .await;
+            .build(Handle::current());
 
         assert!(client.is_ok());
     }
@@ -356,6 +352,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -382,6 +381,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -408,6 +410,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -434,6 +439,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -460,6 +468,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -486,6 +497,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -512,6 +526,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -538,6 +555,9 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -675,10 +695,10 @@ mod tests {
             .set_url(&server.url("/"))
             .set_heartbeat(100)
             .set_runtime_id("foo")
-            .build()
-            .await;
+            .build(Handle::current());
 
-        let client = result.unwrap();
+        let (client, mut worker) = result.unwrap();
+        tokio::spawn(async move { worker.run().await });
 
         client.start().await;
         client
@@ -688,6 +708,9 @@ mod tests {
             })
             .unwrap();
         client.shutdown().await;
+        while telemetry_srv.hits_async().await == 0 {
+            sleep(Duration::from_millis(10)).await;
+        }
         // One payload generate-metrics
         telemetry_srv.assert_hits_async(1).await;
     }

data-pipeline/src/trace_exporter/mod.rs

@@ -4,7 +4,7 @@ pub mod agent_response;
 pub mod error;
 use crate::agent_info::{AgentInfoArc, AgentInfoFetcher};
 use crate::stats_exporter::StatsExporter;
-use crate::telemetry::{SendPayloadTelemetry, TelemetryClient, TelemetryClientBuilder};
+use crate::telemetry::{self, SendPayloadTelemetry, TelemetryClient, TelemetryClientBuilder};
 use crate::trace_exporter::error::{RequestError, TraceExporterError};
 use crate::{
     health_metrics, health_metrics::HealthMetric, span_concentrator::SpanConcentrator,
@@ -22,6 +22,7 @@ use ddcommon::header::{
 };
 use ddcommon::tag::Tag;
 use ddcommon::{hyper_migration, parse_uri, tag, Endpoint};
+use ddtelemetry::worker::{self, TelemetryWorker};
 use dogstatsd_client::{new, Client, DogStatsDAction};
 use either::Either;
 use error::BuilderErrorKind;
@@ -68,24 +69,37 @@ pub enum TraceExporterOutputFormat {
 }
 
 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(&self) -> impl std::future::Future<Output = ()> + Send;
+        fn run(&mut self) -> impl std::future::Future<Output = ()> + Send;
+        fn on_pause(&mut self);
     }
 
     impl Worker for StatsExporter {
-        async fn run(&self) {
-            self.run().await
+        async fn run(&mut self) {
+            Self::run(self).await
         }
+        fn on_pause(&mut self) {}
     }
 
     impl Worker for AgentInfoFetcher {
-        async fn run(&self) {
-            self.run().await
+        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();
         }
     }
 
@@ -107,7 +121,7 @@ mod pausable_worker {
         }
 
         pub fn start(&mut self, rt: &Runtime) {
-            if let Self::Paused { worker } = std::mem::replace(self, Self::InvalidState) {
+            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 {
@@ -240,6 +254,7 @@ use pausable_worker::PausableWorker;
 struct TraceExporterWorkers {
     pub info: PausableWorker<AgentInfoFetcher>,
     pub stats: Option<PausableWorker<StatsExporter>>,
+    pub telemetry: Option<PausableWorker<TelemetryWorker>>,
 }
 
 /// The TraceExporter ingest traces from the tracers serialized as messagepack and forward them to
@@ -312,6 +327,12 @@ impl TraceExporter {
         if let Some(stats_worker) = &mut workers.stats {
             stats_worker.start(runtime);
         }
+        if let Some(telemetry_worker) = &mut workers.telemetry {
+            telemetry_worker.start(runtime);
+            if let Some(client) = &self.telemetry {
+                runtime.block_on(client.start());
+            }
+        }
     }
 
     pub fn stop_worker(&self) {
@@ -324,6 +345,9 @@ impl TraceExporter {
                 if let Some(stats_worker) = &mut workers.stats {
                     stats_worker.stop().await
                 };
+                if let Some(telemetry_worker) = &mut workers.telemetry {
+                    telemetry_worker.stop().await
+                };
             });
         }
         // Drop runtime to shutdown all threads
@@ -1137,15 +1161,21 @@ impl TraceExporterBuilder {
                 if let Some(id) = telemetry_config.runtime_id {
                     builder = builder.set_runtime_id(&id);
                 }
-                builder.build().await
+                builder.build(runtime.handle().clone())
             })?)
         } else {
             None
         };
 
-        if let Some(client) = &telemetry {
-            runtime.block_on(client.start());
-        }
+        let (telemetry_client, telemetry_worker) = match telemetry {
+            Some((client, worker)) => {
+                let mut telemetry_worker = PausableWorker::new(worker);
+                telemetry_worker.start(&runtime);
+                runtime.block_on(client.start());
+                (Some(client), Some(telemetry_worker))
+            }
+            None => (None, None),
+        };
 
         Ok(TraceExporter {
             endpoint: Endpoint {
@@ -1177,10 +1207,11 @@ impl TraceExporterBuilder {
             client_side_stats: ArcSwap::new(stats.into()),
             agent_info,
             previous_info_state: ArcSwapOption::new(None),
-            telemetry,
+            telemetry: telemetry_client,
             workers: Arc::new(Mutex::new(TraceExporterWorkers {
                 info: info_fetcher_worker,
                 stats: None,
+                telemetry: telemetry_worker,
             })),
         })
     }