saluki_components/sources/dogstatsd/replay/
reader.rs

1//! Reader for Datadog DogStatsD capture files.
2//!
3//! Decodes a `.dog` or `.dog.zstd` capture file into the sequence of `UnixDogstatsdMsg` records it
4//! contains, plus the optional `TaggerState` trailer.
5
6use std::{fs, path::Path};
7
8use datadog_protos::agent::{TaggerState, UnixDogstatsdMsg};
9use prost::Message;
10use saluki_error::{generic_error, GenericError};
11
12use super::file_header::{file_version, valid_header, DATADOG_HEADER, MIN_NANO_VERSION, MIN_STATE_VERSION};
13
14const ZSTD_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD];
15const LENGTH_PREFIX_SIZE: usize = 4;
16
17/// Timestamp resolution recorded in a capture file.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum TimestampResolution {
20    /// Timestamps are recorded in whole seconds (file version < 3).
21    Seconds,
22    /// Timestamps are recorded in nanoseconds (file version >= 3).
23    Nanoseconds,
24}
25
26/// Reads back a DogStatsD traffic capture file.
27#[derive(Debug)]
28pub struct TrafficCaptureReader {
29    contents: Vec<u8>,
30    version: u8,
31    offset: usize,
32}
33
34impl TrafficCaptureReader {
35    /// Opens a capture file at the given path.
36    ///
37    /// Detects zstd-compressed inputs by magic bytes and decompresses transparently. Validates the
38    /// Datadog capture header and parses the file version.
39    pub fn from_path(path: &Path) -> Result<Self, GenericError> {
40        let raw =
41            fs::read(path).map_err(|e| generic_error!("Failed to read capture file '{}': {}", path.display(), e))?;
42
43        let contents = if has_zstd_magic(&raw) {
44            zstd::stream::decode_all(raw.as_slice())
45                .map_err(|e| generic_error!("Failed to decompress capture file '{}': {}", path.display(), e))?
46        } else {
47            raw
48        };
49
50        if !valid_header(&contents) {
51            return Err(generic_error!(
52                "Capture file '{}' does not begin with a valid Datadog capture header.",
53                path.display()
54            ));
55        }
56
57        let version = file_version(&contents)?;
58
59        Ok(Self {
60            contents,
61            version,
62            offset: DATADOG_HEADER.len(),
63        })
64    }
65
66    /// Returns the capture file version.
67    pub fn version(&self) -> u8 {
68        self.version
69    }
70
71    /// Returns the timestamp resolution implied by the file version.
72    pub fn timestamp_resolution(&self) -> TimestampResolution {
73        if self.version < MIN_NANO_VERSION {
74            TimestampResolution::Seconds
75        } else {
76            TimestampResolution::Nanoseconds
77        }
78    }
79
80    /// Reads the next captured DogStatsD record from the file.
81    ///
82    /// Returns `Ok(None)` when the stream of records is exhausted: at EOF, at the four-byte state
83    /// separator that introduces the tagger trailer, or at a truncated record boundary.
84    pub fn read_next(&mut self) -> Result<Option<UnixDogstatsdMsg>, GenericError> {
85        if self.offset + LENGTH_PREFIX_SIZE > self.contents.len() {
86            return Ok(None);
87        }
88
89        let size_bytes = &self.contents[self.offset..self.offset + LENGTH_PREFIX_SIZE];
90        let size = u32::from_le_bytes(size_bytes.try_into().expect("length prefix is 4 bytes")) as usize;
91        self.offset += LENGTH_PREFIX_SIZE;
92
93        // The writer emits a zero-length prefix to mark the start of the tagger state trailer; treat
94        // that (and any size that would overrun the buffer) as the end of the record stream.
95        if size == 0 || self.offset + size > self.contents.len() {
96            // A zero-length prefix is the legitimate trailer marker. A non-zero `size` that overruns the buffer is a
97            // corrupt/oversized length prefix being silently read as clean EOF, which drops every following
98            // well-formed record. Surface the corrupt case as distinct from a real trailer.
99            saluki_antithesis::always_or_unreachable!(
100                size == 0,
101                "replay read_next stopped at the real trailer, not on a corrupt length prefix",
102                { "size": size, "offset": self.offset, "len": self.contents.len() }
103            );
104
105            return Ok(None);
106        }
107
108        let msg = UnixDogstatsdMsg::decode(&self.contents[self.offset..self.offset + size])
109            .map_err(|e| generic_error!("Failed to decode captured DogStatsD record: {}", e))?;
110        self.offset += size;
111
112        Ok(Some(msg))
113    }
114
115    /// Reads the tagger state trailer from the end of the capture file.
116    ///
117    /// Returns `Ok(None)` if the file version predates state support, or if the trailer is empty.
118    /// Does not modify the read offset, so it can be called independently of `read_next`.
119    pub fn read_state(&self) -> Result<Option<TaggerState>, GenericError> {
120        if self.version < MIN_STATE_VERSION {
121            return Ok(None);
122        }
123
124        let len = self.contents.len();
125        if len < LENGTH_PREFIX_SIZE {
126            return Ok(None);
127        }
128
129        let size_bytes = &self.contents[len - LENGTH_PREFIX_SIZE..len];
130        let size = u32::from_le_bytes(size_bytes.try_into().expect("length suffix is 4 bytes")) as usize;
131        if size == 0 {
132            return Ok(None);
133        }
134
135        if size + LENGTH_PREFIX_SIZE > len {
136            return Err(generic_error!(
137                "Tagger state trailer size ({}) exceeds capture file length ({}).",
138                size,
139                len
140            ));
141        }
142
143        let state_start = len - LENGTH_PREFIX_SIZE - size;
144        let state = TaggerState::decode(&self.contents[state_start..len - LENGTH_PREFIX_SIZE])
145            .map_err(|e| generic_error!("Failed to decode tagger state trailer: {}", e))?;
146        Ok(Some(state))
147    }
148}
149
150fn has_zstd_magic(buf: &[u8]) -> bool {
151    buf.len() >= ZSTD_MAGIC.len() && buf[..ZSTD_MAGIC.len()] == ZSTD_MAGIC
152}
153
154#[cfg(test)]
155mod tests {
156    use std::{
157        fs,
158        path::PathBuf,
159        sync::Arc,
160        thread,
161        time::{Duration, SystemTime, UNIX_EPOCH},
162    };
163
164    use saluki_common::collections::FastHashMap;
165    use saluki_context::{
166        origin::OriginTagCardinality,
167        tags::{SharedTagSet, Tag, TagSet},
168    };
169    use saluki_env::{
170        workload::{origin::ResolvedOrigin, EntityId},
171        WorkloadProvider,
172    };
173
174    use super::*;
175    use crate::sources::dogstatsd::replay::writer::{CaptureRecord, CaptureTargetDir, TrafficCaptureWriter};
176
177    #[derive(Default)]
178    struct MockWorkloadProvider {
179        entities: FastHashMap<EntityId, SharedTagSet>,
180    }
181
182    impl MockWorkloadProvider {
183        fn with_entity(entity_id: EntityId, tags: &[&str]) -> Self {
184            let mut entities = FastHashMap::default();
185            entities.insert(entity_id, shared_tags(tags));
186            Self { entities }
187        }
188    }
189
190    impl WorkloadProvider for MockWorkloadProvider {
191        fn get_tags_for_entity(
192            &self, entity_id: &EntityId, _cardinality: OriginTagCardinality,
193        ) -> Option<SharedTagSet> {
194            self.entities.get(entity_id).cloned()
195        }
196
197        fn get_resolved_origin(&self, _origin: saluki_context::origin::RawOrigin<'_>) -> Option<ResolvedOrigin> {
198            None
199        }
200    }
201
202    #[test]
203    fn plain_capture_round_trip() {
204        let (path, _dir_guard) = run_capture(1, false, &[sample_record(100, b"metric.a:1|c", 11)]);
205
206        let mut reader = TrafficCaptureReader::from_path(&path).expect("reader should open");
207        assert_eq!(reader.version(), 3);
208        assert_eq!(reader.timestamp_resolution(), TimestampResolution::Nanoseconds);
209
210        let msg = reader.read_next().expect("read should succeed").expect("first record");
211        assert_eq!(msg.timestamp, 100);
212        assert_eq!(msg.payload, b"metric.a:1|c");
213        assert_eq!(msg.payload_size, msg.payload.len() as i32);
214        assert_eq!(msg.pid, 11);
215
216        assert!(reader.read_next().expect("read should succeed").is_none());
217    }
218
219    #[test]
220    fn compressed_capture_round_trip() {
221        let (path, _dir_guard) = run_capture(
222            1,
223            true,
224            &[
225                sample_record(1, b"metric.a:1|c", 1),
226                sample_record(2, b"metric.b:2|c", 1),
227                sample_record(3, b"metric.c:3|c", 1),
228            ],
229        );
230
231        let mut reader = TrafficCaptureReader::from_path(&path).expect("reader should open");
232
233        for expected_ts in [1, 2, 3] {
234            let msg = reader.read_next().expect("read should succeed").expect("record");
235            assert_eq!(msg.timestamp, expected_ts);
236        }
237
238        assert!(reader.read_next().expect("read should succeed").is_none());
239    }
240
241    #[test]
242    fn read_next_stops_at_state_separator() {
243        let (path, _dir_guard) = run_capture(1, false, &[sample_record(7, b"x:1|c", 5)]);
244
245        let mut reader = TrafficCaptureReader::from_path(&path).expect("reader should open");
246        assert!(reader.read_next().expect("first record").is_some());
247
248        // Subsequent calls must yield None rather than attempting to decode the trailer as a record.
249        assert!(reader.read_next().expect("trailer boundary").is_none());
250        assert!(reader.read_next().expect("idempotent EOF").is_none());
251    }
252
253    #[test]
254    fn truncated_record_returns_none() {
255        let (path, _dir_guard) = run_capture(1, false, &[sample_record(1, b"metric.a:1|c", 1)]);
256
257        let bytes = fs::read(&path).expect("capture readable");
258        let truncated_path = path.with_extension("truncated");
259        // Drop the last 8 bytes so the trailer length prefix is gone and a record is incomplete.
260        fs::write(&truncated_path, &bytes[..bytes.len().saturating_sub(8)]).expect("write truncated");
261
262        let mut reader = TrafficCaptureReader::from_path(&truncated_path).expect("reader should open");
263        let _ = reader.read_next();
264        // Whatever the reader recovered, the next call must terminate cleanly rather than error.
265        assert!(reader.read_next().expect("clean EOF on truncation").is_none());
266    }
267
268    #[test]
269    fn bad_header_is_rejected() {
270        let tmp = unique_path("reader-bad-header");
271        fs::write(&tmp, b"this is not a capture file").expect("write garbage");
272
273        let err = TrafficCaptureReader::from_path(&tmp).expect_err("bad header should fail");
274        assert!(err.to_string().contains("Datadog capture header"));
275
276        let _ = fs::remove_file(&tmp);
277    }
278
279    #[test]
280    fn read_state_recovers_entity_tags() {
281        let target_dir = unique_dir("reader-state");
282        let workload = Arc::new(MockWorkloadProvider::with_entity(
283            EntityId::Container("container-xyz".into()),
284            &["env:prod", "service:api"],
285        ));
286        let writer = TrafficCaptureWriter::with_workload_provider(1, Some(workload));
287
288        let path = writer
289            .start_capture(
290                CaptureTargetDir::Explicit(target_dir.clone()),
291                Duration::from_millis(250),
292                false,
293            )
294            .expect("capture should start");
295        assert!(writer.enqueue(CaptureRecord {
296            timestamp_ns: 1,
297            payload: b"metric:1|c".to_vec(),
298            pid: Some(99),
299            ancillary: Vec::new(),
300            container_id: Some("container_id://container-xyz".to_string()),
301        }));
302        writer.stop_capture();
303        wait_until_inactive(&writer);
304
305        let reader = TrafficCaptureReader::from_path(&path).expect("reader should open");
306        let state = reader
307            .read_state()
308            .expect("state should decode")
309            .expect("state present");
310        let entity = state
311            .state
312            .get("container_id://container-xyz")
313            .expect("captured entity present");
314        assert_eq!(
315            entity.low_cardinality_tags,
316            vec!["env:prod".to_string(), "service:api".to_string()]
317        );
318        assert_eq!(
319            state.pid_map.get(&99).map(String::as_str),
320            Some("container_id://container-xyz")
321        );
322
323        let _ = fs::remove_dir_all(target_dir);
324    }
325
326    fn run_capture(queue_depth: usize, compressed: bool, records: &[CaptureRecord]) -> (PathBuf, DirGuard) {
327        let target_dir = unique_dir("reader-capture");
328        let writer = TrafficCaptureWriter::new(queue_depth);
329
330        let path = writer
331            .start_capture(
332                CaptureTargetDir::Explicit(target_dir.clone()),
333                Duration::from_millis(500),
334                compressed,
335            )
336            .expect("capture should start");
337
338        for record in records {
339            assert!(writer.enqueue(clone_record(record)));
340        }
341        writer.stop_capture();
342        wait_until_inactive(&writer);
343
344        (path, DirGuard { path: target_dir })
345    }
346
347    fn clone_record(record: &CaptureRecord) -> CaptureRecord {
348        CaptureRecord {
349            timestamp_ns: record.timestamp_ns,
350            payload: record.payload.clone(),
351            pid: record.pid,
352            ancillary: record.ancillary.clone(),
353            container_id: record.container_id.clone(),
354        }
355    }
356
357    fn sample_record(timestamp_ns: i64, payload: &[u8], pid: i32) -> CaptureRecord {
358        CaptureRecord {
359            timestamp_ns,
360            payload: payload.to_vec(),
361            pid: Some(pid),
362            ancillary: Vec::new(),
363            container_id: Some(format!("container_id://container-{}", pid)),
364        }
365    }
366
367    fn shared_tags(tags: &[&str]) -> SharedTagSet {
368        TagSet::from_iter(tags.iter().copied().map(Tag::from)).into_shared()
369    }
370
371    fn wait_until_inactive(writer: &TrafficCaptureWriter) {
372        let deadline = std::time::Instant::now() + Duration::from_secs(2);
373        while writer.is_ongoing() && std::time::Instant::now() < deadline {
374            thread::sleep(Duration::from_millis(10));
375        }
376        assert!(!writer.is_ongoing(), "capture writer did not stop in time");
377    }
378
379    fn unique_dir(label: &str) -> PathBuf {
380        let path = unique_path(label);
381        fs::create_dir_all(&path).expect("test directory should be created");
382        path
383    }
384
385    fn unique_path(label: &str) -> PathBuf {
386        let timestamp = SystemTime::now()
387            .duration_since(UNIX_EPOCH)
388            .expect("clock should be after epoch")
389            .as_nanos();
390        std::env::temp_dir().join(format!("saluki-{}-{}-{}", label, std::process::id(), timestamp))
391    }
392
393    struct DirGuard {
394        path: PathBuf,
395    }
396
397    impl Drop for DirGuard {
398        fn drop(&mut self) {
399            let _ = fs::remove_dir_all(&self.path);
400        }
401    }
402}
saluki_components/sources/dogstatsd/replay/reader.rs

saluki_components/sources/dogstatsd/replay/
reader.rs
