Skip to main content

dial9_core/
buffer.rs

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
23/// Marker trait for `SegmentWriter`'s backend mode. Sealed: only [`Disk`]
24/// and [`Memory`] implement it.
25pub trait BufferMode: mode_sealed::Sealed + Send + 'static {
26    /// Whether the writer mode is disk-backed.
27    const IS_DISK: bool;
28}
29
30/// Disk-backed mode (default).
31#[derive(Debug)]
32#[non_exhaustive]
33pub struct Disk;
34/// In-memory mode.
35#[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
48/// Alias for the disk-backed writer (the default mode).
49pub type DiskBuffer = SegmentWriter<Disk>;
50/// Alias for the in-memory writer.
51pub type MemoryBuffer = SegmentWriter<Memory>;
52
53/// Segment-metadata key carrying the crates.io version of
54/// `dial9-tokio-telemetry`. Populated by default, any user-supplied entry with
55/// this key take precedence.
56const DIAL9_VERSION_KEY: &str = "dial9.dial9-tokio-telemetry.version";
57
58/// Compile-time value for `DIAL9_VERSION_KEY`.
59const DIAL9_VERSION_VALUE: &str = env!("CARGO_PKG_VERSION");
60
61/// Segment-metadata key carrying the logical CPU capacity available to the
62/// process. Populated by default when the platform can report it.
63const 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    /// Build segment metadata from user-supplied entries on top of the default
91    /// `dial9.dial9-tokio-telemetry.version` key. User entries with the same key override the default.
92    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    /// Merge incoming entries with existing ones. Incoming entries take priority
99    /// on key conflict; existing entries with keys not in the incoming set are preserved.
100    /// Returns `true` if the resulting entries differ from the previous state.
101    ///
102    /// The "unchanged -> no rewrite" detection (`merged == self.entries`) is a
103    /// positional `Vec` compare. A source re-emitting the same entries is only
104    /// deduped to a no-op if it emits them in a stable order across calls, so
105    /// every `Source::segment_metadata` MUST produce a deterministic order. A
106    /// nondeterministic iteration order (e.g. iterating a `HashMap`) would
107    /// reorder `merged`, fail this compare, and rewrite segment metadata on
108    /// every change-cycle.
109    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
124/// Default rotation period: 1 minute.
125const DEFAULT_ROTATION_PERIOD: Duration = Duration::from_secs(60);
126
127/// Default maximum interval between thread-local buffer drains.
128const DEFAULT_DRAIN_INTERVAL: Duration = Duration::from_secs(30);
129
130/// Default segment filename stem for rotating writers. Segments are then
131/// `trace.0.bin`, `trace.1.bin`, and so on inside the configured directory.
132const SEGMENT_STEM: &str = "trace";
133
134const BYTES_PER_MIB: u64 = 1024 * 1024;
135
136/// Hard cap on the builder-derived per-file size, regardless of the total
137/// disk budget. Time-based rotation should fire first under normal load;
138/// this cap keeps individual segments small enough to remain manageable.
139const MAX_FILE_SIZE_CAP: u64 = 100 * BYTES_PER_MIB;
140
141/// Default per-file rotation threshold derived from the total disk budget.
142/// Picks a quarter of the budget so a single segment never dominates
143/// retention, capped at 100 MiB.
144fn derive_max_file_size(max_total_size: u64) -> u64 {
145    (max_total_size / 4).min(MAX_FILE_SIZE_CAP)
146}
147
148/// A writer that rotates trace segments to bound resource usage and time.
149/// Generic over backend: use [`DiskBuffer`] (files) or [`MemoryBuffer`].
150///
151/// Rotation triggers when *either* condition is met:
152/// - `max_file_size`: the active segment exceeds this many bytes
153/// - `rotation_period`: this much monotonic time has elapsed since the writer
154///   (or the previous rotation) started (default: 1 minute)
155///
156/// **Prefer time-based rotation.** Time-based rotation is coordinated with the
157/// flush loop: thread-local buffers are drained before the segment is sealed,
158/// so each segment contains events from a clean, non-overlapping time window.
159/// Size-based rotation fires immediately when the threshold is crossed and does
160/// not drain thread-local buffers, so segments may contain events that overlap
161/// in time. Set `max_file_size` large enough that time-based rotation fires
162/// first under normal conditions (e.g. 100 MB or more). Size-based rotation
163/// then acts as a safety valve for unexpected data bursts. When using
164/// [`DiskBuffer::builder`] without specifying `max_file_size`, it
165/// defaults to `min(100 MiB, max_total_size / 4)` on disk.
166///
167/// `max_total_size` is the retention budget across closed segments. The
168/// oldest segments are dropped once the total exceeds this budget.
169///
170/// The trace lives in a directory (`dir`); disk segments are named
171/// `{dir}/{stem}.0.bin`, `{dir}/{stem}.1.bin`, etc., each a self-contained
172/// trace with its own header. Rotating writers use the stem `trace`;
173/// [`single_file`](Self::single_file) takes the stem from the given file name.
174pub struct SegmentWriter<Mode: BufferMode = Disk> {
175    /// Directory the segments live in.
176    dir: PathBuf,
177    /// Segment filename stem, e.g. `trace` for `trace.0.bin`.
178    stem: String,
179    max_file_size: u64,
180    max_total_size: u64,
181    /// How often to rotate based on monotonic time. `Duration::MAX` disables
182    /// time-based rotation (used by `single_file()`).
183    rotation_period: Duration,
184    /// The next monotonic instant at which time-based rotation should fire,
185    /// or `None` if time-based rotation is disabled.
186    next_rotation_time: Option<Instant>,
187    /// Tracks (seg_ref, size) of closed segments oldest-first for disk eviction.
188    /// Always empty in memory mode (eviction handled by the memory backend).
189    closed_files: VecDeque<(SegmentRef, u64)>,
190    /// Path of the currently active (being-written) segment.
191    /// Used as a HashMap key in memory mode; a real path in disk mode.
192    active_path: PathBuf,
193    state: WriterState,
194    next_index: u32,
195    /// Metadata written at the start of each segment. Updated by the flush
196    /// thread to include runtime names alongside any user-provided entries.
197    segment_metadata: SegmentMetadata,
198    /// Events silently dropped because the writer was finished/stopped.
199    dropped_events: usize,
200    /// Whether any real (non-metadata) events have been written to the current segment.
201    /// Reset on rotation; used by `finalize()` to avoid sealing empty segments.
202    has_real_events: bool,
203    /// How often the flush loop should drain thread-local buffers, independent
204    /// of rotation. Defaults to `min(rotation_period, 30s)`.
205    drain_interval: Duration,
206    /// Next monotonic instant at which `should_drain()` returns true.
207    next_drain_time: Instant,
208    /// Unified filesystem/channel abstraction.
209    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// the write side is obviously larger than the `Finished` size so clippy warns on this
227// but we don't want to force going through a pointer every time we want to write.
228#[allow(clippy::large_enum_variant)]
229enum WriterState {
230    /// Writer is open and events can be written
231    Active {
232        writer: RawEncoder<BufWriter<ActiveHandle>>,
233        need_metadata: bool,
234    },
235
236    /// Writer has been finalized or stopped — no encoder, no fd, no writes.
237    Finished,
238}
239
240#[bon::bon]
241impl SegmentWriter<Disk> {
242    /// Create a `DiskBufferBuilder` for advanced configuration.
243    ///
244    /// When `max_file_size` is omitted, it defaults to
245    /// `min(100 MiB, max_total_size / 4)`.
246    #[builder(builder_type = DiskBufferBuilder, finish_fn = build)]
247    pub fn builder(
248        base_path: impl Into<PathBuf>,
249        /// Per-file rotation threshold in bytes. Defaults to
250        /// `min(100 MiB, max_total_size / 4)` when not set.
251        max_file_size: Option<u64>,
252        max_total_size: u64,
253        /// How often to rotate, measured in monotonic time since the writer
254        /// (or the previous rotation) started. Defaults to 60 seconds.
255        /// `Duration::MAX` disables time-based rotation.
256        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        // The trace is a directory of `{stem}.{index}.bin` segments.
281        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        // Enforce the budget immediately so artifacts from prior writer
320        // lifetimes don't push us over the cap before we even rotate once.
321        writer.evict_oldest()?;
322        Ok(writer)
323    }
324
325    /// Set the namespace for this writer:
326    /// - `boot_id`: The boot id for the namespace.
327    /// - `lock`: The lock file for the namespace.
328    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    /// Create a writer that writes to a single file with no rotation or eviction.
334    /// The segment is written to `{stem}.0.bin.active` while active, then sealed
335    /// to `{stem}.0.bin` on `finalize`. The background worker will symbolize
336    /// and gzip it to `{stem}.0.bin.gz`.
337    ///
338    /// Note: This API does not allow the ability to provide custom segment metadata.
339    /// Time-based rotation is disabled.
340    pub fn single_file(path: impl Into<PathBuf>) -> std::io::Result<Self> {
341        let path = path.into();
342        // Unlike rotating writers, `single_file` takes both the directory and
343        // the stem from the caller's file path.
344        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
384/// Default segment size when no explicit segment size is provided.
385/// Always at least 8 slots of burst headroom in the ring.
386fn 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    /// Create an in-memory writer with a total byte budget. Segments live in process heap
394    /// instead of files. Auto-picks a reasonable segment size,
395    /// use [`builder`](Self::builder) for explicit control.
396    ///
397    /// Same rotation semantics as the disk path. Errors when
398    /// `max_total_size == 0`.
399    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 for in-memory writer configuration.
409    #[builder(builder_type = MemoryBufferBuilder, finish_fn = build)]
410    pub fn builder(
411        max_total_size: u64,
412        /// Override the default segment size.
413        max_segment_size: Option<u64>,
414        /// Wall-clock rotation period.
415        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        // The active buffer and the worker's in-flight segment live outside the ring, so the ring needs
451        // room for at least one sealed segment on top of that reserve.
452        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        // The memory backend ignores paths, but `active_path` still needs a
465        // dir/stem to build its HashMap keys.
466        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        // Drain at least as often as we rotate.
473        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    /// Per-process boot identifier, if namespace isolation is active. This is
501    /// the name of the [`trace_dir`](Self::trace_dir) subdirectory.
502    pub fn boot_id(&self) -> Option<&str> {
503        self.boot_id.as_deref()
504    }
505
506    /// Directory this writer's trace segments live in. When namespace
507    /// isolation is active this is the per-process `{configured_dir}/{boot_id}/`
508    /// subdirectory; otherwise it is the configured directory directly. Use
509    /// this to locate the segment files on disk.
510    pub fn trace_dir(&self) -> &Path {
511        &self.dir
512    }
513
514    /// Segment filename stem, e.g. `trace` for `trace.0.bin`.
515    pub fn trace_stem(&self) -> &str {
516        &self.stem
517    }
518
519    /// The path of the currently active (being-written) segment file.
520    pub fn current_active_path(&self) -> &Path {
521        &self.active_path
522    }
523
524    /// Create an encoder, write the file header, segment metadata, and a
525    /// clock-sync anchor, then convert to a [`RawEncoder`] for the
526    /// remainder of the file's lifetime.
527    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    /// Write a `SegmentMetadataEvent` and a fresh `ClockSyncEvent` into
557    /// the current active segment.
558    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    /// Path for a segment that is actively being written.
578    fn active_path(dir: &Path, stem: &str, index: u32) -> PathBuf {
579        dir.join(format!("{stem}.{index}.bin.active"))
580    }
581
582    /// Compute the next rotation deadline as `now + period`, or `None` when
583    /// `period == Duration::MAX` (time-based rotation disabled).
584    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        // Advance timers up front. If anything below fails the flush loop must
594        // NOT see should_drain() return true on the next 5ms tick — otherwise
595        // it busy-spins re-attempting the same failing rotate.
596        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        // Take ownership of the encoder (state is Finished until new segment opens).
601        let WriterState::Active {
602            writer: mut raw, ..
603        } = std::mem::replace(&mut self.state, WriterState::Finished)
604        else {
605            return Ok(());
606        };
607
608        // Best-effort flush. If the underlying file is gone the buffered bytes
609        // are already lost; proceed to rotate rather than erroring.
610        let _ = raw.flush();
611        let closed_size = raw.bytes_written();
612        let current_index = self.next_index - 1;
613
614        // Extract the ActiveHandle for sealing.
615        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        // Seal the current segment. If `.active` was removed externally
621        // (disk only: operator, log rotation, container teardown) abandon the
622        // segment and start a fresh one.
623        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                // state is already Finished from mem::replace above
640                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        // Open the new active segment. The backend self-heals if a disk
648        // parent directory was removed underneath us, any other failure
649        // leaves state = Finished so the writer stops cleanly rather than
650        // retrying every drain cycle.
651        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    /// Total size across all closed + active segments (disk mode only).
672    /// Always returns 0 in memory mode, eviction is handled by the memory backend.
673    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        // Always keep at least the current file.
690        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 even the current file alone exceeds total budget, stop writing.
696        if self.total_size() > self.max_total_size {
697            self.state = WriterState::Finished;
698        }
699        Ok(())
700    }
701
702    /// Rotate if the current file exceeds max_file_size.
703    /// Called after writing a complete logical unit (def + event).
704    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    /// The filesystem backend, handed to the worker so it can drain sealed
717    /// segments. Only needed when the pipeline worker is compiled in.
718    #[cfg(feature = "pipeline")]
719    pub(crate) fn fs_handle(&self) -> Option<Arc<Fs>> {
720        Some(Arc::clone(&self.fs))
721    }
722
723    /// Flush buffered data to the underlying storage.
724    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    /// Merge the segment metadata entries written into the next rotated segment.
737    ///
738    /// Accepts any iterator so callers can drain a reused buffer (retaining its
739    /// capacity) instead of handing over an owned `Vec`. A `Vec` still works.
740    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        // Periodic drain without rotation; advance the drain timer.
770        self.next_drain_time = now + self.drain_interval;
771        Ok(false)
772    }
773
774    /// Finalize the writer: flush, seal the active segment, and prevent further
775    /// writes. Terminal — the writer is inert afterward.
776    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        // Best-effort flush: if the file is gone the bytes are already lost.
785        let _ = self.flush();
786
787        // Take ownership of the encoder (state -> Finished).
788        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            // No real events — just header + metadata. Remove instead of
826            // sealing so the background worker doesn't upload an empty segment.
827            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        // Final sealed segment must count toward the eviction budget too,
840        // otherwise finalize can leave the directory over `max_total_size`.
841        // No-ops on memory mode (`!M::IS_DISK`).
842        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    // `pub(crate)` in production: the flush loop drives it. `pub` under
851    // `test-util` so sibling-crate tests/benches can write pre-encoded batches.
852    crate::test_util_pub! {
853    /// Transcode an encoded batch into the active segment.
854    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            // Note: we do NOT advance next_rotation_time or next_drain_time
862            // when the first event arrives in an empty segment, even if the
863            // timers are stale. The drain state machine (Idle → EpochBumped →
864            // drain) takes 3 flush cycles (~15ms) to complete, so by the time
865            // drained() is called there will be multiple batches in the segment,
866            // not a single event. Advancing the timers here would skip rotation
867            // windows and produce fewer segments than expected.
868            // Raw-copy the thread-local batch. Each batch is self-contained
869            // (starts with its own header), so the next batch's header acts as
870            // the reset frame for decoders.
871            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    /// A minimal data event for exercising the writer, distinct from the bus's
903    /// own framing events (`ClockSyncEvent`/`SegmentMetadataEvent`) so the
904    /// decode helper can tell real events from the per-segment framing.
905    #[derive(TraceEvent)]
906    #[traceevent(wire_slot)]
907    struct TestEvent {
908        #[traceevent(timestamp)]
909        timestamp_ns: u64,
910        value: u64,
911    }
912
913    /// Decoded view of a trace, classified by frame name via the registry.
914    #[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    /// Encode a single event into a self-contained batch (header + event),
977    /// matching the format produced by ThreadLocalBuffer.
978    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    /// Read all data (non-framing) events from a trace file.
993    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    /// Total size of all trace files (.bin and .active) in a directory.
1002    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    /// Write one batch to a temp file and return the file size.
1016    /// This captures the actual overhead (header + schema + event) so tests
1017    /// don't depend on hardcoded format sizes.
1018    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        // Two events should be larger than one event
1103        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        // No real events were written, so finalize removes the empty segment.
1132        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        // Set max_file_size to fit ~1 event so rotation triggers quickly
1147        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        // All 3 events should be readable across rotated files
1161        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        // Key invariant: total disk usage stays within budget
1194        assert!(total_disk_usage(dir.path()) <= max_total_size);
1195
1196        // Oldest files should be evicted
1197        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        // Small file size to force rotation, total budget fits ~1 file
1206        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        // Should have stopped writing — total events across all files < 100
1221        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    /// Bug: write_encoded_batch sets stopped=true when total_size slightly exceeds
1238    /// max_total_size, without attempting eviction. This happens right after
1239    /// rotate() + evict_oldest() brings total_size just under budget, then the
1240    /// first batch in the new file pushes it a few bytes over. The writer
1241    /// permanently stops even though eviction could free space.
1242    ///
1243    /// Reproduces the stress test failure: 64-worker runtime with 1MB segments
1244    /// and 100MB budget stops producing segments after ~100 rotations.
1245    #[test]
1246    fn test_writer_stops_on_tiny_overshoot_after_eviction() {
1247        let dir = TempDir::new().unwrap();
1248        // Use max_file_size that doesn't evenly divide by batch size,
1249        // so files end up slightly under max_file_size (with leftover bytes).
1250        // Over 100 files, these leftovers accumulate and push total_size
1251        // past max_total_size after eviction.
1252        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        // Write many batches. The batch size doesn't divide evenly into
1263        // (max_file_size - header), so each file wastes a few bytes. After
1264        // 100 rotations, total_size drifts above max_total_size.
1265        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        // Should have created multiple files with sequential naming
1298        assert!(
1299            std::path::Path::new(&rotating_file(&base, 0)).exists(),
1300            "File 0 should exist"
1301        );
1302        // All events should be readable
1303        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        // All 3 events should be readable across the rotated files.
1334        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        // Each rotated file must be a self-contained, readable trace.
1365        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() // panics if corrupt
1370                } 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        // Total budget smaller than one file — stops immediately
1382        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        // Repeated flush after stop should not error
1393        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        // All events should be readable across files.
1415        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        // Exactly fits one event file — second event triggers rotation
1431        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        // Both events readable across files
1444        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        // Current file should have .active suffix
1470        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        // Write 2 events — triggers rotation after first
1488        writer.write_encoded_batch(&test_batch()).unwrap();
1489        writer.write_encoded_batch(&test_batch()).unwrap();
1490        writer.flush().unwrap();
1491
1492        // First file should be sealed (.bin), second should be active
1493        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        // Small max_file_size so one event triggers rotation.
1537        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        // Write an event — this fills segment 0 and triggers rotation to segment 1.
1544        writer.write_encoded_batch(&test_batch()).unwrap();
1545        // Segment 0 is sealed, segment 1 is active with only header + metadata.
1546        assert!(dir.path().join("trace.0.bin").exists());
1547        assert!(dir.path().join("trace.1.bin.active").exists());
1548
1549        // Finalize should remove the empty segment 1 instead of sealing it.
1550        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        // Segment 0 should still exist.
1560        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        // single_file seals to test.0.bin after finalize, no leftover .active
1573        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        // Simulate the flush thread merging static + runtime→worker entries.
1695        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        // Write enough events to trigger rotation — rotated segments should
1700        // contain both static and dynamic metadata.
1701        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        // First segment was constructed before update_dynamic_metadata, so
1717        // it only has static metadata. Rotated segments have both.
1718        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        // Metadata should be present and carry only the built-in dial9.dial9-tokio-telemetry.version entry
1761        // (no user-supplied entries via single_file()).
1762        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    /// When the background worker has renamed a sealed `.bin` to `.bin.gz`,
1777    /// eviction should clean up the `.gz` variant instead of silently leaking it.
1778    #[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        // Budget fits many files so segment 0 is not immediately evicted.
1784        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        // Write two batches: the first fills segment 0, the second triggers
1793        // rotation (sealing segment 0 as trace.0.bin) and starts segment 1.
1794        writer.write_encoded_batch(&test_batch()).unwrap();
1795        writer.write_encoded_batch(&test_batch()).unwrap();
1796        // Segment 0 is now sealed as trace.0.bin.
1797
1798        // Simulate the background worker renaming trace.0.bin → trace.0.bin.gz.
1799        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        // Now shrink the budget so the next rotation triggers eviction of
1805        // segment 0 (which has been renamed to .bin.gz on disk).
1806        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        // The .bin.gz file should have been cleaned up by eviction.
1813        assert!(!seg0_gz.exists(), "trace.0.bin.gz should have been evicted");
1814    }
1815
1816    /// Eviction must never drop below the most-recent segment, even when that
1817    /// single segment alone exceeds `max_total_size`. In that case it retains
1818    /// the segment on disk (so on-disk usage legitimately exceeds the budget)
1819    /// and signals "stop writing" by transitioning to `Finished`.
1820    ///
1821    /// This is the floor that makes an end-to-end `on-disk bytes <=
1822    /// max_total_size` assertion unsound — see `tests/writeback_no_leaked_gz.rs`.
1823    #[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        // No rotation (huge per-file size) so the single active segment is the
1828        // only one; a budget smaller than one segment forces the floor.
1829        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        // The lone active segment already exceeds the total budget.
1844        assert!(
1845            writer.total_size() > max_total_size,
1846            "single segment ({}) should exceed budget ({max_total_size})",
1847            writer.total_size()
1848        );
1849
1850        // Eviction has no closed segments to drop and must NOT delete the
1851        // current (most-recent) segment. It signals "stop" instead.
1852        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        // The most-recent segment is retained on disk despite exceeding the
1859        // budget — eviction never drops below one segment.
1860        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    // ---- Time-based rotation tests ----
1871
1872    #[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        // Advance past the 60s boundary
1892        tokio::time::advance(Duration::from_secs(61)).await;
1893
1894        // Time-based rotation is now driven by drained(), not write_encoded_batch.
1895        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    /// The first rotation must happen exactly `rotation_period` after the writer
1919    /// is created, not earlier due to wall-clock alignment. Starting at a non-aligned
1920    /// wall-clock time (UNIX_EPOCH + 22s) with a 60s period and advancing 50s must
1921    /// NOT rotate. So only 50s of monotonic time have elapsed since the writer started.
1922    #[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        // 50s of monotonic time have elapsed under the 60s period, so no rotation.
1942        // On the old wall-clock-aligned implementation this would advance past the
1943        // 60s wall-clock boundary (22s + 50s = 72s ≥ 60s) and incorrectly rotate.
1944        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        // after the period DOES elapse, rotation fires.
1956        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        // Advance past the boundary without writing any events
1983        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        // Advance well past the boundary with no events
2129        tokio::time::advance(Duration::from_secs(300)).await;
2130
2131        // First event after the gap — should NOT trigger rotation
2132        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        // Second event shortly after — still within the new boundary
2139        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        // ClockSync must precede the first data event so a streaming
2171        // decoder never sees a data timestamp without an anchor.
2172        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        // SegmentMetadata.timestamp_ns should remain monotonic-scale,
2210        // not epoch wall-clock.
2211        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    /// A hand-built legacy-shaped buffer (SegmentMetadata + data event,
2262    /// no ClockSyncEvent) must still round-trip through the decoder.
2263    #[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        // Use a real monotonic reading so reconstruction lands near now.
2303        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    /// S3-style metadata set via `update_segment_metadata` before any events
2350    /// are written must appear in the segment's SegmentMetadata event.
2351    #[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        // Simulate the recorder builder setting S3 metadata
2363        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    /// Simulates the flush loop pattern: S3 metadata is set once, then
2397    /// runtime entries are merged repeatedly. S3 metadata must survive.
2398    #[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        // Step 1: S3 metadata set (like the recorder builder)
2410        writer.update_segment_metadata(vec![
2411            ("bucket".into(), "my-bucket".into()),
2412            ("service_name".into(), "my-svc".into()),
2413        ]);
2414
2415        // Step 2: flush loop merges only runtime entries — S3 metadata
2416        // set in step 1 must be preserved by the merge logic.
2417        writer.update_segment_metadata(vec![("runtime.main".into(), "0,1".into())]);
2418
2419        // Write enough to trigger rotation
2420        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        // Rotated segments should contain both S3 and runtime metadata
2436        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    /// Repeated calls to `update_segment_metadata` with identical entries
2460    /// should not set `need_metadata`, avoiding redundant writes.
2461    #[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        // First batch writes metadata
2475        writer.write_encoded_batch(&test_batch()).unwrap();
2476
2477        // Same entries again — should be a no-op
2478        writer.update_segment_metadata(entries.clone());
2479        // Second batch should NOT write another metadata event
2480        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    /// The crates.io version of the writer's crate is embedded in every
2496    /// segment's metadata under `dial9.dial9-tokio-telemetry.version`. Regression test for
2497    /// https://github.com/dial9-rs/dial9/issues/423.
2498    #[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    /// The process's available logical CPU capacity is embedded in every
2521    /// segment's metadata when the platform can report it.
2522    #[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    /// User-supplied `dial9.dial9-tokio-telemetry.version` entries win over the built-in default,
2554    /// both at builder time and via `update_segment_metadata`.
2555    #[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        // Rotate and then runtime-override on the next segment.
2569        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    /// Regression test for https://github.com/dial9-rs/dial9/issues/386
2591    ///
2592    /// If the `.active` file is removed externally (e.g. by an operator,
2593    /// log-rotation tool, or container teardown) the flush loop calls
2594    /// `drained()` → `rotate()` → `fs::rename(.active, .bin)` which fails
2595    /// with `NotFound`. Without recovery, `next_drain_time` is never
2596    /// advanced, so `should_drain()` returns true on every subsequent
2597    /// 5ms tick and the flush thread busy-loops.
2598    ///
2599    /// `drained()` must recover by abandoning the missing segment, opening a
2600    /// fresh one, and advancing the drain/rotation timers.
2601    #[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        // Simulate external deletion of the .active file.
2619        let active_path = writer.current_active_path().to_owned();
2620        assert!(active_path.exists());
2621        std::fs::remove_file(&active_path).unwrap();
2622
2623        // Cross the rotation boundary so drained() will try to rotate.
2624        tokio::time::advance(Duration::from_secs(61)).await;
2625
2626        assert!(writer.should_drain(), "should_drain should fire");
2627
2628        // drained() must succeed despite the missing .active file. Returning
2629        // an error here is what causes the flush thread to busy-loop because
2630        // the timers are never advanced.
2631        writer
2632            .drained()
2633            .expect("drained() must recover from missing .active file");
2634
2635        // After recovery, should_drain() must return false — otherwise the
2636        // flush thread would spin calling drained() every 5ms.
2637        assert!(
2638            !writer.should_drain(),
2639            "should_drain must return false after recovery (otherwise flush loop spins)"
2640        );
2641
2642        // The writer must still be usable: a fresh active file exists and
2643        // subsequent writes succeed.
2644        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    /// Companion to `test_drained_recovers_when_active_file_deleted` covering
2655    /// the more realistic case where the entire trace directory has been
2656    /// removed (e.g. `rm -rf /var/log/dial9/`). Both the rename AND the
2657    /// `File::create` for the new segment fail with `NotFound`. `drained()`
2658    /// must still advance timers so `should_drain()` stops firing — the
2659    /// writer can transition to `Finished`, but the flush loop must NOT
2660    /// busy-spin.
2661    #[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        // `drained()` may surface the underlying error, but the critical
2687        // invariant is that `should_drain()` must NOT fire on the next tick —
2688        // otherwise the flush thread busy-loops.
2689        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        // Subsequent drained() calls must not re-fire either.
2697        tokio::time::advance(Duration::from_millis(5)).await;
2698        let _ = writer.drained();
2699        assert!(!writer.should_drain());
2700    }
2701
2702    /// Across a process restart, retained `.bin`/`.bin.gz` artifacts from the
2703    /// previous lifetime must count toward `max_total_size`. Without this, a
2704    /// crash-restart loop grows the trace directory unbounded.
2705    #[test]
2706    fn test_restart_seeds_closed_files_and_evicts() {
2707        let dir = TempDir::new().unwrap();
2708        let base = dir.path().join("trace");
2709        // Lifetime 1: write a few sealed segments.
2710        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        // Lifetime 2: shrink the budget so existing artifacts must be evicted.
2732        let new_budget = one_event + 1; // fits ~1 retained segment + the new active one
2733        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        // Discovery + immediate evict_oldest should have shed older segments.
2740        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        // Next active index must not collide with retained segments.
2746        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    /// Stale `.active` files from a dead writer can't be processed by the
2756    /// worker — they must be cleaned up on startup so the next writer doesn't
2757    /// trip over orphaned indices.
2758    #[test]
2759    fn test_restart_discards_stale_active_files() {
2760        let dir = TempDir::new().unwrap();
2761        // Simulate an orphan from a previous, crashed writer.
2762        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    /// `.bin.gz` write-back siblings must count toward the eviction budget so
2778    /// post-processing doesn't push retention past the cap.
2779    #[test]
2780    fn test_restart_counts_gz_siblings_toward_budget() {
2781        let dir = TempDir::new().unwrap();
2782        // Simulate a previous lifetime where WriteBack produced a .bin.gz.
2783        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        // Budget too small for both. Restart must evict the whole family.
2789        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    /// finalize() must run eviction so the final sealed segment counts toward
2800    /// the budget. Without it, finalize can leave the directory over cap.
2801    #[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        // Below the boundary: rejected (no room for even one ring slot).
2854        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        // At boundary: accepted (1 active + 1 in-flight + 1 ring slot).
2861        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    /// End-to-end writer -> worker tests: drive a memory writer through the
2869    /// real worker pipeline and assert every written event reaches a processor.
2870    #[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        /// Captures each processed segment's payload bytes.
2881        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        /// Exercises the full seam: write -> Fs::Mem seal -> ring -> finalize
2915        /// (mark_writer_done) -> WorkerLoop::run drain-to-empty -> processor.
2916        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            // Seals the active segment onto the ring and signals writer_done.
2922            writer.finalize().unwrap();
2923
2924            let (capture, captured) = CapturingProcessor::new();
2925            // stop is never cancelled: the loop exits via writer_done only.
2926            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        /// Count decoded payload events across `segments`, dropping the
2944        /// per-segment metadata/clock-sync framing the writer emits.
2945        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        /// Same, but a tiny `max_segment_size` forces several rotations so the
2968        /// worker delivers multiple sealed segments.
2969        #[tokio::test]
2970        async fn mem_writer_e2e_delivers_all_events_across_rotations() {
2971            const EVENTS: usize = 60;
2972
2973            // Huge ring (nothing evicts) + tiny segments (rotate every few batches).
2974            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}