Skip to main content

fast_telemetry_export/
dogstatsd.rs

1//! DogStatsD metrics exporter.
2//!
3//! Exports metrics to a DogStatsD-compatible agent (Datadog, StatsD) over UDP.
4//! The exporter handles UDP socket management, packet batching (respecting MTU),
5//! and newline-delimited packet splitting via `memchr`.
6//!
7//! The actual metric serialization is provided by the caller via a closure,
8//! making this exporter work with any metrics struct.
9
10use std::time::Duration;
11
12/// Configuration for the DogStatsD exporter.
13#[derive(Clone)]
14pub struct DogStatsDConfig {
15    /// DogStatsD endpoint (host:port), e.g. "127.0.0.1:8125"
16    pub endpoint: String,
17    /// Export interval
18    pub interval: Duration,
19    /// Maximum UDP packet size (default: 8000 bytes)
20    pub max_packet_size: usize,
21}
22
23impl Default for DogStatsDConfig {
24    fn default() -> Self {
25        Self {
26            endpoint: "127.0.0.1:8125".to_string(),
27            interval: Duration::from_secs(10),
28            max_packet_size: 8000,
29        }
30    }
31}
32
33impl DogStatsDConfig {
34    pub fn new(endpoint: impl Into<String>) -> Self {
35        Self {
36            endpoint: endpoint.into(),
37            ..Default::default()
38        }
39    }
40
41    pub fn with_interval(mut self, interval: Duration) -> Self {
42        self.interval = interval;
43        self
44    }
45
46    pub fn with_max_packet_size(mut self, size: usize) -> Self {
47        self.max_packet_size = size;
48        self
49    }
50}
51
52/// Run the DogStatsD export loop.
53///
54/// `export_fn` is called each cycle with a `&mut String` buffer. The closure
55/// should append newline-delimited DogStatsD metric lines (typically via
56/// `ExportMetrics::export_dogstatsd_delta`). The exporter handles UDP batching
57/// and sending.
58///
59/// Runs until `cancel` is triggered. On cancellation, returns immediately
60/// without a final export (DogStatsD is fire-and-forget).
61///
62/// # Example
63///
64/// ```ignore
65/// use std::sync::Arc;
66///
67/// use fast_telemetry_export::dogstatsd::{DogStatsDConfig, run};
68/// use tokio_util::sync::CancellationToken;
69///
70/// let metrics = Arc::new(MyMetrics::new());
71/// let mut state = MyMetricsExportState::new();
72/// let tags: Vec<(&str, &str)> = vec![("service", "myapp")];
73/// let cancel = CancellationToken::new();
74/// let config = DogStatsDConfig::new("127.0.0.1:8125");
75///
76/// let m = metrics.clone();
77/// tokio::spawn(run(config, cancel, move |output| {
78///     m.export_dogstatsd_delta(output, &tags, &mut state);
79/// }));
80/// ```
81///
82/// `MyMetricsExportState` is the derive-generated state type for delta
83/// DogStatsD export. Keep one state instance per export sink.
84#[cfg(feature = "tokio-runtime")]
85pub async fn run<F>(
86    config: DogStatsDConfig,
87    cancel: tokio_util::sync::CancellationToken,
88    mut export_fn: F,
89) where
90    F: FnMut(&mut String),
91{
92    use tokio::net::UdpSocket;
93    use tokio::time::MissedTickBehavior;
94
95    log::info!("Starting DogStatsD exporter, endpoint={}", config.endpoint);
96
97    let socket = match UdpSocket::bind("0.0.0.0:0").await {
98        Ok(s) => s,
99        Err(e) => {
100            log::error!("Failed to bind UDP socket for DogStatsD export: {e}");
101            return;
102        }
103    };
104
105    if let Err(e) = socket.connect(&config.endpoint).await {
106        log::error!("Failed to connect UDP socket to {}: {e}", config.endpoint);
107        return;
108    }
109
110    let max_packet_size = config.max_packet_size;
111    let mut output = String::with_capacity(16384);
112    let mut batch = Vec::<u8>::with_capacity(max_packet_size);
113
114    let mut interval = tokio::time::interval(config.interval);
115    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
116    interval.tick().await;
117
118    loop {
119        tokio::select! {
120            _ = interval.tick() => {}
121            _ = cancel.cancelled() => {
122                log::info!("DogStatsD exporter shutting down");
123                return;
124            }
125        }
126
127        output.clear();
128        export_fn(&mut output);
129
130        if output.is_empty() {
131            continue;
132        }
133
134        let output_bytes = output.as_bytes();
135        batch.clear();
136
137        let mut total_sent = 0usize;
138        let mut batch_count = 0usize;
139        let mut metric_count = 0usize;
140        let mut start = 0usize;
141
142        for nl in memchr::memchr_iter(b'\n', output_bytes) {
143            let end = nl + 1;
144            let line = &output_bytes[start..end];
145            let line_len = line.len();
146            metric_count += 1;
147
148            if line_len > max_packet_size {
149                log::warn!(
150                    "Dropping oversized metric line ({line_len} bytes, max {max_packet_size})"
151                );
152                start = end;
153                continue;
154            }
155
156            if !batch.is_empty() && batch.len() + line_len > max_packet_size {
157                match socket.send(&batch).await {
158                    Ok(n) => {
159                        total_sent += n;
160                        batch_count += 1;
161                    }
162                    Err(e) => log::warn!("Failed to send DogStatsD batch: {e}"),
163                }
164                batch.clear();
165            }
166
167            batch.extend_from_slice(line);
168            start = end;
169        }
170
171        // Handle trailing bytes if output didn't end with '\n'
172        if start < output_bytes.len() {
173            let line = &output_bytes[start..];
174            let line_len = line.len();
175            metric_count += 1;
176
177            if line_len <= max_packet_size {
178                if !batch.is_empty() && batch.len() + line_len > max_packet_size {
179                    match socket.send(&batch).await {
180                        Ok(n) => {
181                            total_sent += n;
182                            batch_count += 1;
183                        }
184                        Err(e) => log::warn!("Failed to send DogStatsD batch: {e}"),
185                    }
186                    batch.clear();
187                }
188                batch.extend_from_slice(line);
189            } else {
190                log::warn!("Dropping oversized trailing metric ({line_len} bytes)");
191            }
192        }
193
194        if !batch.is_empty() {
195            match socket.send(&batch).await {
196                Ok(n) => {
197                    total_sent += n;
198                    batch_count += 1;
199                }
200                Err(e) => log::warn!("Failed to send final DogStatsD batch: {e}"),
201            }
202        }
203
204        log::debug!(
205            "DogStatsD export: {metric_count} metrics, {batch_count} batches, {total_sent} bytes"
206        );
207    }
208}
209
210/// Run the DogStatsD export loop on a monoio runtime.
211///
212/// This is the monoio-native counterpart to [`run`]. It uses
213/// [`monoio::net::udp::UdpSocket`] and [`monoio::time::interval`], so the
214/// caller must run it inside a monoio runtime with timers enabled.
215#[cfg(feature = "monoio")]
216pub async fn run_monoio<F>(
217    config: DogStatsDConfig,
218    cancel: tokio_util::sync::CancellationToken,
219    mut export_fn: F,
220) where
221    F: FnMut(&mut String),
222{
223    use std::net::{SocketAddr, ToSocketAddrs};
224
225    use monoio::net::udp::UdpSocket;
226    use monoio::time::MissedTickBehavior;
227
228    log::info!(
229        "Starting monoio DogStatsD exporter, endpoint={}",
230        config.endpoint
231    );
232
233    let endpoint = match config.endpoint.to_socket_addrs() {
234        Ok(mut addrs) => match addrs.next() {
235            Some(addr) => addr,
236            None => {
237                log::error!("DogStatsD endpoint resolved to no addresses");
238                return;
239            }
240        },
241        Err(e) => {
242            log::error!(
243                "Failed to resolve DogStatsD endpoint {}: {e}",
244                config.endpoint
245            );
246            return;
247        }
248    };
249
250    let bind_addr: SocketAddr = if endpoint.is_ipv4() {
251        "0.0.0.0:0"
252    } else {
253        "[::]:0"
254    }
255    .parse()
256    .expect("valid UDP bind address");
257
258    let socket = match UdpSocket::bind(bind_addr) {
259        Ok(s) => s,
260        Err(e) => {
261            log::error!("Failed to bind monoio UDP socket for DogStatsD export: {e}");
262            return;
263        }
264    };
265
266    if let Err(e) = socket.connect(endpoint).await {
267        log::error!("Failed to connect monoio UDP socket to {endpoint}: {e}");
268        return;
269    }
270
271    let max_packet_size = config.max_packet_size;
272    let mut output = String::with_capacity(16384);
273    let mut batch = Vec::<u8>::with_capacity(max_packet_size);
274
275    let mut interval = monoio::time::interval(config.interval);
276    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
277    interval.tick().await;
278
279    loop {
280        monoio::select! {
281            _ = interval.tick() => {}
282            _ = cancel.cancelled() => {
283                log::info!("monoio DogStatsD exporter shutting down");
284                return;
285            }
286        }
287
288        output.clear();
289        export_fn(&mut output);
290
291        if output.is_empty() {
292            continue;
293        }
294
295        let output_bytes = output.as_bytes();
296        batch.clear();
297
298        let mut total_sent = 0usize;
299        let mut batch_count = 0usize;
300        let mut metric_count = 0usize;
301        let mut start = 0usize;
302
303        for nl in memchr::memchr_iter(b'\n', output_bytes) {
304            let end = nl + 1;
305            let line = &output_bytes[start..end];
306            let line_len = line.len();
307            metric_count += 1;
308
309            if line_len > max_packet_size {
310                log::warn!(
311                    "Dropping oversized metric line ({line_len} bytes, max {max_packet_size})"
312                );
313                start = end;
314                continue;
315            }
316
317            if !batch.is_empty()
318                && batch.len() + line_len > max_packet_size
319                && let Some(n) = send_monoio_batch(&socket, &mut batch, "DogStatsD batch").await
320            {
321                total_sent += n;
322                batch_count += 1;
323            }
324
325            batch.extend_from_slice(line);
326            start = end;
327        }
328
329        if start < output_bytes.len() {
330            let line = &output_bytes[start..];
331            let line_len = line.len();
332            metric_count += 1;
333
334            if line_len <= max_packet_size {
335                if !batch.is_empty()
336                    && batch.len() + line_len > max_packet_size
337                    && let Some(n) = send_monoio_batch(&socket, &mut batch, "DogStatsD batch").await
338                {
339                    total_sent += n;
340                    batch_count += 1;
341                }
342                batch.extend_from_slice(line);
343            } else {
344                log::warn!("Dropping oversized trailing metric ({line_len} bytes)");
345            }
346        }
347
348        if !batch.is_empty()
349            && let Some(n) = send_monoio_batch(&socket, &mut batch, "final DogStatsD batch").await
350        {
351            total_sent += n;
352            batch_count += 1;
353        }
354
355        log::debug!(
356            "monoio DogStatsD export: {metric_count} metrics, {batch_count} batches, {total_sent} bytes"
357        );
358    }
359}
360
361/// Run the DogStatsD export loop on a compio runtime.
362///
363/// This is the compio-native counterpart to `run`. It uses
364/// [`compio::net::UdpSocket`] and [`compio::time`], so the caller must run it
365/// inside a compio runtime. `cancel` may be any future that completes when the
366/// exporter should shut down.
367#[cfg(feature = "compio")]
368pub async fn run_compio<F>(
369    config: DogStatsDConfig,
370    cancel: impl std::future::Future<Output = ()>,
371    mut export_fn: F,
372) where
373    F: FnMut(&mut String),
374{
375    use std::net::{SocketAddr, ToSocketAddrs};
376
377    use compio::net::UdpSocket;
378    use futures_util::{FutureExt as _, select};
379
380    log::info!(
381        "Starting compio DogStatsD exporter, endpoint={}",
382        config.endpoint
383    );
384
385    let endpoint = match config.endpoint.to_socket_addrs() {
386        Ok(mut addrs) => match addrs.next() {
387            Some(addr) => addr,
388            None => {
389                log::error!("DogStatsD endpoint resolved to no addresses");
390                return;
391            }
392        },
393        Err(e) => {
394            log::error!(
395                "Failed to resolve DogStatsD endpoint {}: {e}",
396                config.endpoint
397            );
398            return;
399        }
400    };
401
402    let bind_addr: SocketAddr = if endpoint.is_ipv4() {
403        "0.0.0.0:0"
404    } else {
405        "[::]:0"
406    }
407    .parse()
408    .expect("valid UDP bind address");
409
410    let socket = match UdpSocket::bind(bind_addr).await {
411        Ok(s) => s,
412        Err(e) => {
413            log::error!("Failed to bind compio UDP socket for DogStatsD export: {e}");
414            return;
415        }
416    };
417
418    if let Err(e) = socket.connect(endpoint).await {
419        log::error!("Failed to connect compio UDP socket to {endpoint}: {e}");
420        return;
421    }
422
423    let max_packet_size = config.max_packet_size;
424    let mut output = String::with_capacity(16384);
425    let mut batch = Vec::<u8>::with_capacity(max_packet_size);
426    let mut interval = compio::time::interval(config.interval);
427    interval.tick().await;
428    let cancel = cancel.fuse();
429    let mut cancel = std::pin::pin!(cancel);
430
431    loop {
432        let tick = interval.tick();
433        let tick = std::pin::pin!(tick);
434
435        select! {
436            _ = tick.fuse() => {},
437            _ = cancel.as_mut() => {
438                log::info!("compio DogStatsD exporter shutting down");
439                return;
440            }
441        }
442
443        output.clear();
444        export_fn(&mut output);
445
446        if output.is_empty() {
447            continue;
448        }
449
450        let output_bytes = output.as_bytes();
451        batch.clear();
452
453        let mut total_sent = 0usize;
454        let mut batch_count = 0usize;
455        let mut metric_count = 0usize;
456        let mut start = 0usize;
457
458        for nl in memchr::memchr_iter(b'\n', output_bytes) {
459            let end = nl + 1;
460            let line = &output_bytes[start..end];
461            let line_len = line.len();
462            metric_count += 1;
463
464            if line_len > max_packet_size {
465                log::warn!(
466                    "Dropping oversized metric line ({line_len} bytes, max {max_packet_size})"
467                );
468                start = end;
469                continue;
470            }
471
472            if !batch.is_empty()
473                && batch.len() + line_len > max_packet_size
474                && let Some(n) = send_compio_batch(&socket, &mut batch, "DogStatsD batch").await
475            {
476                total_sent += n;
477                batch_count += 1;
478            }
479
480            batch.extend_from_slice(line);
481            start = end;
482        }
483
484        if start < output_bytes.len() {
485            let line = &output_bytes[start..];
486            let line_len = line.len();
487            metric_count += 1;
488
489            if line_len <= max_packet_size {
490                if !batch.is_empty()
491                    && batch.len() + line_len > max_packet_size
492                    && let Some(n) = send_compio_batch(&socket, &mut batch, "DogStatsD batch").await
493                {
494                    total_sent += n;
495                    batch_count += 1;
496                }
497                batch.extend_from_slice(line);
498            } else {
499                log::warn!("Dropping oversized trailing metric ({line_len} bytes)");
500            }
501        }
502
503        if !batch.is_empty()
504            && let Some(n) = send_compio_batch(&socket, &mut batch, "final DogStatsD batch").await
505        {
506            total_sent += n;
507            batch_count += 1;
508        }
509
510        log::debug!(
511            "compio DogStatsD export: {metric_count} metrics, {batch_count} batches, {total_sent} bytes"
512        );
513    }
514}
515
516#[cfg(feature = "monoio")]
517async fn send_monoio_batch(
518    socket: &monoio::net::udp::UdpSocket,
519    batch: &mut Vec<u8>,
520    context: &str,
521) -> Option<usize> {
522    let send_buf = std::mem::take(batch);
523    let (result, mut send_buf) = socket.send(send_buf).await;
524    send_buf.clear();
525    *batch = send_buf;
526
527    match result {
528        Ok(n) => Some(n),
529        Err(e) => {
530            log::warn!("Failed to send monoio {context}: {e}");
531            None
532        }
533    }
534}
535
536#[cfg(feature = "compio")]
537async fn send_compio_batch(
538    socket: &compio::net::UdpSocket,
539    batch: &mut Vec<u8>,
540    context: &str,
541) -> Option<usize> {
542    let send_buf = std::mem::take(batch);
543    let compio::BufResult(result, mut send_buf) = socket.send(send_buf).await;
544    send_buf.clear();
545    *batch = send_buf;
546
547    match result {
548        Ok(n) => Some(n),
549        Err(e) => {
550            log::warn!("Failed to send compio {context}: {e}");
551            None
552        }
553    }
554}