feat(dogstatsd): add shared_runtime buffered sink

Author: VianneyRuhlmann

Cargo.lock

@@ -41,7 +41,7 @@ libdd-trace-stats = { version = "6.0.0", path = "../libdd-trace-stats", default-
 libdd-trace-utils = { version = "9.0.0", path = "../libdd-trace-utils", default-features = false }
 libdd-trace-obfuscation = { version = "5.0.0", path = "../libdd-trace-obfuscation", default-features = false, optional = true }
 libdd-ddsketch = { version = "1.1.0", path = "../libdd-ddsketch" }
-libdd-dogstatsd-client = { version = "4.0.0", path = "../libdd-dogstatsd-client", default-features = false }
+libdd-dogstatsd-client = { version = "4.0.0", path = "../libdd-dogstatsd-client", default-features = false, features = ["shared-runtime"] }
 libdd-tinybytes = { version = "1.1.1", path = "../libdd-tinybytes", features = [
     "bytes_string",
     "serialization",

libdd-data-pipeline/Cargo.toml

@@ -575,9 +575,13 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
             })?),
         };
 
-        let dogstatsd = self
+        let (dogstatsd, dogstatsd_worker) = self
             .dogstatsd_url
-            .and_then(|u| DogStatsDClient::new(Endpoint::from_slice(&u)).ok());
+            .and_then(|u| {
+                DogStatsDClient::new_with_shared_runtime(Endpoint::from_slice(&u), &*shared_runtime)
+                    .ok()
+            })
+            .unzip();
 
         let base_url = self.url.as_deref().unwrap_or(DEFAULT_AGENT_URL);
 
@@ -836,6 +840,7 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
             capabilities,
             #[cfg(not(target_arch = "wasm32"))]
             workers: TraceExporterWorkers {
+                dogstatsd: dogstatsd_handle,
                 info_fetcher: info_fetcher_handle,
                 #[cfg(feature = "telemetry")]
                 telemetry: telemetry_handle,

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

@@ -189,6 +189,7 @@ pub use libdd_trace_utils::tracer_metadata::TracerMetadata;
 #[cfg(not(target_arch = "wasm32"))]
 #[derive(Debug)]
 pub(crate) struct TraceExporterWorkers {
+    dogstatsd: Option<WorkerHandle>,
     /// `None` when no background `/info` fetcher is started (agentless trace
     /// export mode, log-export mode).
     info_fetcher: Option<WorkerHandle>,
@@ -345,6 +346,10 @@ impl<
                 join_set.spawn(async move { handle.stop().await });
             }
 
+            if let Some(dogstatsd) = self.dogstatsd {
+                join_set.spawn(async move { dogstatsd.stop().await });
+            }
+
             if let Some(info_fetcher) = self.workers.info_fetcher {
                 join_set.spawn(async move { info_fetcher.stop().await });
             }

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

@@ -18,12 +18,15 @@ serde = { version = "1.0", features = ["derive", "rc"] }
 tracing = { version = "0.1", default-features = false }
 anyhow = { version = "1.0" }
 http = "1.1"
+libdd-shared-runtime = { version = "2.0.0", path = "../libdd-shared-runtime", optional = true }
+tokio = { version = "1.23", features = ["sync"], optional = true }
+async-trait = { version = "0.1", optional = true }
 
 [features]
 default = ["https"]
 https = ["libdd-common/https"]
 fips = ["libdd-common/fips"]
-
+shared-runtime = ["dep:libdd-shared-runtime", "dep:tokio", "dep:async-trait"]
 
 [dev-dependencies]
 tokio = {version = "1.23", features = ["rt", "time", "test-util", "rt-multi-thread"], default-features = false}

libdd-dogstatsd-client/Cargo.toml

@@ -6,6 +6,8 @@ use cadence::prelude::*;
 use cadence::{Metric, MetricBuilder, QueuingMetricSink, StatsdClient};
 use libdd_common::tag::Tag;
 use libdd_common::Endpoint;
+#[cfg(feature = "shared-runtime")]
+use libdd_shared_runtime::{SharedRuntime, WorkerHandle};
 use sink::create_udp_sink;
 #[cfg(unix)]
 use sink::create_unix_sink;
@@ -15,6 +17,10 @@ use tracing::error;
 /// Provides transport sink which can be wrapped in buffered sink
 mod sink;
 
+/// Provides a buffered sink running on a [libdd_shared_runtime::SharedRuntime]
+#[cfg(feature = "shared-runtime")]
+mod shared_runtime_sink;
+
 // Queue with a maximum capacity of 32K elements
 const QUEUE_SIZE: usize = 32 * 1024;
 
@@ -47,6 +53,31 @@ impl DogStatsDClient {
         })
     }
 
