diff --git a/Cargo.lock b/Cargo.lock index 2917cbba83..35ec8ecced 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3091,9 +3091,11 @@ name = "libdd-dogstatsd-client" version = "4.0.0" dependencies = [ "anyhow", + "async-trait", "cadence", "http", "libdd-common", + "libdd-shared-runtime", "serde", "tokio", "tracing", diff --git a/libdd-data-pipeline/Cargo.toml b/libdd-data-pipeline/Cargo.toml index 6425345d22..453d780021 100644 --- a/libdd-data-pipeline/Cargo.toml +++ b/libdd-data-pipeline/Cargo.toml @@ -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", diff --git a/libdd-data-pipeline/src/trace_exporter/builder.rs b/libdd-data-pipeline/src/trace_exporter/builder.rs index 552cdfb1f7..15c4968357 100644 --- a/libdd-data-pipeline/src/trace_exporter/builder.rs +++ b/libdd-data-pipeline/src/trace_exporter/builder.rs @@ -575,9 +575,13 @@ impl TraceExporterBuilder { })?), }; - let dogstatsd = self + let (dogstatsd, dogstatsd_handle) = 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 TraceExporterBuilder { capabilities, #[cfg(not(target_arch = "wasm32"))] workers: TraceExporterWorkers { + dogstatsd: dogstatsd_handle, info_fetcher: info_fetcher_handle, #[cfg(feature = "telemetry")] telemetry: telemetry_handle, diff --git a/libdd-data-pipeline/src/trace_exporter/mod.rs b/libdd-data-pipeline/src/trace_exporter/mod.rs index 1722c2dbe1..06434e9331 100644 --- a/libdd-data-pipeline/src/trace_exporter/mod.rs +++ b/libdd-data-pipeline/src/trace_exporter/mod.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, /// `None` when no background `/info` fetcher is started (agentless trace /// export mode, log-export mode). info_fetcher: Option, @@ -345,6 +346,10 @@ impl< join_set.spawn(async move { handle.stop().await }); } + if let Some(dogstatsd) = self.workers.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 }); } diff --git a/libdd-dogstatsd-client/Cargo.toml b/libdd-dogstatsd-client/Cargo.toml index 4182c59f64..4d4c0f530c 100644 --- a/libdd-dogstatsd-client/Cargo.toml +++ b/libdd-dogstatsd-client/Cargo.toml @@ -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", default-features = false, 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} diff --git a/libdd-dogstatsd-client/src/client/mod.rs b/libdd-dogstatsd-client/src/client/mod.rs index b06cd3641e..2a8fd590dd 100644 --- a/libdd-dogstatsd-client/src/client/mod.rs +++ b/libdd-dogstatsd-client/src/client/mod.rs @@ -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) { @@ -239,4 +270,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)); + } } diff --git a/libdd-dogstatsd-client/src/client/shared_runtime_sink.rs b/libdd-dogstatsd-client/src/client/shared_runtime_sink.rs new file mode 100644 index 0000000000..f8033a1df5 --- /dev/null +++ b/libdd-dogstatsd-client/src/client/shared_runtime_sink.rs @@ -0,0 +1,149 @@ +// 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, + sink: Arc, +} + +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 { + 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, + sink: Arc, + pending: Option, +} + +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; + } + } + } + + /// 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 = 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)) +}