1use dial9_trace_format::encoder::{Encoder, RawEncoder};
2
3use crate::clock::clock_pair;
4use crate::collector::Batch;
5use crate::format::{ClockSyncEvent, SegmentMetadataEvent};
6use crate::fs::{ActiveHandle, Fs, RemoveReason};
7use crate::primitives::fs;
8use crate::rate_limit::rate_limited;
9use crate::sealed::SegmentRef;
10use std::collections::VecDeque;
11use std::io::BufWriter;
12use std::marker::PhantomData;
13use std::path::{Path, PathBuf};
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16
17use metrique_timesource::time_source;
18
19mod mode_sealed {
20 pub trait Sealed {}
21}
22
23pub trait BufferMode: mode_sealed::Sealed + Send + 'static {
26 const IS_DISK: bool;
28}
29
30#[derive(Debug)]
32#[non_exhaustive]
33pub struct Disk;
34#[derive(Debug)]
36#[non_exhaustive]
37pub struct Memory;
38
39impl mode_sealed::Sealed for Disk {}
40impl mode_sealed::Sealed for Memory {}
41impl BufferMode for Disk {
42 const IS_DISK: bool = true;
43}
44impl BufferMode for Memory {
45 const IS_DISK: bool = false;
46}
47
48pub type DiskBuffer = SegmentWriter<Disk>;
50pub type MemoryBuffer = SegmentWriter<Memory>;
52
53const DIAL9_VERSION_KEY: &str = "dial9.dial9-tokio-telemetry.version";
57
58const DIAL9_VERSION_VALUE: &str = env!("CARGO_PKG_VERSION");
60
61const PROCESS_AVAILABLE_PARALLELISM_KEY: &str = "process.available_parallelism";
64
65#[derive(Clone)]
66struct SegmentMetadata {
67 entries: Vec<(String, String)>,
68}
69
70impl Default for SegmentMetadata {
71 fn default() -> Self {
72 let mut entries = vec![(
73 DIAL9_VERSION_KEY.to_string(),
74 DIAL9_VERSION_VALUE.to_string(),
75 )];
76 match std::thread::available_parallelism() {
77 Ok(parallelism) => entries.push((
78 PROCESS_AVAILABLE_PARALLELISM_KEY.to_string(),
79 parallelism.get().to_string(),
80 )),
81 Err(e) => rate_limited!(Duration::from_secs(60), {
82 tracing::warn!("failed to read process available parallelism: {e}");
83 }),
84 }
85 Self { entries }
86 }
87}
88
89impl SegmentMetadata {
90 fn new(user_entries: Vec<(String, String)>) -> Self {
93 let mut s = Self::default();
94 s.merge(user_entries.into_iter());
95 s
96 }
97
98 fn merge(&mut self, entries: impl Iterator<Item = (String, String)>) -> bool {
110 let mut merged: Vec<(String, String)> = entries.collect();
111 for (k, v) in &self.entries {
112 if !merged.iter().any(|(mk, _)| mk == k) {
113 merged.push((k.clone(), v.clone()));
114 }
115 }
116 if merged == self.entries {
117 return false;
118 }
119 self.entries = merged;
120 true
121 }
122}
123
124const DEFAULT_ROTATION_PERIOD: Duration = Duration::from_secs(60);
126
127const DEFAULT_DRAIN_INTERVAL: Duration = Duration::from_secs(30);
129
130const SEGMENT_STEM: &str = "trace";
133
134const BYTES_PER_MIB: u64 = 1024 * 1024;
135
136const MAX_FILE_SIZE_CAP: u64 = 100 * BYTES_PER_MIB;
140
141fn derive_max_file_size(max_total_size: u64) -> u64 {
145 (max_total_size / 4).min(MAX_FILE_SIZE_CAP)
146}
147
148pub struct SegmentWriter<Mode: BufferMode = Disk> {
175 dir: PathBuf,
177 stem: String,
179 max_file_size: u64,
180 max_total_size: u64,
181 rotation_period: Duration,
184 next_rotation_time: Option<Instant>,
187 closed_files: VecDeque<(SegmentRef, u64)>,
190 active_path: PathBuf,
193 state: WriterState,
194 next_index: u32,
195 segment_metadata: SegmentMetadata,
198 dropped_events: usize,
200 has_real_events: bool,
203 drain_interval: Duration,
206 next_drain_time: Instant,
208 fs: Arc<Fs>,
210 boot_id: Option<String>,
211 _namespace_lock: Option<std::fs::File>,
212 _mode: PhantomData<Mode>,
213}
214
215impl<M: BufferMode> std::fmt::Debug for SegmentWriter<M> {
216 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
217 f.debug_struct("SegmentWriter")
218 .field("dir", &self.dir)
219 .field("stem", &self.stem)
220 .field("max_file_size", &self.max_file_size)
221 .field("max_total_size", &self.max_total_size)
222 .finish_non_exhaustive()
223 }
224}
225
226#[allow(clippy::large_enum_variant)]
229enum WriterState {
230 Active {
232 writer: RawEncoder<BufWriter<ActiveHandle>>,
233 need_metadata: bool,
234 },
235
236 Finished,
238}
239
240#[bon::bon]
241impl SegmentWriter<Disk> {
242 #[builder(builder_type = DiskBufferBuilder, finish_fn = build)]
247 pub fn builder(
248 base_path: impl Into<PathBuf>,
249 max_file_size: Option<u64>,
252 max_total_size: u64,
253 rotation_period: Option<Duration>,
257 segment_metadata: Option<Vec<(String, String)>>,
258 ) -> std::io::Result<Self> {
259 Self::create(
260 base_path,
261 max_file_size.unwrap_or_else(|| derive_max_file_size(max_total_size)),
262 max_total_size,
263 rotation_period.unwrap_or(DEFAULT_ROTATION_PERIOD),
264 segment_metadata
265 .map(SegmentMetadata::new)
266 .unwrap_or_default(),
267 )
268 }
269
270 fn create(
271 base_path: impl Into<PathBuf>,
272 max_file_size: u64,
273 max_total_size: u64,
274 rotation_period: Duration,
275 segment_metadata: SegmentMetadata,
276 ) -> std::io::Result<Self> {
277 if rotation_period == Duration::from_secs(0) {
278 return Err(std::io::Error::other("Rotation period must not be zero"));
279 }
280 let dir = base_path.into();
282 if !dir.as_os_str().is_empty() {
283 fs::create_dir_all(&dir)?;
284 }
285 let stem = SEGMENT_STEM.to_string();
286 let fs = Fs::new_disk(&dir, stem.as_str());
287 let discovered = fs.discover_existing()?;
288 let first_index = discovered.next_active_index;
289 let next_index = first_index
290 .checked_add(1)
291 .ok_or_else(|| std::io::Error::other("trace segment index overflow"))?;
292 let first_path = Self::active_path(&dir, &stem, first_index);
293 let handle = fs.create_segment(&first_path)?;
294 let state = Self::prepare_segment(BufWriter::new(handle))?;
295 let now = time_source().instant().as_std();
296 let drain_interval = rotation_period.min(DEFAULT_DRAIN_INTERVAL);
297
298 let mut writer = Self {
299 dir,
300 stem,
301 max_file_size,
302 max_total_size,
303 rotation_period,
304 next_rotation_time: Self::next_rotation_from(now, rotation_period),
305 closed_files: discovered.closed_files,
306 active_path: first_path,
307 state,
308 next_index,
309 segment_metadata,
310 dropped_events: 0,
311 has_real_events: false,
312 drain_interval,
313 next_drain_time: now + drain_interval,
314 fs,
315 boot_id: None,
316 _namespace_lock: None,
317 _mode: PhantomData,
318 };
319 writer.evict_oldest()?;
322 Ok(writer)
323 }
324
325 pub fn set_namespace(&mut self, boot_id: String, lock: std::fs::File) {
329 self.boot_id = Some(boot_id);
330 self._namespace_lock = Some(lock);
331 }
332
333 pub fn single_file(path: impl Into<PathBuf>) -> std::io::Result<Self> {
341 let path = path.into();
342 let dir = path
345 .parent()
346 .filter(|p| !p.as_os_str().is_empty())
347 .unwrap_or(Path::new("."))
348 .to_path_buf();
349 let stem = path
350 .file_stem()
351 .and_then(|s| s.to_str())
352 .unwrap_or(SEGMENT_STEM)
353 .to_string();
354 let fs = Fs::new_disk(&dir, stem.as_str());
355 let active_path = Self::active_path(&dir, &stem, 0);
356 let handle = fs.create_segment(&active_path)?;
357 let state = Self::prepare_segment(BufWriter::new(handle))?;
358 let now = time_source().instant().as_std();
359
360 Ok(Self {
361 dir,
362 stem,
363 max_file_size: u64::MAX,
364 max_total_size: u64::MAX,
365 rotation_period: Duration::MAX,
366 next_rotation_time: None,
367 closed_files: VecDeque::new(),
368 active_path,
369 state,
370 next_index: 1,
371 segment_metadata: SegmentMetadata::default(),
372 dropped_events: 0,
373 has_real_events: false,
374 drain_interval: DEFAULT_DRAIN_INTERVAL,
375 next_drain_time: now + DEFAULT_DRAIN_INTERVAL,
376 fs,
377 boot_id: None,
378 _namespace_lock: None,
379 _mode: PhantomData,
380 })
381 }
382}
383
384fn pick_segment_size(max_total_size: u64) -> u64 {
387 const MIN_SLOTS: u64 = 8;
388 (max_total_size / MIN_SLOTS).max(1)
389}
390
391#[bon::bon]
392impl SegmentWriter<Memory> {
393 pub fn new(max_total_size: u64) -> std::io::Result<Self> {
400 Self::create_in_memory(
401 max_total_size,
402 pick_segment_size(max_total_size),
403 DEFAULT_ROTATION_PERIOD,
404 SegmentMetadata::default(),
405 )
406 }
407
408 #[builder(builder_type = MemoryBufferBuilder, finish_fn = build)]
410 pub fn builder(
411 max_total_size: u64,
412 max_segment_size: Option<u64>,
414 rotation_period: Option<Duration>,
416 segment_metadata: Option<Vec<(String, String)>>,
417 ) -> std::io::Result<Self> {
418 let seg_size = max_segment_size.unwrap_or_else(|| pick_segment_size(max_total_size));
419 Self::create_in_memory(
420 max_total_size,
421 seg_size,
422 rotation_period.unwrap_or(DEFAULT_ROTATION_PERIOD),
423 segment_metadata
424 .map(SegmentMetadata::new)
425 .unwrap_or_default(),
426 )
427 }
428
429 fn create_in_memory(
430 max_total_size: u64,
431 max_segment_size: u64,
432 rotation_period: Duration,
433 segment_metadata: SegmentMetadata,
434 ) -> std::io::Result<Self> {
435 if max_total_size == 0 {
436 return Err(std::io::Error::new(
437 std::io::ErrorKind::InvalidInput,
438 "max_total_size must be > 0",
439 ));
440 }
441 if max_segment_size == 0 {
442 return Err(std::io::Error::new(
443 std::io::ErrorKind::InvalidInput,
444 "max_segment_size must be > 0",
445 ));
446 }
447 if rotation_period == Duration::from_secs(0) {
448 return Err(std::io::Error::other("Rotation period must not be zero"));
449 }
450 let min_total = (crate::fs::PIPELINE_RESERVE_SEGMENTS + 1).saturating_mul(max_segment_size);
453 if max_total_size < min_total {
454 return Err(std::io::Error::new(
455 std::io::ErrorKind::InvalidInput,
456 format!(
457 "max_total_size ({max_total_size}) must be >= {min_total} \
458 ({} × max_segment_size: 1 active + 1 in-flight + 1 ring slot)",
459 crate::fs::PIPELINE_RESERVE_SEGMENTS + 1
460 ),
461 ));
462 }
463 let fs = Fs::new_in_memory(max_total_size, max_segment_size)?;
464 let dir = PathBuf::from("mem");
467 let stem = SEGMENT_STEM.to_string();
468 let active_path = Self::active_path(&dir, &stem, 0);
469 let handle = fs.create_segment(&active_path)?;
470 let state = Self::prepare_segment(BufWriter::new(handle))?;
471 let now = time_source().instant().as_std();
472 let drain_interval = rotation_period.min(DEFAULT_DRAIN_INTERVAL);
474
475 Ok(Self {
476 dir,
477 stem,
478 max_file_size: max_segment_size,
479 max_total_size,
480 rotation_period,
481 next_rotation_time: Self::next_rotation_from(now, rotation_period),
482 closed_files: VecDeque::new(),
483 active_path,
484 state,
485 next_index: 1,
486 segment_metadata,
487 dropped_events: 0,
488 has_real_events: false,
489 drain_interval,
490 next_drain_time: now + drain_interval,
491 fs,
492 boot_id: None,
493 _namespace_lock: None,
494 _mode: PhantomData,
495 })
496 }
497}
498
499impl<M: BufferMode> SegmentWriter<M> {
500 pub fn boot_id(&self) -> Option<&str> {
503 self.boot_id.as_deref()
504 }
505
506 pub fn trace_dir(&self) -> &Path {
511 &self.dir
512 }
513
514 pub fn trace_stem(&self) -> &str {
516 &self.stem
517 }
518
519 pub fn current_active_path(&self) -> &Path {
521 &self.active_path
522 }
523
524 fn prepare_segment(writer: BufWriter<ActiveHandle>) -> std::io::Result<WriterState> {
528 let mut encoder = Encoder::new_to(writer)?;
529 let (mono, real) = clock_pair();
530 encoder.write(&ClockSyncEvent {
531 timestamp_ns: mono,
532 realtime_ns: real,
533 })?;
534 Ok(WriterState::Active {
535 writer: encoder.into_raw_encoder(),
536 need_metadata: true,
537 })
538 }
539
540 fn write_metadata_if_needed(&mut self) -> std::io::Result<()> {
541 match &mut self.state {
542 WriterState::Active {
543 writer,
544 need_metadata,
545 } => {
546 if *need_metadata {
547 Self::write_segment_metadata(writer, &self.segment_metadata.entries)?;
548 }
549 *need_metadata = false;
550 Ok(())
551 }
552 WriterState::Finished => Ok(()),
553 }
554 }
555
556 fn write_segment_metadata(
559 writer: &mut RawEncoder<BufWriter<ActiveHandle>>,
560 entries: &[(String, String)],
561 ) -> std::io::Result<()> {
562 let mut enc = Encoder::new();
563 let entries = entries.to_vec();
564 let (mono, real) = clock_pair();
565 enc.write(&SegmentMetadataEvent {
566 timestamp_ns: mono,
567 entries,
568 })?;
569 enc.write(&ClockSyncEvent {
570 timestamp_ns: mono,
571 realtime_ns: real,
572 })?;
573 writer.write_raw(&enc.finish())?;
574 Ok(())
575 }
576
577 fn active_path(dir: &Path, stem: &str, index: u32) -> PathBuf {
579 dir.join(format!("{stem}.{index}.bin.active"))
580 }
581
582 fn next_rotation_from(now: Instant, period: Duration) -> Option<Instant> {
585 (period != Duration::MAX).then(|| now + period)
586 }
587
588 fn rotate(&mut self) -> std::io::Result<()> {
589 if matches!(self.state, WriterState::Finished) {
590 return Ok(());
591 }
592
593 let now = time_source().instant().as_std();
597 self.next_rotation_time = Self::next_rotation_from(now, self.rotation_period);
598 self.next_drain_time = now + self.drain_interval;
599
600 let WriterState::Active {
602 writer: mut raw, ..
603 } = std::mem::replace(&mut self.state, WriterState::Finished)
604 else {
605 return Ok(());
606 };
607
608 let _ = raw.flush();
611 let closed_size = raw.bytes_written();
612 let current_index = self.next_index - 1;
613
614 let bw: BufWriter<ActiveHandle> = raw.into_inner();
616 let handle: ActiveHandle = bw
617 .into_inner()
618 .unwrap_or_else(|e| e.into_inner().into_parts().0);
619
620 match self.fs.seal(handle, &self.active_path, current_index) {
624 Ok(seg_ref) => {
625 if M::IS_DISK {
626 self.closed_files.push_back((seg_ref, closed_size));
627 }
628 }
629 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
630 rate_limited!(Duration::from_secs(60), {
631 tracing::warn!(
632 "active trace file {} disappeared before sealing; \
633 abandoning segment and starting a fresh one",
634 self.active_path.display()
635 );
636 });
637 }
638 Err(e) => {
639 return Err(e);
641 }
642 }
643
644 let new_path = Self::active_path(&self.dir, &self.stem, self.next_index);
645 self.next_index += 1;
646
647 let handle: ActiveHandle = self.fs.create_segment(&new_path)?;
652
653 self.state = match Self::prepare_segment(BufWriter::new(handle)) {
654 Ok(s) => s,
655 Err(e) => {
656 let _ = self.fs.remove_active(&new_path);
657 return Err(e);
658 }
659 };
660 self.active_path = new_path;
661 self.has_real_events = false;
662
663 tracing::debug!(
664 segment_index = self.next_index - 1,
665 "rotated to new trace segment"
666 );
667 self.evict_oldest()?;
668 Ok(())
669 }
670
671 fn total_size(&self) -> u64 {
674 if !M::IS_DISK {
675 return 0;
676 }
677 let closed: u64 = self.closed_files.iter().map(|(_, s)| s).sum();
678 let active = match &self.state {
679 WriterState::Active { writer, .. } => writer.bytes_written(),
680 WriterState::Finished => 0,
681 };
682 closed + active
683 }
684
685 fn evict_oldest(&mut self) -> std::io::Result<()> {
686 if !M::IS_DISK {
687 return Ok(());
688 }
689 while self.total_size() > self.max_total_size && !self.closed_files.is_empty() {
691 if let Some((seg_ref, _size)) = self.closed_files.pop_front() {
692 self.fs.remove_sealed(&seg_ref, RemoveReason::Eviction);
693 }
694 }
695 if self.total_size() > self.max_total_size {
697 self.state = WriterState::Finished;
698 }
699 Ok(())
700 }
701
702 fn maybe_rotate(&mut self) -> std::io::Result<()> {
705 let WriterState::Active { writer: raw, .. } = &self.state else {
706 return Ok(());
707 };
708 if raw.bytes_written() > self.max_file_size {
709 self.rotate()?;
710 }
711 Ok(())
712 }
713}
714
715impl<M: BufferMode> SegmentWriter<M> {
716 #[cfg(feature = "pipeline")]
719 pub(crate) fn fs_handle(&self) -> Option<Arc<Fs>> {
720 Some(Arc::clone(&self.fs))
721 }
722
723 pub fn flush(&mut self) -> std::io::Result<()> {
725 if let WriterState::Active { writer: raw, .. } = &mut self.state {
726 raw.flush()?;
727 }
728 Ok(())
729 }
730
731 #[cfg(test)]
732 pub(crate) fn segment_metadata(&self) -> &[(String, String)] {
733 &self.segment_metadata.entries
734 }
735
736 pub fn update_segment_metadata(&mut self, entries: impl IntoIterator<Item = (String, String)>) {
741 if self.segment_metadata.merge(entries.into_iter()) {
742 match &mut self.state {
743 WriterState::Active { need_metadata, .. } => *need_metadata = true,
744 WriterState::Finished => {}
745 }
746 }
747 }
748
749 pub(crate) fn write_current_segment_metadata(&mut self) -> std::io::Result<()> {
750 self.write_metadata_if_needed()
751 }
752
753 pub(crate) fn should_drain(&self) -> bool {
754 self.has_real_events && time_source().instant().as_std() >= self.next_drain_time
755 }
756
757 pub(crate) fn drained(&mut self) -> std::io::Result<bool> {
758 if !self.has_real_events {
759 return Ok(false);
760 }
761 let now = time_source().instant().as_std();
762 if self
763 .next_rotation_time
764 .is_some_and(|deadline| now >= deadline)
765 {
766 self.rotate()?;
767 return Ok(true);
768 }
769 self.next_drain_time = now + self.drain_interval;
771 Ok(false)
772 }
773
774 pub fn finalize(&mut self) -> std::io::Result<()> {
777 if matches!(self.state, WriterState::Finished) {
778 rate_limited!(Duration::from_secs(60), {
779 tracing::warn!("writer is already closed.");
780 });
781 self.fs.mark_writer_done();
782 return Ok(());
783 }
784 let _ = self.flush();
786
787 let WriterState::Active { writer: raw, .. } =
789 std::mem::replace(&mut self.state, WriterState::Finished)
790 else {
791 self.fs.mark_writer_done();
792 return Ok(());
793 };
794
795 let bytes_written = raw.bytes_written();
796 let bw: BufWriter<ActiveHandle> = raw.into_inner();
797 let handle: ActiveHandle = bw
798 .into_inner()
799 .unwrap_or_else(|e| e.into_inner().into_parts().0);
800
801 let current_index = self.next_index - 1;
802
803 if self.has_real_events {
804 match self.fs.seal(handle, &self.active_path, current_index) {
805 Ok(seg_ref) => {
806 if M::IS_DISK {
807 self.closed_files.push_back((seg_ref, bytes_written));
808 }
809 }
810 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
811 rate_limited!(Duration::from_secs(60), {
812 tracing::warn!(
813 "active trace file {} disappeared before finalize; \
814 dropping segment",
815 self.active_path.display()
816 );
817 });
818 }
819 Err(e) => {
820 self.fs.mark_writer_done();
821 return Err(e);
822 }
823 }
824 } else {
825 tracing::debug!(
828 "removing empty final segment {}",
829 self.active_path.display()
830 );
831 if let Err(e) = self.fs.remove_active(&self.active_path)
832 && e.kind() != std::io::ErrorKind::NotFound
833 {
834 self.fs.mark_writer_done();
835 return Err(e);
836 }
837 }
838
839 if let Err(e) = self.evict_oldest() {
843 self.fs.mark_writer_done();
844 return Err(e);
845 }
846 self.fs.mark_writer_done();
847 Ok(())
848 }
849
850 crate::test_util_pub! {
853 fn write_encoded_batch(&mut self, batch: &Batch) -> std::io::Result<()> {
855 self.write_metadata_if_needed()?;
856 let WriterState::Active { writer: raw, .. } = &mut self.state else {
857 self.dropped_events += batch.event_count() as usize;
858 return Ok(());
859 };
860 if batch.event_count() > 0 {
861 raw.write_raw(batch.encoded_bytes())?;
872 self.has_real_events = true;
873 self.maybe_rotate()?;
874 }
875 Ok(())
876 }
877 }
878}
879
880impl<M: BufferMode> Drop for SegmentWriter<M> {
881 fn drop(&mut self) {
882 if self.dropped_events > 0 {
883 rate_limited!(Duration::from_secs(60), {
884 tracing::info!(
885 target: "dial9_telemetry",
886 dropped_events = self.dropped_events,
887 "SegmentWriter dropped events after finalization"
888 );
889 });
890 }
891 }
892}
893
894#[cfg(test)]
895mod tests {
896 use super::*;
897 use dial9_trace_format::TraceEvent;
898 use std::collections::HashMap;
899 use std::io::Read;
900 use tempfile::TempDir;
901
902 #[derive(TraceEvent)]
906 #[traceevent(wire_slot)]
907 struct TestEvent {
908 #[traceevent(timestamp)]
909 timestamp_ns: u64,
910 value: u64,
911 }
912
913 #[derive(Debug)]
915 enum Decoded {
916 ClockSync {
917 timestamp_ns: u64,
918 realtime_ns: u64,
919 },
920 SegmentMetadata {
921 timestamp_ns: u64,
922 entries: HashMap<String, String>,
923 },
924 Data {
925 timestamp_ns: u64,
926 },
927 }
928
929 fn decode_all(data: &[u8]) -> Vec<Decoded> {
930 use dial9_trace_format::decoder::{DecodedFrameRef, Decoder};
931 use dial9_trace_format::types::FieldValueRef;
932
933 let mut dec = Decoder::new(data).expect("valid trace header");
934 let mut out = Vec::new();
935 while let Some(frame) = dec.next_frame_ref().expect("decode frame") {
936 let DecodedFrameRef::Event {
937 type_id,
938 timestamp_ns,
939 values,
940 } = frame
941 else {
942 continue;
943 };
944 let ts = timestamp_ns;
945 let name = dec.registry().get(type_id).map(|s| s.name());
946 match name {
947 Some("ClockSyncEvent") => {
948 let realtime_ns = match values.first() {
949 Some(FieldValueRef::Varint(v)) => *v,
950 other => panic!("ClockSyncEvent realtime_ns: {other:?}"),
951 };
952 out.push(Decoded::ClockSync {
953 timestamp_ns: ts,
954 realtime_ns,
955 });
956 }
957 Some("SegmentMetadataEvent") => {
958 let entries = match values.first() {
959 Some(FieldValueRef::StringMap(m)) => m
960 .iter()
961 .map(|(k, v)| (k.to_string(), v.to_string()))
962 .collect(),
963 other => panic!("SegmentMetadataEvent entries: {other:?}"),
964 };
965 out.push(Decoded::SegmentMetadata {
966 timestamp_ns: ts,
967 entries,
968 });
969 }
970 _ => out.push(Decoded::Data { timestamp_ns: ts }),
971 }
972 }
973 out
974 }
975
976 fn test_batch() -> Batch {
979 let mut enc = Encoder::new_to(Vec::new()).unwrap();
980 enc.write(&TestEvent {
981 timestamp_ns: 1000,
982 value: 0,
983 })
984 .unwrap();
985 Batch::new(enc.into_inner(), 1)
986 }
987
988 fn rotating_file(base: &std::path::Path, i: u32) -> String {
989 format!("{}.{}.bin", base.display(), i)
990 }
991
992 fn read_trace_events(path: &str) -> Vec<Decoded> {
994 let data = std::fs::read(path).unwrap();
995 decode_all(&data)
996 .into_iter()
997 .filter(|e| matches!(e, Decoded::Data { .. }))
998 .collect()
999 }
1000
1001 fn total_disk_usage(dir: &std::path::Path) -> u64 {
1003 std::fs::read_dir(dir)
1004 .unwrap()
1005 .filter_map(|e| e.ok())
1006 .filter(|e| {
1007 let p = e.path();
1008 p.extension()
1009 .is_some_and(|ext| ext == "bin" || ext == "active")
1010 })
1011 .map(|e| e.metadata().unwrap().len())
1012 .sum()
1013 }
1014
1015 fn single_event_file_size() -> u64 {
1019 let dir = TempDir::new().unwrap();
1020 let path = dir.path().join("probe.bin");
1021 let mut w = DiskBuffer::single_file(&path).unwrap();
1022 w.write_encoded_batch(&test_batch()).unwrap();
1023 w.flush().unwrap();
1024 std::fs::metadata(w.current_active_path()).unwrap().len()
1025 }
1026
1027 #[test]
1028 fn test_writer_creation() {
1029 let dir = TempDir::new().unwrap();
1030 let path = dir.path().join("test_trace_v2.bin");
1031 let writer = DiskBuffer::single_file(&path);
1032 assert!(writer.is_ok());
1033 }
1034
1035 #[test]
1036 fn derive_max_file_size_caps_large_budgets_at_100_mib() {
1037 assert_eq!(
1038 derive_max_file_size(1024 * BYTES_PER_MIB),
1039 100 * BYTES_PER_MIB
1040 );
1041 }
1042
1043 #[test]
1044 fn derive_max_file_size_uses_quarter_of_small_budgets() {
1045 assert_eq!(derive_max_file_size(64 * BYTES_PER_MIB), 16 * BYTES_PER_MIB);
1046 }
1047
1048 #[test]
1049 fn builder_defaults_max_file_size_from_total_size() {
1050 let dir = TempDir::new().unwrap();
1051 let total = 64 * BYTES_PER_MIB;
1052 let writer = DiskBuffer::builder()
1053 .base_path(dir.path())
1054 .max_total_size(total)
1055 .build()
1056 .expect("builder should succeed without max_file_size");
1057 assert_eq!(writer.max_file_size, derive_max_file_size(total));
1058 }
1059
1060 #[test]
1061 fn builder_honors_explicit_max_file_size() {
1062 let dir = TempDir::new().unwrap();
1063 let writer = DiskBuffer::builder()
1064 .base_path(dir.path())
1065 .max_file_size(7 * BYTES_PER_MIB)
1066 .max_total_size(64 * BYTES_PER_MIB)
1067 .build()
1068 .expect("builder should succeed");
1069 assert_eq!(writer.max_file_size, 7 * BYTES_PER_MIB);
1070 }
1071
1072 #[test]
1073 fn test_write_event() {
1074 let dir = TempDir::new().unwrap();
1075 let path = dir.path().join("test_event_v2.bin");
1076 let mut writer = DiskBuffer::single_file(&path).unwrap();
1077
1078 writer.write_encoded_batch(&test_batch()).unwrap();
1079 writer.flush().unwrap();
1080
1081 let metadata = std::fs::metadata(writer.current_active_path()).unwrap();
1082 assert!(
1083 metadata.len() > 0,
1084 "file should not be empty after writing an event"
1085 );
1086 }
1087
1088 #[test]
1089 fn test_write_batch_sizes() {
1090 let dir = TempDir::new().unwrap();
1091 let path = dir.path().join("test_batch_v2.bin");
1092 let mut writer = DiskBuffer::single_file(&path).unwrap();
1093
1094 let one_event_size = single_event_file_size();
1095
1096 for _ in 0..2 {
1097 writer.write_encoded_batch(&test_batch()).unwrap();
1098 }
1099 writer.flush().unwrap();
1100
1101 let metadata = std::fs::metadata(writer.current_active_path()).unwrap();
1102 assert!(metadata.len() > one_event_size);
1104 }
1105
1106 #[test]
1107 fn test_binary_format_header() {
1108 let dir = TempDir::new().unwrap();
1109 let path = dir.path().join("test_format_v2.bin");
1110 let writer = DiskBuffer::single_file(&path).unwrap();
1111 let active = writer.current_active_path().to_owned();
1112 drop(writer);
1113
1114 let mut file = std::fs::File::open(&active).unwrap();
1115 let mut magic = [0u8; 4];
1116 file.read_exact(&mut magic).unwrap();
1117 assert_eq!(&magic, b"TRC\0");
1118 }
1119
1120 #[test]
1121 fn test_rotating_writer_creation() {
1122 let dir = TempDir::new().unwrap();
1123 let mut writer = DiskBuffer::builder()
1124 .base_path(dir.path())
1125 .max_file_size(1024)
1126 .max_total_size(4096)
1127 .build()
1128 .unwrap();
1129 writer.finalize().unwrap();
1130
1131 assert!(
1133 !dir.path().join("trace.0.bin").exists(),
1134 "empty segment should not be sealed"
1135 );
1136 assert!(
1137 !dir.path().join("trace.0.bin.active").exists(),
1138 "active file should be removed"
1139 );
1140 }
1141
1142 #[test]
1143 fn test_rotating_writer_rotation() {
1144 let dir = TempDir::new().unwrap();
1145 let base = dir.path().join("trace");
1146 let one_event = single_event_file_size();
1148 let mut writer = DiskBuffer::builder()
1149 .base_path(dir.path())
1150 .max_file_size(one_event)
1151 .max_total_size(100_000)
1152 .build()
1153 .unwrap();
1154
1155 for _ in 0..3 {
1156 writer.write_encoded_batch(&test_batch()).unwrap();
1157 }
1158 writer.finalize().unwrap();
1159
1160 let total: usize = (0..10)
1162 .map(|i| {
1163 let f = rotating_file(&base, i);
1164 if std::path::Path::new(&f).exists() {
1165 read_trace_events(&f).len()
1166 } else {
1167 0
1168 }
1169 })
1170 .sum();
1171 assert_eq!(total, 3);
1172 }
1173
1174 #[test]
1175 fn test_rotating_writer_eviction() {
1176 let dir = TempDir::new().unwrap();
1177 let base = dir.path().join("trace");
1178 let one_event = single_event_file_size();
1179 let max_file_size = one_event;
1180 let max_total_size = max_file_size * 3;
1181 let mut writer = DiskBuffer::builder()
1182 .base_path(dir.path())
1183 .max_file_size(max_file_size)
1184 .max_total_size(max_total_size)
1185 .build()
1186 .unwrap();
1187
1188 for _ in 0..10 {
1189 writer.write_encoded_batch(&test_batch()).unwrap();
1190 }
1191 writer.finalize().unwrap();
1192
1193 assert!(total_disk_usage(dir.path()) <= max_total_size);
1195
1196 assert!(!std::path::Path::new(&rotating_file(&base, 0)).exists());
1198 }
1199
1200 #[test]
1201 fn test_rotating_writer_stops_when_over_budget() {
1202 let dir = TempDir::new().unwrap();
1203 let base = dir.path().join("trace");
1204 let one_event = single_event_file_size();
1205 let max_file_size = one_event;
1207 let max_total_size = one_event + 5;
1208 let mut writer = DiskBuffer::builder()
1209 .base_path(dir.path())
1210 .max_file_size(max_file_size)
1211 .max_total_size(max_total_size)
1212 .build()
1213 .unwrap();
1214
1215 for _ in 0..100 {
1216 writer.write_encoded_batch(&test_batch()).unwrap();
1217 }
1218 writer.finalize().unwrap();
1219
1220 let total: usize = (0..100)
1222 .map(|i| {
1223 let f = rotating_file(&base, i);
1224 if std::path::Path::new(&f).exists() {
1225 read_trace_events(&f).len()
1226 } else {
1227 0
1228 }
1229 })
1230 .sum();
1231 assert!(
1232 total < 100,
1233 "should have stopped writing, got {total} events"
1234 );
1235 }
1236
1237 #[test]
1246 fn test_writer_stops_on_tiny_overshoot_after_eviction() {
1247 let dir = TempDir::new().unwrap();
1248 let max_file_size = 200;
1253 let num_files = 100u64;
1254 let max_total_size = max_file_size * num_files;
1255 let mut writer = DiskBuffer::builder()
1256 .base_path(dir.path())
1257 .max_file_size(max_file_size)
1258 .max_total_size(max_total_size)
1259 .build()
1260 .unwrap();
1261
1262 for i in 0..5000 {
1266 writer.write_encoded_batch(&test_batch()).unwrap();
1267 if matches!(writer.state, WriterState::Finished) {
1268 panic!(
1269 "Writer stopped at batch {i}! total_size={}, max_total_size={}, \
1270 closed_files={}. \
1271 write_encoded_batch should try eviction before stopping.",
1272 writer.total_size(),
1273 max_total_size,
1274 writer.closed_files.len()
1275 );
1276 }
1277 }
1278 }
1279
1280 #[test]
1281 fn test_rotating_writer_file_naming() {
1282 let dir = TempDir::new().unwrap();
1283 let base = dir.path().join("trace");
1284 let one_event = single_event_file_size();
1285 let mut writer = DiskBuffer::builder()
1286 .base_path(dir.path())
1287 .max_file_size(one_event)
1288 .max_total_size(100_000)
1289 .build()
1290 .unwrap();
1291
1292 for _ in 0..5 {
1293 writer.write_encoded_batch(&test_batch()).unwrap();
1294 }
1295 writer.finalize().unwrap();
1296
1297 assert!(
1299 std::path::Path::new(&rotating_file(&base, 0)).exists(),
1300 "File 0 should exist"
1301 );
1302 let total: usize = (0..10)
1304 .map(|i| {
1305 let f = rotating_file(&base, i);
1306 if std::path::Path::new(&f).exists() {
1307 read_trace_events(&f).len()
1308 } else {
1309 0
1310 }
1311 })
1312 .sum();
1313 assert_eq!(total, 5);
1314 }
1315
1316 #[test]
1317 fn test_write_batch_across_rotation_boundary() {
1318 let dir = TempDir::new().unwrap();
1319 let base = dir.path().join("trace");
1320 let one_event = single_event_file_size();
1321 let mut writer = DiskBuffer::builder()
1322 .base_path(dir.path())
1323 .max_file_size(one_event)
1324 .max_total_size(100_000)
1325 .build()
1326 .unwrap();
1327
1328 for _ in 0..3 {
1329 writer.write_encoded_batch(&test_batch()).unwrap();
1330 }
1331 writer.finalize().unwrap();
1332
1333 let total: usize = (0..10)
1335 .map(|i| {
1336 let f = rotating_file(&base, i);
1337 if std::path::Path::new(&f).exists() {
1338 read_trace_events(&f).len()
1339 } else {
1340 0
1341 }
1342 })
1343 .sum();
1344 assert_eq!(total, 3);
1345 }
1346
1347 #[test]
1348 fn test_rotated_files_have_valid_headers() {
1349 let dir = TempDir::new().unwrap();
1350 let base = dir.path().join("trace");
1351 let one_event = single_event_file_size();
1352 let mut writer = DiskBuffer::builder()
1353 .base_path(dir.path())
1354 .max_file_size(one_event)
1355 .max_total_size(100_000)
1356 .build()
1357 .unwrap();
1358
1359 for _ in 0..3 {
1360 writer.write_encoded_batch(&test_batch()).unwrap();
1361 }
1362 writer.finalize().unwrap();
1363
1364 let total: usize = (0..10)
1366 .map(|i| {
1367 let f = rotating_file(&base, i);
1368 if std::path::Path::new(&f).exists() {
1369 read_trace_events(&f).len() } else {
1371 0
1372 }
1373 })
1374 .sum();
1375 assert_eq!(total, 3);
1376 }
1377
1378 #[test]
1379 fn test_flush_after_stop() {
1380 let dir = TempDir::new().unwrap();
1381 let mut writer = DiskBuffer::builder()
1383 .base_path(dir.path())
1384 .max_file_size(10_000)
1385 .max_total_size(50)
1386 .build()
1387 .unwrap();
1388
1389 for _ in 0..5 {
1390 writer.write_encoded_batch(&test_batch()).unwrap();
1391 }
1392 assert!(writer.flush().is_ok());
1394 assert!(writer.flush().is_ok());
1395 }
1396
1397 #[test]
1398 fn test_mixed_event_sizes() {
1399 let dir = TempDir::new().unwrap();
1400 let base = dir.path().join("trace");
1401 let one_event = single_event_file_size();
1402 let mut writer = DiskBuffer::builder()
1403 .base_path(dir.path())
1404 .max_file_size(one_event)
1405 .max_total_size(100_000)
1406 .build()
1407 .unwrap();
1408
1409 for _ in 0..3 {
1410 writer.write_encoded_batch(&test_batch()).unwrap();
1411 }
1412 writer.finalize().unwrap();
1413
1414 let mut total = 0;
1416 for i in 0..10 {
1417 let f = rotating_file(&base, i);
1418 if std::path::Path::new(&f).exists() {
1419 total += read_trace_events(&f).len();
1420 }
1421 }
1422 assert_eq!(total, 3);
1423 }
1424
1425 #[test]
1426 fn test_event_exactly_on_max_file_size_boundary() {
1427 let dir = TempDir::new().unwrap();
1428 let base = dir.path().join("trace");
1429 let one_event = single_event_file_size();
1430 let mut writer = DiskBuffer::builder()
1432 .base_path(dir.path())
1433 .max_file_size(one_event)
1434 .max_total_size(100_000)
1435 .build()
1436 .unwrap();
1437
1438 for _ in 0..2 {
1439 writer.write_encoded_batch(&test_batch()).unwrap();
1440 }
1441 writer.finalize().unwrap();
1442
1443 let total: usize = (0..10)
1445 .map(|i| {
1446 let f = rotating_file(&base, i);
1447 if std::path::Path::new(&f).exists() {
1448 read_trace_events(&f).len()
1449 } else {
1450 0
1451 }
1452 })
1453 .sum();
1454 assert_eq!(total, 2);
1455 }
1456
1457 #[test]
1458 fn test_active_suffix_while_writing() {
1459 let dir = TempDir::new().unwrap();
1460 let mut writer = DiskBuffer::builder()
1461 .base_path(dir.path())
1462 .max_file_size(1024)
1463 .max_total_size(100000)
1464 .build()
1465 .unwrap();
1466 writer.write_encoded_batch(&test_batch()).unwrap();
1467 writer.flush().unwrap();
1468
1469 let active = dir.path().join("trace.0.bin.active");
1471 assert!(active.exists(), "active file should exist while writing");
1472 let sealed = dir.path().join("trace.0.bin");
1473 assert!(!sealed.exists(), "sealed file should not exist yet");
1474 }
1475
1476 #[test]
1477 fn test_rotation_seals_previous_file() {
1478 let dir = TempDir::new().unwrap();
1479 let one_event = single_event_file_size();
1480 let mut writer = DiskBuffer::builder()
1481 .base_path(dir.path())
1482 .max_file_size(one_event)
1483 .max_total_size(100_000)
1484 .build()
1485 .unwrap();
1486
1487 writer.write_encoded_batch(&test_batch()).unwrap();
1489 writer.write_encoded_batch(&test_batch()).unwrap();
1490 writer.flush().unwrap();
1491
1492 assert!(
1494 dir.path().join("trace.0.bin").exists(),
1495 "rotated file should be sealed"
1496 );
1497 assert!(
1498 !dir.path().join("trace.0.bin.active").exists(),
1499 "rotated file should not be active"
1500 );
1501 assert!(
1502 dir.path().join("trace.1.bin.active").exists(),
1503 "current file should be active"
1504 );
1505 assert!(
1506 !dir.path().join("trace.1.bin").exists(),
1507 "current file should not be sealed"
1508 );
1509 }
1510
1511 #[test]
1512 fn test_finalize_renames_current_file() {
1513 let dir = TempDir::new().unwrap();
1514 let mut writer = DiskBuffer::builder()
1515 .base_path(dir.path())
1516 .max_file_size(1024)
1517 .max_total_size(100000)
1518 .build()
1519 .unwrap();
1520 writer.write_encoded_batch(&test_batch()).unwrap();
1521 writer.finalize().unwrap();
1522
1523 assert!(
1524 dir.path().join("trace.0.bin").exists(),
1525 "file should be sealed after finalize()"
1526 );
1527 assert!(
1528 !dir.path().join("trace.0.bin.active").exists(),
1529 "active file should be gone after finalize()"
1530 );
1531 }
1532
1533 #[test]
1534 fn test_finalize_removes_empty_segment_after_rotation() {
1535 let dir = TempDir::new().unwrap();
1536 let mut writer = DiskBuffer::builder()
1538 .base_path(dir.path())
1539 .max_file_size(1)
1540 .max_total_size(100_000)
1541 .build()
1542 .unwrap();
1543 writer.write_encoded_batch(&test_batch()).unwrap();
1545 assert!(dir.path().join("trace.0.bin").exists());
1547 assert!(dir.path().join("trace.1.bin.active").exists());
1548
1549 writer.finalize().unwrap();
1551 assert!(
1552 !dir.path().join("trace.1.bin").exists(),
1553 "empty segment should not be sealed"
1554 );
1555 assert!(
1556 !dir.path().join("trace.1.bin.active").exists(),
1557 "empty active file should be removed"
1558 );
1559 assert!(dir.path().join("trace.0.bin").exists());
1561 }
1562
1563 #[test]
1564 fn test_single_file_no_active_suffix() {
1565 let dir = TempDir::new().unwrap();
1566 let path = dir.path().join("test.bin");
1567 let mut writer = DiskBuffer::single_file(&path).unwrap();
1568 writer.write_encoded_batch(&test_batch()).unwrap();
1569 writer.flush().unwrap();
1570 writer.finalize().unwrap();
1571
1572 assert!(dir.path().join("test.0.bin").exists());
1574 assert!(!dir.path().join("test.0.bin.active").exists());
1575 }
1576
1577 #[test]
1578 #[cfg(feature = "pipeline")]
1579 fn test_single_file_sealed_segment_discoverable_by_worker() {
1580 use crate::sealed::find_sealed_segments;
1581
1582 let dir = TempDir::new().unwrap();
1583 let path = dir.path().join("trace.bin");
1584 let mut writer = DiskBuffer::single_file(&path).unwrap();
1585 writer.write_encoded_batch(&test_batch()).unwrap();
1586 writer.flush().unwrap();
1587 writer.finalize().unwrap();
1588
1589 let segments = find_sealed_segments(dir.path(), "trace").unwrap();
1590 assert_eq!(
1591 segments.len(),
1592 1,
1593 "worker should find exactly one sealed segment"
1594 );
1595 assert_eq!(segments[0].path, dir.path().join("trace.0.bin"));
1596 }
1597
1598 #[test]
1599 fn test_segment_metadata_roundtrip() {
1600 let dir = TempDir::new().unwrap();
1601 let base = dir.path().join("trace");
1602 let mut writer = DiskBuffer::builder()
1603 .base_path(dir.path())
1604 .max_file_size(100_000)
1605 .max_total_size(100_000)
1606 .segment_metadata(vec![
1607 ("service".into(), "checkout-api".into()),
1608 ("host".into(), "i-0abc123".into()),
1609 ])
1610 .build()
1611 .unwrap();
1612 writer.write_encoded_batch(&test_batch()).unwrap();
1613 writer.flush().unwrap();
1614 writer.finalize().unwrap();
1615
1616 let all_events = decode_all(&std::fs::read(format!("{}.0.bin", base.display())).unwrap());
1617 let metadata: Vec<_> = all_events
1618 .iter()
1619 .filter_map(|e| match e {
1620 Decoded::SegmentMetadata { entries, .. } => Some(entries.clone()),
1621 _ => None,
1622 })
1623 .collect();
1624 assert_eq!(metadata.len(), 1);
1625 assert!(
1626 metadata[0].get("service").map(String::as_str) == Some("checkout-api"),
1627 "missing service entry: {:?}",
1628 metadata[0]
1629 );
1630 assert!(
1631 metadata[0].get("host").map(String::as_str) == Some("i-0abc123"),
1632 "missing host entry: {:?}",
1633 metadata[0]
1634 );
1635 assert_eq!(
1636 metadata[0].get(DIAL9_VERSION_KEY).map(String::as_str),
1637 Some(DIAL9_VERSION_VALUE),
1638 "missing built-in dial9.dial9-tokio-telemetry.version: {:?}",
1639 metadata[0]
1640 );
1641 }
1642
1643 #[test]
1644 fn test_segment_metadata_written_in_every_rotated_file() {
1645 let dir = TempDir::new().unwrap();
1646 let one_event = single_event_file_size();
1647 let mut writer = DiskBuffer::builder()
1648 .base_path(dir.path())
1649 .max_file_size(one_event)
1650 .max_total_size(100_000)
1651 .segment_metadata(vec![("k".into(), "v".into())])
1652 .build()
1653 .unwrap();
1654
1655 for _ in 0..5 {
1656 writer.write_encoded_batch(&test_batch()).unwrap();
1657 }
1658 writer.flush().unwrap();
1659 writer.finalize().unwrap();
1660
1661 let mut files: Vec<_> = std::fs::read_dir(dir.path())
1662 .unwrap()
1663 .filter_map(|e| e.ok())
1664 .map(|e| e.path())
1665 .filter(|p| p.extension().is_some_and(|ext| ext == "bin"))
1666 .collect();
1667 files.sort();
1668 assert!(files.len() >= 2, "expected at least 2 files from rotation");
1669
1670 for file in &files {
1671 let all_events = decode_all(&std::fs::read(file).unwrap());
1672 let has_metadata = all_events.iter().any(|e| match e {
1673 Decoded::SegmentMetadata { entries, .. } => {
1674 entries.get("k").map(String::as_str) == Some("v")
1675 }
1676 _ => false,
1677 });
1678 assert!(has_metadata, "{}: expected SegmentMetadata", file.display());
1679 }
1680 }
1681
1682 #[test]
1683 fn test_dynamic_metadata_merged_on_rotation() {
1684 let dir = TempDir::new().unwrap();
1685 let one_event = single_event_file_size();
1686 let mut writer = DiskBuffer::builder()
1687 .base_path(dir.path())
1688 .max_file_size(one_event)
1689 .max_total_size(100_000)
1690 .segment_metadata(vec![("service".into(), "myapp".into())])
1691 .build()
1692 .unwrap();
1693
1694 let mut merged = writer.segment_metadata().to_vec();
1696 merged.push(("runtime.main".into(), "0,1,2,3".into()));
1697 writer.update_segment_metadata(merged);
1698
1699 for _ in 0..4 {
1702 writer.write_encoded_batch(&test_batch()).unwrap();
1703 }
1704 writer.flush().unwrap();
1705 writer.finalize().unwrap();
1706
1707 let mut files: Vec<_> = std::fs::read_dir(dir.path())
1708 .unwrap()
1709 .filter_map(|e| e.ok())
1710 .map(|e| e.path())
1711 .filter(|p| p.extension().is_some_and(|ext| ext == "bin"))
1712 .collect();
1713 files.sort();
1714 assert!(files.len() >= 2, "expected at least 2 files from rotation");
1715
1716 for file in &files[1..] {
1719 let all_events = decode_all(&std::fs::read(file).unwrap());
1720 let meta: Vec<_> = all_events
1721 .iter()
1722 .filter_map(|e| match e {
1723 Decoded::SegmentMetadata { entries, .. } => Some(entries.clone()),
1724 _ => None,
1725 })
1726 .collect();
1727 assert_eq!(
1728 meta.len(),
1729 1,
1730 "{}: expected 1 metadata event",
1731 file.display()
1732 );
1733 assert!(
1734 meta[0].get("service").map(String::as_str) == Some("myapp"),
1735 "{}: missing static metadata",
1736 file.display()
1737 );
1738 assert!(
1739 meta[0].get("runtime.main").map(String::as_str) == Some("0,1,2,3"),
1740 "{}: missing dynamic runtime worker metadata",
1741 file.display()
1742 );
1743 }
1744 }
1745
1746 #[test]
1747 fn test_segment_metadata_empty_entries() {
1748 let dir = TempDir::new().unwrap();
1749 let path = dir.path().join("trace.bin");
1750 let mut writer = DiskBuffer::single_file(&path).unwrap();
1751 writer.write_encoded_batch(&test_batch()).unwrap();
1752 writer.flush().unwrap();
1753
1754 let all_events = decode_all(&std::fs::read(writer.current_active_path()).unwrap());
1755 let data_count = all_events
1756 .iter()
1757 .filter(|e| matches!(e, Decoded::Data { .. }))
1758 .count();
1759 assert_eq!(data_count, 1);
1760 let metadata: Vec<_> = all_events
1763 .iter()
1764 .filter_map(|e| match e {
1765 Decoded::SegmentMetadata { entries, .. } => Some(entries),
1766 _ => None,
1767 })
1768 .collect();
1769 assert_eq!(metadata.len(), 1);
1770 assert_eq!(
1771 metadata[0].get(DIAL9_VERSION_KEY).map(String::as_str),
1772 Some(DIAL9_VERSION_VALUE)
1773 );
1774 }
1775
1776 #[test]
1779 fn test_eviction_removes_gz_variant() {
1780 let dir = TempDir::new().unwrap();
1781 let one_event = single_event_file_size();
1782 let max_file_size = one_event;
1783 let max_total_size = max_file_size * 100;
1785 let mut writer = DiskBuffer::builder()
1786 .base_path(dir.path())
1787 .max_file_size(max_file_size)
1788 .max_total_size(max_total_size)
1789 .build()
1790 .unwrap();
1791
1792 writer.write_encoded_batch(&test_batch()).unwrap();
1795 writer.write_encoded_batch(&test_batch()).unwrap();
1796 let seg0 = dir.path().join("trace.0.bin");
1800 let seg0_gz = dir.path().join("trace.0.bin.gz");
1801 assert!(seg0.exists(), "trace.0.bin should exist after rotation");
1802 std::fs::rename(&seg0, &seg0_gz).unwrap();
1803
1804 writer.max_total_size = max_file_size;
1807 for _ in 0..3 {
1808 writer.write_encoded_batch(&test_batch()).unwrap();
1809 }
1810 writer.finalize().unwrap();
1811
1812 assert!(!seg0_gz.exists(), "trace.0.bin.gz should have been evicted");
1814 }
1815
1816 #[test]
1824 fn test_eviction_keeps_most_recent_segment_when_over_budget() {
1825 let dir = TempDir::new().unwrap();
1826 let one_event = single_event_file_size();
1827 let max_file_size = u64::MAX;
1830 let max_total_size = one_event / 2;
1831 assert!(
1832 max_total_size < one_event,
1833 "test setup: budget must be smaller than a single segment"
1834 );
1835 let mut writer = DiskBuffer::builder()
1836 .base_path(dir.path())
1837 .max_file_size(max_file_size)
1838 .max_total_size(max_total_size)
1839 .build()
1840 .unwrap();
1841
1842 writer.write_encoded_batch(&test_batch()).unwrap();
1843 assert!(
1845 writer.total_size() > max_total_size,
1846 "single segment ({}) should exceed budget ({max_total_size})",
1847 writer.total_size()
1848 );
1849
1850 writer.evict_oldest().unwrap();
1853
1854 assert!(
1855 matches!(writer.state, WriterState::Finished),
1856 "writer should stop once even the most-recent segment exceeds budget"
1857 );
1858 assert!(
1861 std::path::Path::new(&writer.current_active_path()).exists(),
1862 "the most-recent segment must not be evicted"
1863 );
1864 assert!(
1865 total_disk_usage(dir.path()) > max_total_size,
1866 "retained segment is expected to push on-disk usage over the budget"
1867 );
1868 }
1869
1870 #[tokio::test(start_paused = true)]
1873 async fn test_time_rotation_triggers_on_expired_boundary() {
1874 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
1875 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
1876
1877 let dir = TempDir::new().unwrap();
1878 let base = dir.path().join("trace");
1879 let mut writer = DiskBuffer::builder()
1880 .base_path(dir.path())
1881 .max_file_size(u64::MAX)
1882 .max_total_size(100_000)
1883 .rotation_period(Duration::from_secs(60))
1884 .build()
1885 .unwrap();
1886
1887 writer.write_encoded_batch(&test_batch()).unwrap();
1888 writer.flush().unwrap();
1889 let initial_index = writer.next_index;
1890
1891 tokio::time::advance(Duration::from_secs(61)).await;
1893
1894 writer.write_encoded_batch(&test_batch()).unwrap();
1896 writer.flush().unwrap();
1897 writer.drained().unwrap();
1898
1899 assert!(
1900 writer.next_index > initial_index,
1901 "expected time-based rotation to trigger"
1902 );
1903 writer.finalize().unwrap();
1904
1905 let total: usize = (0..10)
1906 .map(|i| {
1907 let f = rotating_file(&base, i);
1908 if std::path::Path::new(&f).exists() {
1909 read_trace_events(&f).len()
1910 } else {
1911 0
1912 }
1913 })
1914 .sum();
1915 assert_eq!(total, 2);
1916 }
1917
1918 #[tokio::test(start_paused = true)]
1923 async fn test_first_rotation_uses_monotonic_period_not_wallclock_alignment() {
1924 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
1925 let start_wall = std::time::UNIX_EPOCH + Duration::from_secs(22);
1926 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(start_wall));
1927
1928 let dir = TempDir::new().unwrap();
1929 let mut writer = DiskBuffer::builder()
1930 .base_path(dir.path())
1931 .max_file_size(u64::MAX)
1932 .max_total_size(100_000)
1933 .rotation_period(Duration::from_secs(60))
1934 .build()
1935 .unwrap();
1936
1937 writer.write_encoded_batch(&test_batch()).unwrap();
1938 writer.flush().unwrap();
1939 let initial_index = writer.next_index;
1940
1941 tokio::time::advance(Duration::from_secs(50)).await;
1945
1946 writer.write_encoded_batch(&test_batch()).unwrap();
1947 writer.flush().unwrap();
1948 writer.drained().unwrap();
1949
1950 assert_eq!(
1951 writer.next_index, initial_index,
1952 "rotation must not fire before one full rotation_period of monotonic time has elapsed",
1953 );
1954
1955 tokio::time::advance(Duration::from_secs(11)).await;
1957 writer.write_encoded_batch(&test_batch()).unwrap();
1958 writer.flush().unwrap();
1959 writer.drained().unwrap();
1960 assert!(
1961 writer.next_index > initial_index,
1962 "rotation should fire once a full rotation_period of monotonic time has elapsed",
1963 );
1964
1965 writer.finalize().unwrap();
1966 }
1967
1968 #[tokio::test(start_paused = true)]
1969 async fn test_time_rotation_skips_when_no_real_events() {
1970 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
1971 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
1972
1973 let dir = TempDir::new().unwrap();
1974 let mut writer = DiskBuffer::builder()
1975 .base_path(dir.path())
1976 .max_file_size(u64::MAX)
1977 .max_total_size(100_000)
1978 .rotation_period(Duration::from_secs(60))
1979 .build()
1980 .unwrap();
1981
1982 tokio::time::advance(Duration::from_secs(120)).await;
1984
1985 let empty_batch = Batch::new(vec![], 0);
1986 writer.write_encoded_batch(&empty_batch).unwrap();
1987
1988 assert_eq!(
1989 writer.next_index, 1,
1990 "should not rotate when no real events exist"
1991 );
1992 writer.finalize().unwrap();
1993 }
1994
1995 #[test]
1996 fn test_size_rotation_still_works_with_time_disabled() {
1997 let dir = TempDir::new().unwrap();
1998 let base = dir.path().join("trace");
1999 let one_event = single_event_file_size();
2000 let mut writer = DiskBuffer::builder()
2001 .base_path(dir.path())
2002 .max_file_size(one_event)
2003 .max_total_size(100_000)
2004 .rotation_period(std::time::Duration::MAX)
2005 .build()
2006 .unwrap();
2007
2008 for _ in 0..3 {
2009 writer.write_encoded_batch(&test_batch()).unwrap();
2010 }
2011 writer.finalize().unwrap();
2012
2013 let total: usize = (0..10)
2014 .map(|i| {
2015 let f = rotating_file(&base, i);
2016 if std::path::Path::new(&f).exists() {
2017 read_trace_events(&f).len()
2018 } else {
2019 0
2020 }
2021 })
2022 .sum();
2023 assert_eq!(total, 3);
2024 }
2025
2026 #[tokio::test(start_paused = true)]
2027 async fn test_time_rotation_respects_eviction_budget() {
2028 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
2029 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
2030
2031 let dir = TempDir::new().unwrap();
2032 let one_event = single_event_file_size();
2033 let mut writer = DiskBuffer::builder()
2034 .base_path(dir.path())
2035 .max_file_size(u64::MAX)
2036 .max_total_size(one_event * 3)
2037 .rotation_period(Duration::from_secs(60))
2038 .build()
2039 .unwrap();
2040
2041 writer.write_encoded_batch(&test_batch()).unwrap();
2042 for _ in 0..5 {
2043 tokio::time::advance(Duration::from_secs(61)).await;
2044 writer.write_encoded_batch(&test_batch()).unwrap();
2045 writer.drained().unwrap();
2046 }
2047 writer.finalize().unwrap();
2048
2049 assert!(
2050 total_disk_usage(dir.path()) <= one_event * 3,
2051 "disk usage should stay within budget"
2052 );
2053 }
2054
2055 #[test]
2056 fn test_builder_rotation_period_default() {
2057 let dir = TempDir::new().unwrap();
2058 let writer = DiskBuffer::builder()
2059 .base_path(dir.path())
2060 .max_file_size(1024)
2061 .max_total_size(100_000)
2062 .build()
2063 .unwrap();
2064 assert_eq!(writer.rotation_period, DEFAULT_ROTATION_PERIOD);
2065 }
2066
2067 #[test]
2068 fn test_new_uses_default_rotation_period() {
2069 let dir = TempDir::new().unwrap();
2070 let writer = DiskBuffer::builder()
2071 .base_path(dir.path())
2072 .max_file_size(1024)
2073 .max_total_size(100_000)
2074 .build()
2075 .unwrap();
2076 assert_eq!(writer.rotation_period, DEFAULT_ROTATION_PERIOD);
2077 }
2078
2079 #[tokio::test(start_paused = true)]
2080 async fn test_finalize_after_time_rotation() {
2081 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
2082 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
2083
2084 let dir = TempDir::new().unwrap();
2085 let base = dir.path().join("trace");
2086 let mut writer = DiskBuffer::builder()
2087 .base_path(dir.path())
2088 .max_file_size(u64::MAX)
2089 .max_total_size(100_000)
2090 .rotation_period(Duration::from_secs(60))
2091 .build()
2092 .unwrap();
2093
2094 writer.write_encoded_batch(&test_batch()).unwrap();
2095 tokio::time::advance(Duration::from_secs(61)).await;
2096 writer.write_encoded_batch(&test_batch()).unwrap();
2097 writer.drained().unwrap();
2098 writer.finalize().unwrap();
2099
2100 let total: usize = (0..10)
2101 .map(|i| {
2102 let f = rotating_file(&base, i);
2103 if std::path::Path::new(&f).exists() {
2104 read_trace_events(&f).len()
2105 } else {
2106 0
2107 }
2108 })
2109 .sum();
2110 assert_eq!(total, 2);
2111 }
2112
2113 #[tokio::test(start_paused = true)]
2114 async fn test_stale_boundary_does_not_rotate_first_event() {
2115 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
2116 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
2117
2118 let dir = TempDir::new().unwrap();
2119 let base = dir.path().join("trace");
2120 let mut writer = DiskBuffer::builder()
2121 .base_path(dir.path())
2122 .max_file_size(u64::MAX)
2123 .max_total_size(100_000)
2124 .rotation_period(Duration::from_secs(60))
2125 .build()
2126 .unwrap();
2127
2128 tokio::time::advance(Duration::from_secs(300)).await;
2130
2131 writer.write_encoded_batch(&test_batch()).unwrap();
2133 assert_eq!(
2134 writer.next_index, 1,
2135 "first event after idle gap should not trigger immediate rotation"
2136 );
2137
2138 writer.write_encoded_batch(&test_batch()).unwrap();
2140 assert_eq!(
2141 writer.next_index, 1,
2142 "second event should still be in the same segment"
2143 );
2144
2145 writer.finalize().unwrap();
2146
2147 let events = read_trace_events(&rotating_file(&base, 0));
2148 assert_eq!(events.len(), 2, "both events should be in segment 0");
2149 }
2150
2151 #[test]
2152 fn test_clock_sync_precedes_first_data_event() {
2153 use crate::sealed::LEGACY_EPOCH_NS_FLOOR;
2154
2155 let dir = TempDir::new().unwrap();
2156 let base = dir.path().join("trace");
2157 let mut writer = DiskBuffer::builder()
2158 .base_path(dir.path())
2159 .max_file_size(100_000)
2160 .max_total_size(100_000)
2161 .build()
2162 .unwrap();
2163 writer.write_encoded_batch(&test_batch()).unwrap();
2164 writer.flush().unwrap();
2165 writer.finalize().unwrap();
2166
2167 let data = std::fs::read(rotating_file(&base, 0)).unwrap();
2168 let all = decode_all(&data);
2169
2170 let first_data_idx = all
2173 .iter()
2174 .position(|e| matches!(e, Decoded::Data { .. }))
2175 .expect("expected at least one data event");
2176 let first_clock_sync_idx = all
2177 .iter()
2178 .position(|e| matches!(e, Decoded::ClockSync { .. }))
2179 .expect("expected a ClockSyncEvent in the file");
2180 assert!(first_clock_sync_idx < first_data_idx);
2181
2182 match &all[first_clock_sync_idx] {
2183 Decoded::ClockSync { realtime_ns, .. } => {
2184 assert!(*realtime_ns >= LEGACY_EPOCH_NS_FLOOR);
2185 }
2186 _ => unreachable!(),
2187 }
2188 }
2189
2190 #[test]
2191 fn test_segment_metadata_timestamp_is_monotonic_scale() {
2192 use crate::sealed::LEGACY_EPOCH_NS_FLOOR;
2193
2194 let dir = TempDir::new().unwrap();
2195 let base = dir.path().join("trace");
2196 let mut writer = DiskBuffer::builder()
2197 .base_path(dir.path())
2198 .max_file_size(100_000)
2199 .max_total_size(100_000)
2200 .build()
2201 .unwrap();
2202 writer.write_encoded_batch(&test_batch()).unwrap();
2203 writer.flush().unwrap();
2204 writer.finalize().unwrap();
2205
2206 let data = std::fs::read(rotating_file(&base, 0)).unwrap();
2207 let all = decode_all(&data);
2208
2209 let seg_ts = all
2212 .iter()
2213 .find_map(|e| match e {
2214 Decoded::SegmentMetadata { timestamp_ns, .. } => Some(*timestamp_ns),
2215 _ => None,
2216 })
2217 .expect("SegmentMetadata");
2218 assert!(
2219 seg_ts < LEGACY_EPOCH_NS_FLOOR,
2220 "SegmentMetadata.timestamp_nanos ({seg_ts}) should be monotonic-scale"
2221 );
2222 }
2223
2224 #[test]
2225 fn test_clock_sync_written_in_every_rotated_file() {
2226 let dir = TempDir::new().unwrap();
2227 let one_event = single_event_file_size();
2228 let mut writer = DiskBuffer::builder()
2229 .base_path(dir.path())
2230 .max_file_size(one_event)
2231 .max_total_size(100_000)
2232 .build()
2233 .unwrap();
2234
2235 for _ in 0..5 {
2236 writer.write_encoded_batch(&test_batch()).unwrap();
2237 }
2238 writer.flush().unwrap();
2239 writer.finalize().unwrap();
2240
2241 let mut files: Vec<_> = std::fs::read_dir(dir.path())
2242 .unwrap()
2243 .filter_map(|e| e.ok())
2244 .map(|e| e.path())
2245 .filter(|p| p.extension().is_some_and(|ext| ext == "bin"))
2246 .collect();
2247 files.sort();
2248 assert!(files.len() >= 2, "expected at least 2 files from rotation");
2249
2250 for file in &files {
2251 let all = decode_all(&std::fs::read(file).unwrap());
2252 let has_clock_sync = all.iter().any(|e| matches!(e, Decoded::ClockSync { .. }));
2253 assert!(
2254 has_clock_sync,
2255 "{}: expected ClockSyncEvent",
2256 file.display()
2257 );
2258 }
2259 }
2260
2261 #[test]
2264 fn test_legacy_trace_without_clock_sync_still_decodes() {
2265 let mut enc = Encoder::new_to(Vec::new()).unwrap();
2266 enc.write(&SegmentMetadataEvent {
2267 timestamp_ns: 1,
2268 entries: vec![("k".into(), "v".into())],
2269 })
2270 .unwrap();
2271 enc.write(&TestEvent {
2272 timestamp_ns: 1000,
2273 value: 0,
2274 })
2275 .unwrap();
2276 let buf = enc.into_inner();
2277
2278 let all = decode_all(&buf);
2279 assert!(
2280 all.iter().any(|e| matches!(e, Decoded::Data { .. })),
2281 "expected data event to decode"
2282 );
2283 assert!(
2284 !all.iter().any(|e| matches!(e, Decoded::ClockSync { .. })),
2285 "legacy trace must not contain ClockSync"
2286 );
2287 }
2288
2289 #[test]
2290 fn test_clock_sync_offset_recovers_wall_clock_for_recent_event() {
2291 use std::time::{SystemTime, UNIX_EPOCH};
2292
2293 let dir = TempDir::new().unwrap();
2294 let base = dir.path().join("trace");
2295 let mut writer = DiskBuffer::builder()
2296 .base_path(dir.path())
2297 .max_file_size(100_000)
2298 .max_total_size(100_000)
2299 .build()
2300 .unwrap();
2301
2302 let park_ts = crate::clock::clock_monotonic_ns();
2304 let mut enc = Encoder::new_to(Vec::new()).unwrap();
2305 enc.write(&TestEvent {
2306 timestamp_ns: park_ts,
2307 value: 0,
2308 })
2309 .unwrap();
2310 writer
2311 .write_encoded_batch(&Batch::new(enc.into_inner(), 1))
2312 .unwrap();
2313 writer.flush().unwrap();
2314 writer.finalize().unwrap();
2315
2316 let all = decode_all(&std::fs::read(rotating_file(&base, 0)).unwrap());
2317
2318 let (sync_mono, sync_real) = all
2319 .iter()
2320 .find_map(|e| match e {
2321 Decoded::ClockSync {
2322 timestamp_ns,
2323 realtime_ns,
2324 } => Some((*timestamp_ns, *realtime_ns)),
2325 _ => None,
2326 })
2327 .expect("ClockSync");
2328 let park_from_file = all
2329 .iter()
2330 .find_map(|e| match e {
2331 Decoded::Data { timestamp_ns } => Some(*timestamp_ns),
2332 _ => None,
2333 })
2334 .expect("data event");
2335
2336 let offset = sync_real as i128 - sync_mono as i128;
2337 let reconstructed_wall_ns = park_from_file as i128 + offset;
2338 let now_ns = SystemTime::now()
2339 .duration_since(UNIX_EPOCH)
2340 .unwrap()
2341 .as_nanos() as i128;
2342 let diff = (reconstructed_wall_ns - now_ns).abs();
2343 assert!(
2344 diff < 5_000_000_000,
2345 "reconstructed wall clock {reconstructed_wall_ns} diverges from now {now_ns} by {diff}ns"
2346 );
2347 }
2348
2349 #[test]
2352 fn test_update_segment_metadata_appears_in_trace() {
2353 let dir = TempDir::new().unwrap();
2354 let base = dir.path().join("trace");
2355 let mut writer = DiskBuffer::builder()
2356 .base_path(dir.path())
2357 .max_file_size(100_000)
2358 .max_total_size(100_000)
2359 .build()
2360 .unwrap();
2361
2362 writer.update_segment_metadata(vec![
2364 ("bucket".into(), "my-bucket".into()),
2365 ("service_name".into(), "my-svc".into()),
2366 ]);
2367
2368 writer.write_encoded_batch(&test_batch()).unwrap();
2369 writer.flush().unwrap();
2370 writer.finalize().unwrap();
2371
2372 let all = decode_all(&std::fs::read(rotating_file(&base, 0)).unwrap());
2373 let metadata: Vec<_> = all
2374 .iter()
2375 .filter_map(|e| match e {
2376 Decoded::SegmentMetadata { entries, .. } => Some(entries.clone()),
2377 _ => None,
2378 })
2379 .collect();
2380 assert!(!metadata.is_empty(), "expected SegmentMetadata event");
2381 assert!(
2382 metadata.last().unwrap().get("bucket").map(String::as_str) == Some("my-bucket"),
2383 "S3 metadata should be in segment"
2384 );
2385 assert!(
2386 metadata
2387 .last()
2388 .unwrap()
2389 .get("service_name")
2390 .map(String::as_str)
2391 == Some("my-svc"),
2392 "S3 metadata should be in segment"
2393 );
2394 }
2395
2396 #[test]
2399 fn test_merge_preserves_s3_metadata_across_runtime_updates() {
2400 let dir = TempDir::new().unwrap();
2401 let one_event = single_event_file_size();
2402 let mut writer = DiskBuffer::builder()
2403 .base_path(dir.path())
2404 .max_file_size(one_event)
2405 .max_total_size(100_000)
2406 .build()
2407 .unwrap();
2408
2409 writer.update_segment_metadata(vec![
2411 ("bucket".into(), "my-bucket".into()),
2412 ("service_name".into(), "my-svc".into()),
2413 ]);
2414
2415 writer.update_segment_metadata(vec![("runtime.main".into(), "0,1".into())]);
2418
2419 for _ in 0..4 {
2421 writer.write_encoded_batch(&test_batch()).unwrap();
2422 }
2423 writer.flush().unwrap();
2424 writer.finalize().unwrap();
2425
2426 let mut files: Vec<_> = std::fs::read_dir(dir.path())
2427 .unwrap()
2428 .filter_map(|e| e.ok())
2429 .map(|e| e.path())
2430 .filter(|p| p.extension().is_some_and(|ext| ext == "bin"))
2431 .collect();
2432 files.sort();
2433 assert!(files.len() >= 2, "expected rotation");
2434
2435 for file in &files[1..] {
2437 let all = decode_all(&std::fs::read(file).unwrap());
2438 let meta: Vec<_> = all
2439 .iter()
2440 .filter_map(|e| match e {
2441 Decoded::SegmentMetadata { entries, .. } => Some(entries.clone()),
2442 _ => None,
2443 })
2444 .collect();
2445 let last = meta.last().expect("expected SegmentMetadata");
2446 assert!(
2447 last.get("bucket").map(String::as_str) == Some("my-bucket"),
2448 "{}: S3 metadata lost after merge",
2449 file.display()
2450 );
2451 assert!(
2452 last.get("runtime.main").map(String::as_str) == Some("0,1"),
2453 "{}: runtime metadata missing",
2454 file.display()
2455 );
2456 }
2457 }
2458
2459 #[test]
2462 fn test_update_segment_metadata_no_op_when_unchanged() {
2463 let dir = TempDir::new().unwrap();
2464 let base = dir.path().join("trace");
2465 let mut writer = DiskBuffer::builder()
2466 .base_path(dir.path())
2467 .max_file_size(100_000)
2468 .max_total_size(100_000)
2469 .build()
2470 .unwrap();
2471
2472 let entries = vec![("k".into(), "v".into())];
2473 writer.update_segment_metadata(entries.clone());
2474 writer.write_encoded_batch(&test_batch()).unwrap();
2476
2477 writer.update_segment_metadata(entries.clone());
2479 writer.write_encoded_batch(&test_batch()).unwrap();
2481 writer.flush().unwrap();
2482 writer.finalize().unwrap();
2483
2484 let all = decode_all(&std::fs::read(rotating_file(&base, 0)).unwrap());
2485 let metadata_count = all
2486 .iter()
2487 .filter(|e| matches!(e, Decoded::SegmentMetadata { .. }))
2488 .count();
2489 assert_eq!(
2490 metadata_count, 1,
2491 "identical update_segment_metadata should not trigger another write"
2492 );
2493 }
2494
2495 #[test]
2499 fn test_dial9_version_in_segment_metadata() {
2500 let dir = TempDir::new().unwrap();
2501 let path = dir.path().join("trace.bin");
2502 let mut writer = DiskBuffer::single_file(&path).unwrap();
2503 writer.write_encoded_batch(&test_batch()).unwrap();
2504 writer.flush().unwrap();
2505 writer.finalize().unwrap();
2506
2507 let sealed = dir.path().join("trace.0.bin");
2508 let all = decode_all(&std::fs::read(&sealed).unwrap());
2509 let version_value = all.iter().find_map(|e| match e {
2510 Decoded::SegmentMetadata { entries, .. } => entries.get(DIAL9_VERSION_KEY).cloned(),
2511 _ => None,
2512 });
2513 assert_eq!(
2514 version_value.as_deref(),
2515 Some(env!("CARGO_PKG_VERSION")),
2516 "expected dial9.dial9-tokio-telemetry.version entry matching CARGO_PKG_VERSION"
2517 );
2518 }
2519
2520 #[test]
2523 fn test_available_parallelism_in_segment_metadata() {
2524 let expected = std::thread::available_parallelism().map(|n| n.get().to_string());
2525 let dir = TempDir::new().unwrap();
2526 let path = dir.path().join("trace.bin");
2527 let mut writer = DiskBuffer::single_file(&path).unwrap();
2528 writer.write_encoded_batch(&test_batch()).unwrap();
2529 writer.flush().unwrap();
2530 writer.finalize().unwrap();
2531
2532 let sealed = dir.path().join("trace.0.bin");
2533 let all = decode_all(&std::fs::read(&sealed).unwrap());
2534 let value = all.iter().find_map(|e| match e {
2535 Decoded::SegmentMetadata { entries, .. } => {
2536 entries.get(PROCESS_AVAILABLE_PARALLELISM_KEY).cloned()
2537 }
2538 _ => None,
2539 });
2540 match expected {
2541 Ok(expected) => assert_eq!(
2542 value.as_deref(),
2543 Some(expected.as_str()),
2544 "expected process.available_parallelism entry matching std::thread::available_parallelism()"
2545 ),
2546 Err(_) => assert!(
2547 value.is_none(),
2548 "process.available_parallelism should be omitted when available_parallelism() fails"
2549 ),
2550 }
2551 }
2552
2553 #[test]
2556 fn test_dial9_version_user_override_wins() {
2557 let dir = TempDir::new().unwrap();
2558 let base = dir.path().join("trace");
2559 let mut writer = DiskBuffer::builder()
2560 .base_path(dir.path())
2561 .max_file_size(100_000)
2562 .max_total_size(100_000)
2563 .segment_metadata(vec![(DIAL9_VERSION_KEY.into(), "builder-override".into())])
2564 .build()
2565 .unwrap();
2566 writer.write_encoded_batch(&test_batch()).unwrap();
2567 writer.flush().unwrap();
2568 writer.rotate().unwrap();
2570 writer.update_segment_metadata(vec![(DIAL9_VERSION_KEY.into(), "runtime-override".into())]);
2571 writer.write_encoded_batch(&test_batch()).unwrap();
2572 writer.flush().unwrap();
2573 writer.finalize().unwrap();
2574
2575 let read_version = |idx: u32| -> String {
2576 let all = decode_all(&std::fs::read(rotating_file(&base, idx)).unwrap());
2577 all.iter()
2578 .find_map(|e| match e {
2579 Decoded::SegmentMetadata { entries, .. } => {
2580 entries.get(DIAL9_VERSION_KEY).cloned()
2581 }
2582 _ => None,
2583 })
2584 .expect("expected dial9.dial9-tokio-telemetry.version entry")
2585 };
2586 assert_eq!(read_version(0), "builder-override");
2587 assert_eq!(read_version(1), "runtime-override");
2588 }
2589
2590 #[tokio::test(start_paused = true)]
2602 async fn test_drained_recovers_when_active_file_deleted() {
2603 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
2604 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
2605
2606 let dir = TempDir::new().unwrap();
2607 let mut writer = DiskBuffer::builder()
2608 .base_path(dir.path())
2609 .max_file_size(u64::MAX)
2610 .max_total_size(100_000)
2611 .rotation_period(Duration::from_secs(60))
2612 .build()
2613 .unwrap();
2614
2615 writer.write_encoded_batch(&test_batch()).unwrap();
2616 writer.flush().unwrap();
2617
2618 let active_path = writer.current_active_path().to_owned();
2620 assert!(active_path.exists());
2621 std::fs::remove_file(&active_path).unwrap();
2622
2623 tokio::time::advance(Duration::from_secs(61)).await;
2625
2626 assert!(writer.should_drain(), "should_drain should fire");
2627
2628 writer
2632 .drained()
2633 .expect("drained() must recover from missing .active file");
2634
2635 assert!(
2638 !writer.should_drain(),
2639 "should_drain must return false after recovery (otherwise flush loop spins)"
2640 );
2641
2642 writer.write_encoded_batch(&test_batch()).unwrap();
2645 writer.flush().unwrap();
2646 assert!(
2647 writer.current_active_path().exists(),
2648 "writer must have a fresh active file after recovery"
2649 );
2650
2651 writer.finalize().unwrap();
2652 }
2653
2654 #[tokio::test(start_paused = true)]
2662 async fn test_drained_recovers_when_parent_dir_deleted() {
2663 use metrique_timesource::{TimeSource, tokio::set_time_source_for_current_runtime};
2664 let _guard = set_time_source_for_current_runtime(TimeSource::tokio(std::time::UNIX_EPOCH));
2665
2666 let dir = TempDir::new().unwrap();
2667 let trace_dir = dir.path().join("traces");
2668 std::fs::create_dir_all(&trace_dir).unwrap();
2669 let mut writer = DiskBuffer::builder()
2670 .base_path(&trace_dir)
2671 .max_file_size(u64::MAX)
2672 .max_total_size(100_000)
2673 .rotation_period(Duration::from_secs(60))
2674 .build()
2675 .unwrap();
2676
2677 writer.write_encoded_batch(&test_batch()).unwrap();
2678 writer.flush().unwrap();
2679
2680 std::fs::remove_dir_all(&trace_dir).unwrap();
2681 assert!(!writer.current_active_path().exists());
2682
2683 tokio::time::advance(Duration::from_secs(61)).await;
2684 assert!(writer.should_drain());
2685
2686 let _ = writer.drained();
2690 assert!(
2691 !writer.should_drain(),
2692 "should_drain must return false after a failed rotation \
2693 (otherwise the flush loop spins on every 5ms tick)"
2694 );
2695
2696 tokio::time::advance(Duration::from_millis(5)).await;
2698 let _ = writer.drained();
2699 assert!(!writer.should_drain());
2700 }
2701
2702 #[test]
2706 fn test_restart_seeds_closed_files_and_evicts() {
2707 let dir = TempDir::new().unwrap();
2708 let base = dir.path().join("trace");
2709 let one_event = single_event_file_size();
2711 {
2712 let mut w = DiskBuffer::builder()
2713 .base_path(dir.path())
2714 .max_file_size(one_event)
2715 .max_total_size(100_000)
2716 .build()
2717 .unwrap();
2718 for _ in 0..4 {
2719 w.write_encoded_batch(&test_batch()).unwrap();
2720 }
2721 w.finalize().unwrap();
2722 }
2723 let bin_count_before = (0..20)
2724 .filter(|i| std::path::Path::new(&rotating_file(&base, *i)).exists())
2725 .count();
2726 assert!(
2727 bin_count_before >= 2,
2728 "lifetime 1 should leave multiple sealed segments"
2729 );
2730
2731 let new_budget = one_event + 1; let writer = DiskBuffer::builder()
2734 .base_path(dir.path())
2735 .max_file_size(one_event)
2736 .max_total_size(new_budget)
2737 .build()
2738 .unwrap();
2739 assert!(
2741 total_disk_usage(dir.path()) <= new_budget,
2742 "disk usage exceeds shrunk budget after restart: {}",
2743 total_disk_usage(dir.path())
2744 );
2745 let next_active_path = writer.current_active_path();
2747 assert!(next_active_path.exists());
2748 assert!(
2749 next_active_path
2750 .to_str()
2751 .is_some_and(|s| s.ends_with(".bin.active"))
2752 );
2753 }
2754
2755 #[test]
2759 fn test_restart_discards_stale_active_files() {
2760 let dir = TempDir::new().unwrap();
2761 let orphan = dir.path().join("trace.99.bin.active");
2763 std::fs::write(&orphan, b"orphaned").unwrap();
2764
2765 let _w = DiskBuffer::builder()
2766 .base_path(dir.path())
2767 .max_file_size(1024)
2768 .max_total_size(100_000)
2769 .build()
2770 .unwrap();
2771 assert!(
2772 !orphan.exists(),
2773 "stale .active should be discarded on construction"
2774 );
2775 }
2776
2777 #[test]
2780 fn test_restart_counts_gz_siblings_toward_budget() {
2781 let dir = TempDir::new().unwrap();
2782 let bin = dir.path().join("trace.0.bin");
2784 let gz = dir.path().join("trace.0.bin.gz");
2785 std::fs::write(&bin, vec![0u8; 4096]).unwrap();
2786 std::fs::write(&gz, vec![0u8; 1024]).unwrap();
2787
2788 let _w = DiskBuffer::builder()
2790 .base_path(dir.path())
2791 .max_file_size(100_000)
2792 .max_total_size(100)
2793 .build()
2794 .unwrap();
2795 assert!(!bin.exists(), ".bin should be evicted under restart budget");
2796 assert!(!gz.exists(), ".bin.gz must be evicted with its .bin family");
2797 }
2798
2799 #[test]
2802 fn test_finalize_evicts_to_budget() {
2803 let dir = TempDir::new().unwrap();
2804 let one_event = single_event_file_size();
2805 let max_total_size = one_event * 2;
2806 let mut writer = DiskBuffer::builder()
2807 .base_path(dir.path())
2808 .max_file_size(one_event)
2809 .max_total_size(max_total_size)
2810 .build()
2811 .unwrap();
2812
2813 for _ in 0..10 {
2814 writer.write_encoded_batch(&test_batch()).unwrap();
2815 }
2816 writer.finalize().unwrap();
2817
2818 assert!(
2819 total_disk_usage(dir.path()) <= max_total_size,
2820 "finalize must leave disk usage within budget"
2821 );
2822 }
2823
2824 #[test]
2825 fn in_memory_builder_wires_custom_options() {
2826 let writer = MemoryBuffer::builder()
2827 .max_total_size(8 * 1024 * 1024)
2828 .max_segment_size(64 * 1024)
2829 .rotation_period(Duration::from_secs(30))
2830 .segment_metadata(vec![("svc".into(), "test".into())])
2831 .build()
2832 .unwrap();
2833 assert_eq!(writer.max_file_size, 64 * 1024);
2834 assert_eq!(writer.rotation_period, Duration::from_secs(30));
2835 assert!(
2836 writer
2837 .segment_metadata
2838 .entries
2839 .iter()
2840 .any(|(k, v)| k == "svc" && v == "test")
2841 );
2842 }
2843
2844 #[test]
2845 fn in_memory_rejects_zero_total_size() {
2846 let err = MemoryBuffer::new(0).unwrap_err();
2847 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
2848 }
2849
2850 #[test]
2851 fn in_memory_builder_enforces_3x_segment_min_total_size() {
2852 let seg: u64 = 2048;
2853 let err = MemoryBuffer::builder()
2855 .max_total_size(3 * seg - 1)
2856 .max_segment_size(seg)
2857 .build()
2858 .unwrap_err();
2859 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
2860 MemoryBuffer::builder()
2862 .max_total_size(3 * seg)
2863 .max_segment_size(seg)
2864 .build()
2865 .expect("3× segment must be accepted");
2866 }
2867
2868 #[cfg(feature = "pipeline")]
2871 mod mem_e2e_tests {
2872 use super::*;
2873 use crate::pipeline::{ProcessError, SegmentData, SegmentProcessor};
2874 use crate::worker::WorkerLoop;
2875 use std::future::Future;
2876 use std::pin::Pin;
2877 use std::sync::{Arc, Mutex};
2878 use std::time::Duration;
2879
2880 struct CapturingProcessor {
2882 segments: Arc<Mutex<Vec<Vec<u8>>>>,
2883 }
2884
2885 impl CapturingProcessor {
2886 fn new() -> (Self, Arc<Mutex<Vec<Vec<u8>>>>) {
2887 let segments = Arc::new(Mutex::new(Vec::new()));
2888 (
2889 Self {
2890 segments: segments.clone(),
2891 },
2892 segments,
2893 )
2894 }
2895 }
2896
2897 impl SegmentProcessor for CapturingProcessor {
2898 fn name(&self) -> &'static str {
2899 "Capture"
2900 }
2901 fn process(
2902 &mut self,
2903 data: SegmentData,
2904 ) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>>
2905 {
2906 self.segments
2907 .lock()
2908 .unwrap()
2909 .push(data.payload().clone().into_vec());
2910 Box::pin(async move { Ok(data) })
2911 }
2912 }
2913
2914 async fn run_mem_e2e(mut writer: MemoryBuffer, events: usize) -> Vec<Vec<u8>> {
2917 let fs = writer.fs_handle().expect("memory writer exposes its Fs");
2918 for _ in 0..events {
2919 writer.write_encoded_batch(&test_batch()).unwrap();
2920 }
2921 writer.finalize().unwrap();
2923
2924 let (capture, captured) = CapturingProcessor::new();
2925 let stop = tokio_util::sync::CancellationToken::new();
2927 let mut worker = WorkerLoop::new(
2928 fs,
2929 Duration::from_millis(5),
2930 vec![Box::new(capture)],
2931 stop,
2932 metrique_writer::sink::DevNullSink::boxed(),
2933 None,
2934 )
2935 .await
2936 .expect("initialize worker");
2937 worker.run().await;
2938
2939 let segments = captured.lock().unwrap();
2940 segments.clone()
2941 }
2942
2943 fn count_payload_events(segments: &[Vec<u8>]) -> usize {
2946 segments
2947 .iter()
2948 .flat_map(|s| decode_all(s))
2949 .filter(|e| matches!(e, Decoded::Data { .. }))
2950 .count()
2951 }
2952
2953 #[tokio::test]
2954 async fn mem_writer_e2e_delivers_all_events() {
2955 const EVENTS: usize = 25;
2956
2957 let segments = run_mem_e2e(MemoryBuffer::new(1 << 20).unwrap(), EVENTS).await;
2958
2959 assert!(!segments.is_empty(), "worker captured no segments");
2960 assert_eq!(
2961 count_payload_events(&segments),
2962 EVENTS,
2963 "every written event must reach the processor"
2964 );
2965 }
2966
2967 #[tokio::test]
2970 async fn mem_writer_e2e_delivers_all_events_across_rotations() {
2971 const EVENTS: usize = 60;
2972
2973 let writer = MemoryBuffer::builder()
2975 .max_total_size(16 * 1024 * 1024)
2976 .max_segment_size(256)
2977 .build()
2978 .unwrap();
2979 let segments = run_mem_e2e(writer, EVENTS).await;
2980
2981 assert!(
2982 segments.len() >= 2,
2983 "tiny segments must force rotation, got {} segment(s)",
2984 segments.len()
2985 );
2986 assert_eq!(
2987 count_payload_events(&segments),
2988 EVENTS,
2989 "every event across all rotated segments must reach the processor"
2990 );
2991 }
2992 }
2993}