+    /// Create a [`Client`] backed by a [`MetricSinkWorker`] running on the
+    /// provided [`SharedRuntime`].
+    ///
+    /// Returns the client and a [`WorkerHandle`] that can be used to stop the
+    /// worker independently of the runtime.
+    ///
+    /// # Errors
+    /// Returns an error if the endpoint is invalid or the worker cannot be spawned.
+    #[cfg(feature = "shared-runtime")]
+    pub fn new_with_shared_runtime(
+        endpoint: Endpoint,
+        runtime: &impl SharedRuntime,
+    ) -> anyhow::Result<(Self, WorkerHandle)> {
+        let (sink, handle) = shared_runtime_sink::create_shared_runtime_sink(&endpoint, runtime)?;
+
+        let client = Self {
+            inner: Arc::new(InnerClient {
+                client: OnceLock::from(StatsdClient::from_sink("", sink)),
+                endpoint,
+            }),
+        };
+
+        Ok((client, handle))
+    }
+
     /// Send a vector of DogStatsDActionOwned, this is the same as `send` except it uses the
     /// "owned" version of DogStatsDAction. See the docs for DogStatsDActionOwned for details.
     pub fn send_owned(&self, actions: Vec<DogStatsDActionOwned>) {
@@ -243,4 +274,51 @@ mod test {
             task.await.unwrap();
         }
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    #[cfg(feature = "shared-runtime")]
+    fn test_shared_runtime_flusher() {
+        use libdd_shared_runtime::{BasicRuntime, BlockingRuntime, SharedRuntime};
+
+        let socket = net::UdpSocket::bind("127.0.0.1:0").expect("failed to bind host socket");
+        let _ = socket.set_read_timeout(Some(Duration::from_millis(500)));
+
+        let runtime = BasicRuntime::new().unwrap();
+
+        let (flusher, _handle) = DogStatsDClient::new_with_shared_runtime(
+            Endpoint::from_slice(socket.local_addr().unwrap().to_string().as_str()),
+            &runtime,
+        )
+        .unwrap();
+
+        flusher.send(vec![
+            Count("test_count", 3, &vec![tag!("foo", "bar")]),
+            Count("test_neg_count", -2, &vec![]),
+            Distribution("test_distribution", 4.2, &vec![]),
+            Gauge("test_gauge", 7.6, &vec![]),
+            Histogram("test_histogram", 8.0, &vec![]),
+            Set("test_set", 9, &vec![tag!("the", "end")]),
+            Set("test_neg_set", -1, &vec![]),
+        ]);
+
+        runtime
+            .block_on(runtime.shutdown_async())
+            .expect("Failed to shutdown runtime");
+
+        fn read(socket: &net::UdpSocket) -> String {
+            let mut buf = [0; 100];
+            socket.recv(&mut buf).expect("No data");
+            let datagram = String::from_utf8_lossy(buf.strip_suffix(&[0]).unwrap());
+            datagram.trim_matches(char::from(0)).to_string()
+        }
+
+        assert_eq!("test_count:3|c|#foo:bar", read(&socket));
+        assert_eq!("test_neg_count:-2|c", read(&socket));
+        assert_eq!("test_distribution:4.2|d", read(&socket));
+        assert_eq!("test_gauge:7.6|g", read(&socket));
+        assert_eq!("test_histogram:8|h", read(&socket));
+        assert_eq!("test_set:9|s|#the:end", read(&socket));
+        assert_eq!("test_neg_set:-1|s", read(&socket));
+    }
 }

libdd-dogstatsd-client/src/client/mod.rs

