saluki_components/sources/dogstatsd/replay/
capture_api.rs

1//! HTTP API handler for the DogStatsD capture control surface.
2//!
3//! Exposes `POST /dogstatsd/capture/trigger` on the privileged API to start a capture session.
4
5use std::path::Path;
6
7use saluki_api::{
8    extract::State,
9    routing::{post, Router},
10    APIHandler, Json, StatusCode,
11};
12use saluki_config::parse_duration;
13use serde::{Deserialize, Serialize};
14
15use super::DogStatsDCaptureControl;
16
17/// Request body for `POST /dogstatsd/capture/trigger`.
18#[derive(Deserialize)]
19pub struct CaptureTriggerBody {
20    /// Duration of the capture, parsed by `parse_duration` (for example, `"10s"`, `"500ms"`).
21    pub duration: String,
22
23    /// Optional override for the capture output directory. When omitted, the source's
24    /// configured default directory is used.
25    #[serde(default)]
26    pub path: Option<String>,
27
28    /// Whether the capture file should be zstd-compressed.
29    #[serde(default)]
30    pub compressed: bool,
31}
32
33/// Response body for `POST /dogstatsd/capture/trigger`.
34#[derive(Serialize)]
35pub struct CaptureTriggerResponseBody {
36    /// Absolute path the capture is being written to.
37    pub path: String,
38}
39
40/// API handler for the DogStatsD capture control surface.
41#[derive(Clone)]
42pub struct DogStatsDCaptureAPIHandler {
43    capture_control: DogStatsDCaptureControl,
44}
45
46impl DogStatsDCaptureAPIHandler {
47    /// Creates a new handler bound to the given capture control.
48    pub fn new(capture_control: DogStatsDCaptureControl) -> Self {
49        Self { capture_control }
50    }
51
52    async fn trigger_handler(
53        State(capture_control): State<DogStatsDCaptureControl>, Json(body): Json<CaptureTriggerBody>,
54    ) -> Result<Json<CaptureTriggerResponseBody>, (StatusCode, String)> {
55        let duration = parse_duration(&body.duration).map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
56        let requested_dir = body.path.as_deref().map(Path::new);
57
58        let capture_path = capture_control
59            .start_capture(requested_dir, duration, body.compressed)
60            .map_err(|e| (StatusCode::PRECONDITION_FAILED, e.to_string()))?;
61
62        Ok(Json(CaptureTriggerResponseBody {
63            path: capture_path.display().to_string(),
64        }))
65    }
66}
67
68impl APIHandler for DogStatsDCaptureAPIHandler {
69    type State = DogStatsDCaptureControl;
70
71    fn generate_initial_state(&self) -> Self::State {
72        self.capture_control.clone()
73    }
74
75    fn generate_routes(&self) -> Router<Self::State> {
76        Router::new().route("/dogstatsd/capture/trigger", post(Self::trigger_handler))
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::super::test_support::{unique_dir, wait_until_inactive};
83    use super::*;
84    use crate::sources::dogstatsd::replay::TrafficCapture;
85
86    fn trigger_body(duration: &str) -> CaptureTriggerBody {
87        CaptureTriggerBody {
88            duration: duration.to_string(),
89            path: None,
90            compressed: false,
91        }
92    }
93
94    #[tokio::test]
95    async fn trigger_rejects_unbound_control() {
96        // With no capture runtime bound to the control, the handler surfaces the "capture control unavailable" error
97        // as a precondition failure.
98        let control = DogStatsDCaptureControl::new();
99
100        let err = DogStatsDCaptureAPIHandler::trigger_handler(State(control), Json(trigger_body("50ms")))
101            .await
102            .err()
103            .expect("unbound control should fail");
104
105        assert_eq!(err.0, StatusCode::PRECONDITION_FAILED);
106    }
107
108    #[tokio::test]
109    async fn trigger_rejects_invalid_duration() {
110        // A malformed duration is rejected up front (BAD_REQUEST), before the capture runtime is consulted.
111        let control = DogStatsDCaptureControl::new();
112        control.bind(TrafficCapture::new(unique_dir("capture-api-bad-duration"), 1));
113
114        let err = DogStatsDCaptureAPIHandler::trigger_handler(State(control), Json(trigger_body("not-a-duration")))
115            .await
116            .err()
117            .expect("invalid duration should fail");
118
119        assert_eq!(err.0, StatusCode::BAD_REQUEST);
120    }
121
122    #[tokio::test]
123    async fn trigger_starts_capture_and_returns_path() {
124        let target_dir = unique_dir("capture-api-start");
125        let control = DogStatsDCaptureControl::new();
126        control.bind(TrafficCapture::new(target_dir.clone(), 1));
127
128        let Json(response) =
129            DogStatsDCaptureAPIHandler::trigger_handler(State(control.clone()), Json(trigger_body("100ms")))
130                .await
131                .expect("capture should start");
132        assert!(
133            response
134                .path
135                .starts_with(target_dir.to_str().expect("path should be utf-8")),
136            "returned path {:?} should be inside the configured directory {:?}",
137            response.path,
138            target_dir
139        );
140        assert!(response.path.contains("datadog-capture"));
141        assert!(control.is_ongoing().expect("control should be bound"));
142
143        wait_until_inactive(|| control.is_ongoing().expect("control should be bound"));
144        let _ = std::fs::remove_dir_all(target_dir);
145    }
146
147    #[tokio::test]
148    async fn trigger_rejects_concurrent_capture() {
149        let target_dir = unique_dir("capture-api-concurrent");
150        let control = DogStatsDCaptureControl::new();
151        control.bind(TrafficCapture::new(target_dir.clone(), 1));
152
153        let _ = DogStatsDCaptureAPIHandler::trigger_handler(State(control.clone()), Json(trigger_body("1s")))
154            .await
155            .expect("first capture should start");
156
157        let err = DogStatsDCaptureAPIHandler::trigger_handler(State(control.clone()), Json(trigger_body("1s")))
158            .await
159            .err()
160            .expect("second concurrent capture should fail");
161        assert_eq!(err.0, StatusCode::PRECONDITION_FAILED);
162        assert!(
163            err.1.contains("capture already in progress"),
164            "unexpected error: {}",
165            err.1
166        );
167
168        control.stop_capture().expect("stop should succeed");
169        wait_until_inactive(|| control.is_ongoing().expect("control should be bound"));
170        let _ = std::fs::remove_dir_all(target_dir);
171    }
172}