saluki_components/destinations/dsd_stats/
mod.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use saluki_api::{
6    extract::{Query, State},
7    routing::{get, Router},
8    APIHandler, StatusCode,
9};
10use saluki_common::time::get_coarse_unix_timestamp;
11use saluki_context::tags::TagSet;
12use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
13use saluki_core::{
14    components::{
15        destinations::{Destination, DestinationBuilder, DestinationContext},
16        ComponentContext,
17    },
18    data_model::event::{Event, EventType},
19};
20use saluki_error::GenericError;
21use serde::{Deserialize, Serialize, Serializer};
22use serde_json;
23use stringtheory::MetaString;
24use tokio::time::{sleep, Duration, Instant};
25use tokio::{
26    pin,
27    sync::{Mutex, OwnedMutexGuard},
28};
29use tokio::{select, sync::mpsc, sync::oneshot};
30
31type StatsRequestReceiver = mpsc::Receiver<(oneshot::Sender<StatsResponse>, u64)>;
32
33#[derive(Debug, Default, Clone, Serialize)]
34pub struct MetricSample {
35    count: u64,
36    last_seen: u64,
37}
38#[derive(Serialize)]
39enum StatsResponse {
40    /// An existing statistics collection request is running.
41    AlreadyRunning {
42        /// Number of seconds to wait before trying again.
43        try_after: u64,
44    },
45
46    Statistics(CollectedStatistics),
47}
48
49#[derive(Serialize)]
50struct CollectedStatistics {
51    /// Start time of the collected metrics, as a Unix timestamp.
52    start_time_unix: u64,
53
54    /// End time of the collected metrics, as a Unix timestamp.
55    end_time_unix: u64,
56
57    /// Collected statistics.
58    stats: FlattenedStats,
59}
60
61#[derive(Serialize)]
62struct FlattenedMetricStat<'a> {
63    #[serde(flatten)]
64    context: &'a ContextNoOrigin,
65
66    #[serde(flatten)]
67    stats: &'a MetricSample,
68}
69
70struct FlattenedStats(HashMap<ContextNoOrigin, MetricSample>);
71
72impl Serialize for FlattenedStats {
73    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
74    where
75        S: Serializer,
76    {
77        serializer.collect_seq(
78            self.0
79                .iter()
80                .map(|(context, stats)| FlattenedMetricStat { context, stats }),
81        )
82    }
83}
84
85/// Configuration for DogStatsD statistics destination and API handler.
86#[derive(Clone)]
87pub struct DogStatsDStatisticsConfiguration {
88    api_handler: DogStatsDStatsAPIHandler,
89    rx: Arc<Mutex<StatsRequestReceiver>>,
90}
91/// State for the DogStatsD API handler.
92#[derive(Clone)]
93pub struct DogStatsDStatsAPIHandlerState {
94    tx: Arc<mpsc::Sender<(oneshot::Sender<StatsResponse>, u64)>>,
95}
96
97/// API handler for DogStatsD statistics endpoint.
98#[derive(Clone)]
99pub struct DogStatsDStatsAPIHandler {
100    state: DogStatsDStatsAPIHandlerState,
101}
102
103/// DogStatsD destination that collects metrics and processes statistics.
104pub struct DogStatsDStats {
105    rx: OwnedMutexGuard<StatsRequestReceiver>,
106}
107
108#[async_trait::async_trait]
109impl Destination for DogStatsDStats {
110    async fn run(mut self: Box<Self>, mut context: DestinationContext) -> Result<(), GenericError> {
111        let mut health = context.take_health_handle();
112        let mut collection_active = false;
113        let mut stats_response_tx: Option<tokio::sync::oneshot::Sender<StatsResponse>> = None;
114        let mut current_stats: Option<HashMap<ContextNoOrigin, MetricSample>> = None;
115        let mut stats_collection_start_time = 0;
116        let mut stats_collection_end_time: u64 = 0;
117        let collection_done = sleep(std::time::Duration::ZERO);
118        pin!(collection_done);
119
120        health.mark_ready();
121
122        loop {
123            select! {
124                _ = health.live() => {
125                    continue
126                },
127                Some((response_tx, collection_period_secs)) = self.rx.recv() => {
128                    if collection_active {
129                        // We're already collecting statistics for another stats request
130                        // so inform the caller they need to try again later.
131                        let now = get_coarse_unix_timestamp();
132                        saluki_antithesis::always_or_unreachable!(
133                            now >= stats_collection_start_time,
134                            "dsd_stats collection clock did not move backward",
135                            { "now": now, "start_time": stats_collection_start_time }
136                        );
137                        let try_after = stats_collection_end_time.saturating_sub(now);
138
139                        // We don't care if we can successfully send back a response or not.
140                        let _ = response_tx.send(StatsResponse::AlreadyRunning { try_after });
141                    } else {
142                        // Start collection.
143                        collection_active = true;
144                        stats_collection_start_time = get_coarse_unix_timestamp();
145                        stats_collection_end_time = stats_collection_start_time + collection_period_secs;
146                        stats_response_tx = Some(response_tx);
147                        current_stats = Some(HashMap::new());
148                        collection_done.as_mut().reset(Instant::now() + Duration::from_secs(collection_period_secs));
149                    }
150                },
151                maybe_events = context.events().next() => match maybe_events {
152                    Some(events) => {
153                        if let Some(stats) = current_stats.as_mut() {
154                            // We're actively collecting, so process the metrics.
155                            for event in events {
156                                if let Event::Metric(metric) = event {
157
158                                    let context = metric.context();
159                                    let new_context = ContextNoOrigin {
160                                        name: context.name().clone(),
161                                        tags: context.tags().clone(),
162                                    };
163
164                                    let timestamp = get_coarse_unix_timestamp();
165                                    let sample = stats.entry(new_context).or_default();
166                                    sample.count += 1;
167                                    sample.last_seen = timestamp;
168
169                            }
170                        }
171                     }},
172                     None => break,
173                },
174                _ = &mut collection_done, if collection_active => {
175                    collection_active = false;
176
177                    // Build the response.
178                    let stats = match current_stats.take() {
179                        Some(stats) => stats,
180                        None => continue,
181                    };
182
183                    let response = StatsResponse::Statistics(CollectedStatistics {
184                        start_time_unix: stats_collection_start_time,
185                        end_time_unix: stats_collection_end_time,
186                        stats: FlattenedStats(stats),
187                    });
188
189                    let response_tx = match stats_response_tx.take() {
190                        Some(tx) => tx,
191                        None => continue,
192                    };
193
194                    // We don't care if we can successfully send back a response or not.
195                    let _ = response_tx.send(response);
196                }
197
198            }
199        }
200        Ok(())
201    }
202}
203
204#[derive(Eq, Hash, PartialEq, Serialize)]
205struct ContextNoOrigin {
206    name: MetaString,
207    tags: TagSet,
208}
209#[derive(Deserialize)]
210struct StatsQueryParams {
211    collection_duration_secs: u64,
212}
213
214impl DogStatsDStatsAPIHandler {
215    async fn stats_handler(
216        State(state): State<DogStatsDStatsAPIHandlerState>, Query(query): Query<StatsQueryParams>,
217    ) -> (StatusCode, String) {
218        const MAXIMUM_COLLECTION_DURATION_SECS: u64 = 600;
219        if query.collection_duration_secs > MAXIMUM_COLLECTION_DURATION_SECS {
220            return (
221                StatusCode::BAD_REQUEST,
222                format!(
223                    "Collection duration cannot be greater than {} seconds.",
224                    MAXIMUM_COLLECTION_DURATION_SECS
225                ),
226            );
227        }
228
229        let (oneshot_tx, oneshot_rx) = oneshot::channel();
230
231        state
232            .tx
233            .send((oneshot_tx, query.collection_duration_secs))
234            .await
235            .unwrap(); // TODO: use config to set collection period
236
237        match oneshot_rx.await {
238            Ok(stats) => match stats {
239                StatsResponse::Statistics(collected_stats) => match serde_json::to_string(&collected_stats) {
240                    Ok(json) => (StatusCode::OK, json),
241                    Err(e) => (
242                        StatusCode::INTERNAL_SERVER_ERROR,
243                        format!("Failed to serialize stats: {}", e),
244                    ),
245                },
246                StatsResponse::AlreadyRunning { try_after } => (
247                    StatusCode::TOO_MANY_REQUESTS,
248                    format!(
249                        "Statistics collection already active. Please try again in {} seconds.",
250                        try_after
251                    ),
252                ),
253            },
254            Err(_) => (
255                StatusCode::INTERNAL_SERVER_ERROR,
256                "Failed to collect statistics.".to_string(),
257            ),
258        }
259    }
260}
261
262impl APIHandler for DogStatsDStatsAPIHandler {
263    type State = DogStatsDStatsAPIHandlerState;
264
265    fn generate_initial_state(&self) -> Self::State {
266        self.state.clone()
267    }
268
269    fn generate_routes(&self) -> Router<Self::State> {
270        Router::new().route("/dogstatsd/stats", get(Self::stats_handler))
271    }
272}
273
274impl DogStatsDStatisticsConfiguration {
275    /// Creates a new `DogStatsDStatisticsConfiguration`.
276    pub fn new() -> Self {
277        let (tx, rx) = mpsc::channel(4);
278        let state = DogStatsDStatsAPIHandlerState { tx: Arc::new(tx) };
279        let handler = DogStatsDStatsAPIHandler { state };
280
281        Self {
282            api_handler: handler,
283            rx: Arc::new(Mutex::new(rx)),
284        }
285    }
286
287    /// Returns an API handler for DogStatsD API.
288    pub fn api_handler(&self) -> DogStatsDStatsAPIHandler {
289        self.api_handler.clone()
290    }
291}
292
293#[async_trait]
294impl DestinationBuilder for DogStatsDStatisticsConfiguration {
295    fn input_event_type(&self) -> EventType {
296        EventType::Metric
297    }
298
299    async fn build(&self, _context: ComponentContext) -> Result<Box<dyn Destination + Send>, GenericError> {
300        let rx = self.rx.clone().try_lock_owned()?;
301        Ok(Box::new(DogStatsDStats { rx }))
302    }
303}
304
305impl MemoryBounds for DogStatsDStatisticsConfiguration {
306    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
307        builder
308            .minimum()
309            .with_single_value::<DogStatsDStats>("component struct");
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use std::collections::BTreeSet;
316
317    use saluki_context::tags::Tag;
318    use saluki_core::components::ComponentSpawner;
319    use serde_json::json;
320
321    use super::*;
322
323    fn tag_set<const N: usize>(tags: [&'static str; N]) -> TagSet {
324        tags.into_iter().map(Tag::from_static).collect()
325    }
326
327    #[test]
328    fn collected_statistics_serialize_as_flat_metric_entries() {
329        // The `/dogstatsd/stats` endpoint returns `CollectedStatistics`: a start/end window plus a flat array of
330        // per-context samples, where each entry inlines the context (`name`, `tags`) and its `count`/`last_seen`.
331        let mut stats = HashMap::new();
332        stats.insert(
333            ContextNoOrigin {
334                name: MetaString::from("my.counter"),
335                tags: tag_set(["env:prod", "service:web"]),
336            },
337            MetricSample {
338                count: 3,
339                last_seen: 100,
340            },
341        );
342
343        let collected = CollectedStatistics {
344            start_time_unix: 10,
345            end_time_unix: 70,
346            stats: FlattenedStats(stats),
347        };
348        let json = serde_json::to_value(&collected).expect("collected statistics should serialize");
349
350        assert_eq!(json!(10), json["start_time_unix"]);
351        assert_eq!(json!(70), json["end_time_unix"]);
352
353        let entries = json["stats"].as_array().expect("stats should serialize as an array");
354        assert_eq!(1, entries.len());
355        let entry = &entries[0];
356        assert_eq!(json!("my.counter"), entry["name"]);
357        assert_eq!(json!(3), entry["count"]);
358        assert_eq!(json!(100), entry["last_seen"]);
359        let tags = entry["tags"]
360            .as_array()
361            .expect("tags should serialize as an array")
362            .iter()
363            .map(|tag| tag.as_str().expect("each tag should be a string"))
364            .collect::<BTreeSet<_>>();
365        assert_eq!(BTreeSet::from(["env:prod", "service:web"]), tags);
366    }
367
368    #[tokio::test]
369    async fn stats_handler_rejects_excessive_collection_duration() {
370        // The handler caps collection at 600 seconds and short-circuits longer requests with a 400 before any
371        // collection is started.
372        let config = DogStatsDStatisticsConfiguration::new();
373        let state = config.api_handler.state.clone();
374
375        let (status, body) = DogStatsDStatsAPIHandler::stats_handler(
376            State(state),
377            Query(StatsQueryParams {
378                collection_duration_secs: 601,
379            }),
380        )
381        .await;
382
383        assert_eq!(StatusCode::BAD_REQUEST, status);
384        assert_eq!("Collection duration cannot be greater than 600 seconds.", body);
385    }
386
387    #[tokio::test]
388    async fn collection_request_accumulates_metrics_then_responds_on_timeout() {
389        use saluki_core::accounting::{ComponentRegistry, MemoryLimiter};
390        use saluki_core::components::ComponentContext;
391        use saluki_core::data_model::event::metric::Metric;
392        use saluki_core::health::HealthRegistry;
393        use saluki_core::runtime::state::DataspaceRegistry;
394        use saluki_core::runtime::Supervisor;
395        use saluki_core::topology::interconnect::Consumer;
396        use saluki_core::topology::{EventsBuffer, TopologyContext};
397        use tokio::runtime::Handle;
398        use tokio::time::timeout;
399
400        // Build the destination and grab the request sender the API handler would normally use.
401        let config = DogStatsDStatisticsConfiguration::new();
402        let request_tx = config.api_handler.state.tx.clone();
403
404        let component_context = ComponentContext::test_destination("test");
405        let destination = config
406            .build(component_context.clone())
407            .await
408            .expect("dsd_stats destination should build");
409
410        // Wire up the destination context: an events channel we control and an idle health handle.
411        let (events_tx, events_rx) = mpsc::channel::<EventsBuffer>(4);
412        let consumer = Consumer::new(component_context.clone(), events_rx);
413        let topology_context = TopologyContext::new(
414            Arc::from("test"),
415            MemoryLimiter::noop(),
416            HealthRegistry::new(),
417            Handle::current(),
418            DataspaceRegistry::new(),
419        );
420        let health = HealthRegistry::new()
421            .register_component(&saluki_core::support::SubsystemIdentifier::from_dotted("test"))
422            .expect("component was not previously registered");
423        // This component doesn't spawn supervised children yet, so a spawner over a never-run supervisor is
424        // sufficient. Anything that does spawn needs `TestComponentSupervisor` (saluki_core::components::test_util)
425        // instead, otherwise the spawn fails with `SupervisorGone`.
426        let supervisor_handle = Supervisor::new("test").expect("valid supervisor name").handle();
427        let spawner = ComponentSpawner::new(supervisor_handle, Handle::current());
428        let context = DestinationContext::new(
429            &topology_context,
430            &component_context,
431            ComponentRegistry::default(),
432            health,
433            consumer,
434            spawner,
435        );
436
437        let run_handle = tokio::spawn(async move { destination.run(context).await });
438
439        // Start a one-second collection window. Yield afterwards so the current-thread runtime lets the run loop
440        // process the request (marking collection active) before the metrics arrive; otherwise the metrics would be
441        // dropped as "not collecting".
442        let (response_tx, response_rx) = oneshot::channel();
443        request_tx
444            .send((response_tx, 1))
445            .await
446            .expect("collection request should be accepted");
447        tokio::task::yield_now().await;
448
449        // The same context seen twice accumulates a single entry with count 2; a distinct context yields count 1.
450        let mut events = EventsBuffer::default();
451        assert!(events
452            .try_push(Event::Metric(Metric::counter("dsd.stats.repeated", 1.0)))
453            .is_none());
454        assert!(events
455            .try_push(Event::Metric(Metric::counter("dsd.stats.repeated", 1.0)))
456            .is_none());
457        assert!(events
458            .try_push(Event::Metric(Metric::counter("dsd.stats.single", 1.0)))
459            .is_none());
460        events_tx.send(events).await.expect("metrics should be accepted");
461        tokio::task::yield_now().await;
462
463        // The collection window elapses after one second, completing collection and sending the response; the
464        // recv is bounded well above that window so a stalled collection surfaces as a failure, not a hang.
465        let response = timeout(Duration::from_secs(5), response_rx)
466            .await
467            .expect("collection response should arrive before timeout")
468            .expect("collection response channel should remain open");
469
470        let collected = match response {
471            StatsResponse::Statistics(collected) => collected,
472            StatsResponse::AlreadyRunning { .. } => panic!("first request should not report an active collection"),
473        };
474        let samples = collected.stats.0;
475        assert_eq!(2, samples.len(), "each distinct context should have its own sample");
476
477        let repeated = samples
478            .iter()
479            .find(|(ctx, _)| ctx.name.as_ref() == "dsd.stats.repeated")
480            .map(|(_, sample)| sample)
481            .expect("repeated context should be collected");
482        assert_eq!(2, repeated.count, "the repeated context should be counted twice");
483
484        let single = samples
485            .iter()
486            .find(|(ctx, _)| ctx.name.as_ref() == "dsd.stats.single")
487            .map(|(_, sample)| sample)
488            .expect("single context should be collected");
489        assert_eq!(1, single.count);
490
491        // Closing the events channel lets the run loop terminate cleanly.
492        drop(events_tx);
493        timeout(Duration::from_secs(1), run_handle)
494            .await
495            .expect("run task should stop before timeout")
496            .expect("run task should not panic")
497            .expect("run should complete cleanly");
498    }
499}
saluki_components/destinations/dsd_stats/mod.rs

saluki_components/destinations/dsd_stats/
mod.rs