@@ -0,0 +1,150 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::anyhow;
+use async_trait::async_trait;
+use cadence::{MetricSink, SinkStats};
+use libdd_common::Endpoint;
+use libdd_shared_runtime::{worker::Worker, SharedRuntime, WorkerHandle};
+use std::fmt;
+use std::io;
+use std::panic::RefUnwindSafe;
+use std::sync::Arc;
+use tokio::sync::mpsc;
+use tracing::error;
+
+use super::{sink, QUEUE_SIZE};
+
+/// A [`MetricSink`] that offloads sent metrics to a [`SharedRuntime`].
+#[derive(Clone)]
+pub struct SharedRuntimeMetricSink {
+    sender: mpsc::Sender<String>,
+    sink: Arc<dyn MetricSink + Send + Sync + RefUnwindSafe>,
+}
+
+impl fmt::Debug for SharedRuntimeMetricSink {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "SharedRuntimeMetricSink {{ {:?} }}", self.sender)
+    }
+}
+
+impl MetricSink for SharedRuntimeMetricSink {
+    fn emit(&self, metric: &str) -> io::Result<usize> {
+        let len = metric.len();
+        self.sender
+            .try_send(metric.to_owned())
+            .map(|_| len)
+            .map_err(|e| match e {
+                mpsc::error::TrySendError::Full(_) => {
+                    io::Error::new(io::ErrorKind::WouldBlock, "dogstatsd channel full")
+                }
+                mpsc::error::TrySendError::Closed(_) => {
+                    io::Error::new(io::ErrorKind::BrokenPipe, "dogstatsd channel closed")
+                }
+            })
+    }
+
+    fn flush(&self) -> Result<(), std::io::Error> {
+        self.sink.flush()
+    }
+
+    fn stats(&self) -> SinkStats {
+        self.sink.stats()
+    }
+}
+
+/// A [`Worker`] that drains metrics from a channel and forwards them to a
+/// wrapped [`MetricSink`] (e.g. `UdpMetricSink`).
+pub struct MetricSinkWorker {
+    receiver: mpsc::Receiver<String>,
+    sink: Arc<dyn MetricSink + Send + Sync + RefUnwindSafe>,
+    pending: Option<String>,
+}
+
+impl std::fmt::Debug for MetricSinkWorker {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MetricSinkWorker")
+            .field("pending", &self.pending)
+            .finish_non_exhaustive()
+    }
+}
+
+#[async_trait]
+impl Worker for MetricSinkWorker {
+    /// Awaits the next metric from the channel, storing it in `pending`.
+    ///
+    /// If the channel is closed this future never resolves, allowing the
+    /// SharedRuntime to cancel it cleanly at the next yield point.
+    async fn trigger(&mut self) {
+        match self.receiver.recv().await {
+            Some(metric) => self.pending = Some(metric),
+            None => {
+                // Channel closed — park until cancelled by the runtime.
+                std::future::pending::<()>().await;
+            }
+        }
+        println!("Received metric: {:#?}", self.pending);
+    }
+
+    /// Forwards the single metric stored by `trigger` to the wrapped sink.
+    async fn run(&mut self) {
+        if let Some(metric) = self.pending.take() {
+            if let Err(e) = self.sink.emit(&metric) {
+                error!(?e, "MetricSinkWorker: failed to emit metric");
+            }
+        }
+    }
+
+    /// Drains any remaining queued metrics and flushes the wrapped sink.
+    async fn shutdown(&mut self) {
+        // Forward the metric that was sitting in `pending`, if any.
+        if let Some(metric) = self.pending.take() {
+            if let Err(e) = self.sink.emit(&metric) {
+                error!(?e, "MetricSinkWorker: failed to emit metric on shutdown");
+            }
+        }
+        // Drain the channel.
+        while let Ok(metric) = self.receiver.try_recv() {
+            if let Err(e) = self.sink.emit(&metric) {
+                error!(?e, "MetricSinkWorker: failed to emit metric on shutdown");
+            }
+        }
+        if let Err(e) = self.sink.flush() {
+            error!(?e, "MetricSinkWorker: failed to flush sink on shutdown");
+        }
+    }
+
+    /// Reset the worker in the child process after a fork.
+    fn reset(&mut self) {
+        self.pending = None;
+        // Drain the channel
+        while self.receiver.try_recv().is_ok() {}
+    }
+}
+
+pub fn create_shared_runtime_sink(
+    endpoint: &Endpoint,
+    runtime: &impl SharedRuntime,
+) -> anyhow::Result<(SharedRuntimeMetricSink, WorkerHandle)> {
+    let (tx, rx) = mpsc::channel(QUEUE_SIZE);
+
+    let sink: Arc<dyn MetricSink + Send + Sync + RefUnwindSafe> = match endpoint.url.scheme_str() {
+        #[cfg(unix)]
+        Some("unix") => Arc::new(sink::create_unix_sink(endpoint)?),
+        _ => Arc::new(sink::create_udp_sink(endpoint)?),
+    };
+
+    let sink_worker = MetricSinkWorker {
+        receiver: rx,
+        sink: sink.clone(),
+        pending: None,
+    };
+
+    let handle = runtime
+        .spawn_worker(sink_worker, true)
+        .map_err(|e| anyhow!("failed to spawn MetricSinkWorker: {e}"))?;
+
+    let shared_sink = SharedRuntimeMetricSink { sender: tx, sink };
+
+    Ok((shared_sink, handle))
+}