saluki_components/sources/dogstatsd/replay/
capture_api.rs1use 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#[derive(Deserialize)]
19pub struct CaptureTriggerBody {
20 pub duration: String,
22
23 #[serde(default)]
26 pub path: Option<String>,
27
28 #[serde(default)]
30 pub compressed: bool,
31}
32
33#[derive(Serialize)]
35pub struct CaptureTriggerResponseBody {
36 pub path: String,
38}
39
40#[derive(Clone)]
42pub struct DogStatsDCaptureAPIHandler {
43 capture_control: DogStatsDCaptureControl,
44}
45
46impl DogStatsDCaptureAPIHandler {
47 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 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 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}