dial9-core 0.5.1

use crate::collector::CentralCollector;
use crate::encoder;
use crate::encoder::TlBufferHandle;
use crate::metrics::TlDrainStats;
use crate::primitives::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use crate::primitives::sync::{Arc, Mutex};
use std::time::Duration;

/// Recording lifecycle states. `Stopped` is terminal: the flush thread is
/// gone, nothing drains buffers anymore.
#[repr(u8)]
enum State {
    Disabled = 0,
    Enabled = 1,
    Stopped = 2,
}

crate::test_util_pub! {
/// Runtime-agnostic core recording state.
struct SharedState {
    /// Recording lifecycle: `Disabled ⇄ Enabled → Stopped` (terminal).
    state: AtomicU8,
    pub(crate) collector: Arc<CentralCollector>,
    /// Absolute `CLOCK_MONOTONIC` nanosecond timestamp captured at trace start.
    pub(crate) start_time_ns: u64,
    /// Epoch counter bumped by the flush thread every ~30s. Thread-local
    /// buffers stamp this value on each self-flush so the flush thread can
    /// skip busy workers when draining.
    pub(crate) drain_epoch: AtomicU64,
    /// Weak handles to all registered thread-local buffers. The flush thread
    /// uses these to intrusively drain idle/silent buffers.
    tl_buffers: Mutex<Vec<TlBufferHandle>>,
    /// Data sources (CPU profiler, sched profiler, etc.) that the flush thread drains.
    pub(crate) sources: Mutex<Vec<Box<dyn crate::source::Source>>>,
    /// On-demand dump trigger, set once at build time when the runtime is
    /// built with `with_dump_trigger`. Reached by application code through
    /// [`Dial9Handle::dump_trigger`](super::handle::Dial9Handle::dump_trigger).
    #[cfg(feature = "pipeline")]
    dump_trigger: std::sync::OnceLock<crate::dump::DumpTrigger>,
}
}

impl std::fmt::Debug for SharedState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SharedState")
            .field("enabled", &self.is_enabled())
            .field("start_time_ns", &self.start_time_ns)
            .finish_non_exhaustive()
    }
}

impl SharedState {
    crate::test_util_pub! {
        fn new(start_time_ns: u64) -> Self {
            Self {
                state: AtomicU8::new(State::Disabled as u8),
                collector: Arc::new(CentralCollector::new()),
                start_time_ns,
                drain_epoch: AtomicU64::new(0),
                tl_buffers: Mutex::new(Vec::new()),
                sources: Mutex::new(Vec::new()),
                #[cfg(feature = "pipeline")]
                dump_trigger: std::sync::OnceLock::new(),
            }
        }
    }

    crate::test_util_pub! {
        /// Register a data source to be drained by the flush thread each cycle.
        fn push_source(&self, source: Box<dyn crate::source::Source>) {
            self.sources.lock().unwrap().push(source);
        }
    }

    crate::test_util_pub! {
        /// Run `f` against the registered sources. Returns `None` if the lock is
        /// poisoned. Used to drive the per-thread source lifecycle hooks.
        fn with_sources_mut<R>(
            &self,
            f: impl FnOnce(&mut [Box<dyn crate::source::Source>]) -> R,
        ) -> Option<R> {
            self.sources.lock().ok().map(|mut sources| f(&mut sources))
        }
    }

    /// Drop every registered source.
    ///
    /// Sources own OS resources (perf fds, mmap rings) and may hold a
    /// [`Dial9Handle`](crate::handle::Dial9Handle) back to this state, which
    /// would otherwise keep the `Arc` alive forever. Releasing them here bounds
    /// both to the recorder's lifetime.
    pub(crate) fn clear_sources(&self) {
        match self.sources.lock() {
            Ok(mut sources) => sources.clear(),
            Err(_) => tracing::warn!("sources lock poisoned, sources left registered"),
        }
    }

    /// Like [`with_sources_mut`](Self::with_sources_mut), but `f` receives the
    /// list itself so it can register sources too.
    pub(crate) fn with_sources_vec<R>(
        &self,
        f: impl FnOnce(&mut Vec<Box<dyn crate::source::Source>>) -> R,
    ) -> Option<R> {
        self.sources.lock().ok().map(|mut sources| f(&mut sources))
    }

    /// Trace-start `CLOCK_MONOTONIC` timestamp.
    pub(crate) fn start_time_ns(&self) -> u64 {
        self.start_time_ns
    }

    crate::test_util_pub! {
        /// Turn recording on. No-op once stopped.
        fn enable(&self) {
            let _ = self.state.compare_exchange(
                State::Disabled as u8,
                State::Enabled as u8,
                Ordering::Relaxed,
                Ordering::Relaxed,
            );
        }
    }

    /// Turn recording off. No-op once stopped.
    pub(crate) fn disable(&self) {
        let _ = self.state.compare_exchange(
            State::Enabled as u8,
            State::Disabled as u8,
            Ordering::Relaxed,
            Ordering::Relaxed,
        );
    }

    /// Stop for good: recording off, [`enable`](Self::enable) and new attaches
    /// refused. Called at shutdown once the flush thread is gone.
    pub(crate) fn mark_stopped(&self) {
        self.state.store(State::Stopped as u8, Ordering::Relaxed);
    }

    /// Whether the recorder has shut down for good.
    pub(crate) fn is_stopped(&self) -> bool {
        self.state.load(Ordering::Relaxed) == State::Stopped as u8
    }

    /// Install the on-demand dump trigger. Set once at build time by the
    /// facade builder; later calls are ignored. `pub` so the facade (a
    /// sibling crate) can wire the trigger in.
    #[cfg(feature = "pipeline")]
    pub(crate) fn set_dump_trigger(&self, trigger: crate::dump::DumpTrigger) {
        let _ = self.dump_trigger.set(trigger);
    }

    /// The on-demand dump trigger, or `None` when the runtime was built
    /// without `with_dump_trigger`.
    #[cfg(feature = "pipeline")]
    pub(crate) fn dump_trigger(&self) -> Option<&crate::dump::DumpTrigger> {
        self.dump_trigger.get()
    }

    /// Check whether recording is currently enabled.
    ///
    /// For recording paths prefer
    /// [`Dial9Handle::record_event_with`](crate::handle::Dial9Handle::record_event_with),
    /// which builds the event inside the check. Use `is_enabled()` for
    /// control-flow decisions that don't record, such as whether to wrap a
    /// waker in wake-tracking polls.
    pub(crate) fn is_enabled(&self) -> bool {
        self.state.load(Ordering::Relaxed) == State::Enabled as u8
    }

    /// Run `f` only when recording is enabled, passing an [`EventBuffer`]
    /// that provides `record_event` / `record_encodable_event`. Returns
    /// `None` when disabled (no work is done).
    pub(crate) fn if_enabled<R>(&self, f: impl FnOnce(&EventBuffer<'_>) -> R) -> Option<R> {
        if !self.is_enabled() {
            return None;
        }
        Some(f(&EventBuffer(self)))
    }

    /// Test-only shortcut to record an event directly. Production code records
    /// through [`EventBuffer`] via [`if_enabled`](Self::if_enabled).
    #[cfg(test)]
    fn record_encodable_event(&self, event: &dyn encoder::Encodable) {
        if let Some(handle) =
            encoder::record_encodable_event(event, &self.collector, &self.drain_epoch)
        {
            self.tl_buffers.lock().unwrap().push(handle);
        }
    }

    /// Bump the drain epoch and flush all idle/silent thread-local buffers.
    ///
    /// Buffers whose `FlushEpoch` matches the current epoch are skipped
    /// (the owning thread flushed recently, so locking would just add
    /// contention). Dead `Weak` handles are pruned.
    ///
    /// [`bump_drain_epoch`] is called one flush-loop tick
    /// before calling this method. That gives busy worker threads a ~5 ms
    /// grace period to self-flush on their next `record_event`, so the
    /// intrusive drain only needs to lock truly idle/silent buffers.
    ///
    /// Returns per-cycle counters so the flush thread can emit metrics.
    pub(crate) fn drain_all_tl_buffers(&self) -> TlDrainStats {
        let mut stats = TlDrainStats::default();
        let epoch = self.drain_epoch.load(Ordering::Relaxed);

        let handles: Vec<TlBufferHandle> = {
            let guard = self.tl_buffers.lock().unwrap();
            guard
                .iter()
                .map(|h| TlBufferHandle {
                    buffer: h.buffer.clone(),
                    flush_epoch: h.flush_epoch.clone(),
                })
                .collect()
        };

        for handle in &handles {
            // Skip buffers that self-flushed during the current epoch.
            if handle.flush_epoch.load() >= epoch {
                stats.buffers_skipped_busy += 1;
                continue;
            }
            if let Some(arc) = handle.buffer.upgrade() {
                let mut buf = match arc.lock() {
                    Ok(guard) => guard,
                    // Buffer is poisoned (encoder panic); skip rather than
                    // flushing potentially corrupt data.
                    Err(_) => {
                        crate::rate_limit::rate_limited!(Duration::from_secs(60), {
                            tracing::error!(
                                "dial9: thread-local buffer mutex poisoned in drain_all_tl_buffers; skipping flush"
                            );
                        });
                        continue;
                    }
                };
                stats.buffers_locked += 1;
                if buf.has_pending_events() {
                    let batch = buf.flush();
                    stats.events_flushed += batch.event_count();
                    stats.buffers_flushed += 1;
                    self.collector.accept_flush(batch);
                }
                // Stamp so we skip this buffer next cycle if it stays idle.
                handle.flush_epoch.store(epoch);
            }
        }

        // Prune dead handles (Weak refs to threads that have exited).
        let mut guard = self.tl_buffers.lock().unwrap();
        let before = guard.len();
        guard.retain(|h| h.buffer.strong_count() > 0);
        stats.dead_pruned = (before - guard.len()) as u64;

        stats
    }

    /// Advance the global drain epoch so that busy worker threads
    /// self-flush on their next `record_event` call. Call this one
    /// flush-loop tick (~5 ms) before [`drain_all_tl_buffers`] to give
    /// workers a grace period, minimising contention on the intrusive
    /// drain path.
    pub(crate) fn bump_drain_epoch(&self) {
        self.drain_epoch.fetch_add(1, Ordering::Relaxed);
    }

    crate::test_util_pub! {
        /// Drain data sources and write their events into the collector.
        fn flush_sources(&self) {
            let ctx = self.flush_context();
            let mut sources = self.sources.lock().unwrap();
            for source in sources.iter_mut() {
                source.flush(&ctx);
            }
        }
    }

    crate::test_util_pub! {
        /// Build a [`FlushContext`] for this state.
        ///
        /// Used by `flush_sources` and by tests that construct a ctx directly.
        ///
        /// [`FlushContext`]: crate::source::FlushContext
        fn flush_context(&self) -> crate::source::FlushContext<'_> {
            crate::source::FlushContext::new(&self.collector, &self.drain_epoch)
        }
    }
}

/// Handle provided by [`SharedState::if_enabled`] that proves recording is
/// active. All event-recording calls should go through this type so that
/// callers cannot accidentally emit events without an enabled check.
#[derive(Debug)]
pub(crate) struct EventBuffer<'a>(&'a SharedState);

impl EventBuffer<'_> {
    pub(crate) fn record_encodable_event(&self, event: &dyn encoder::Encodable) {
        if let Some(handle) =
            encoder::record_encodable_event(event, &self.0.collector, &self.0.drain_epoch)
        {
            self.0.tl_buffers.lock().unwrap().push(handle);
        }
    }

    pub(crate) fn with_encoder(&self, f: impl FnOnce(&mut encoder::ThreadLocalEncoder<'_>)) {
        if let Some(handle) = encoder::with_encoder(f, &self.0.collector, &self.0.drain_epoch) {
            self.0.tl_buffers.lock().unwrap().push(handle);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::primitives::sync::atomic::AtomicBool;

    fn sample_event() -> crate::format::ClockSyncEvent {
        crate::format::ClockSyncEvent {
            timestamp_ns: 1000,
            realtime_ns: 2000,
        }
    }

    /// Helper: create a SharedState with recording enabled.
    fn enabled_shared_state() -> SharedState {
        let ss = SharedState::new(0);
        ss.enable();
        ss
    }

    #[test]
    fn record_event_registers_tl_buffer_handle() {
        let ss = enabled_shared_state();
        // First event on this thread should register a handle.
        ss.record_encodable_event(&sample_event());
        let handles = ss.tl_buffers.lock().unwrap();
        assert_eq!(handles.len(), 1);
        assert!(handles[0].buffer.upgrade().is_some());
    }

    #[test]
    fn second_record_event_does_not_re_register() {
        let ss = enabled_shared_state();
        ss.record_encodable_event(&sample_event());
        ss.record_encodable_event(&sample_event());
        let handles = ss.tl_buffers.lock().unwrap();
        assert_eq!(handles.len(), 1);
    }

    #[test]
    fn drain_all_tl_buffers_flushes_idle_buffer() {
        let ss = enabled_shared_state();
        // Write an event (won't self-flush — buffer is 1MB).
        ss.record_encodable_event(&sample_event());
        // Nothing in the collector yet (buffer not full).
        assert!(ss.collector.next().is_none());
        // Bump epoch so the idle buffer (epoch 0) is stale, then drain.
        ss.bump_drain_epoch();
        ss.drain_all_tl_buffers();
        let batch = ss.collector.next().expect("expected a batch after drain");
        assert!(batch.event_count() > 0);
    }

    #[test]
    fn drain_all_tl_buffers_from_another_thread() {
        let ss = Arc::new(enabled_shared_state());
        let ss2 = ss.clone();
        // Write events from a spawned thread.
        let handle = std::thread::spawn(move || {
            ss2.record_encodable_event(&sample_event());
            ss2.record_encodable_event(&sample_event());
        });
        handle.join().unwrap();
        // Bump epoch so the buffer is stale, then drain from the main thread.
        ss.bump_drain_epoch();
        ss.drain_all_tl_buffers();
        let batch = ss.collector.next().expect("expected a batch after drain");
        assert_eq!(batch.event_count(), 2);
    }

    #[test]
    fn drain_skips_busy_buffer() {
        let ss = enabled_shared_state();
        ss.record_encodable_event(&sample_event());
        // Bump epoch to 1 (simulates the tick before the drain).
        ss.bump_drain_epoch();
        // Simulate a self-flush by stamping the current epoch.
        {
            let handles = ss.tl_buffers.lock().unwrap();
            handles[0].flush_epoch.store(1);
        }
        ss.drain_all_tl_buffers();
        // Buffer should NOT have been flushed — collector is empty.
        assert!(ss.collector.next().is_none());
    }

    #[test]
    fn drain_prunes_dead_handles() {
        let ss = Arc::new(enabled_shared_state());
        let ss2 = ss.clone();
        let handle = std::thread::spawn(move || {
            ss2.record_encodable_event(&sample_event());
        });
        handle.join().unwrap();
        // Thread exited — its Arc<Mutex<TLB>> was dropped, Weak is dead.
        // But the TLB's Drop impl flushed remaining events, so the handle
        // is dead. Drain should prune it.
        ss.drain_all_tl_buffers();
        let handles = ss.tl_buffers.lock().unwrap();
        assert_eq!(handles.len(), 0, "dead handle should have been pruned");
    }

    /// Intrusive-drain path with a *live* worker thread. Unlike
    /// `drain_all_tl_buffers_from_another_thread`, which joins the worker
    /// before draining (so events reach the collector via the TLB `Drop`
    /// impl, not via the intrusive path), here the worker is parked on a
    /// channel while the main thread bumps+drains, proving that
    /// `drain_all_tl_buffers` upgrades the live `Weak`, locks the mutex
    /// cross-thread, and flushes the pending event.
    #[test]
    fn drain_flushes_live_worker_buffer() {
        let ss = Arc::new(enabled_shared_state());
        let ss2 = ss.clone();
        let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
        let (ready_tx, ready_rx) = std::sync::mpsc::channel::<()>();

        let worker = std::thread::spawn(move || {
            // drain_epoch is 0, so no self-flush happens — the event
            // stays in the buffer.
            ss2.record_encodable_event(&sample_event());
            ready_tx.send(()).unwrap();
            // Park until main thread has drained. The TLB `Drop` impl must
            // not run before the intrusive drain, otherwise we're not
            // testing the intrusive path.
            release_rx.recv().unwrap();
        });

        ready_rx.recv().unwrap();
        // Worker is parked with one event in its TLB and a live handle.
        // Nothing in the collector yet — no self-flush was triggered.
        assert!(ss.collector.next().is_none());

        ss.bump_drain_epoch();
        ss.drain_all_tl_buffers();

        let batch = ss
            .collector
            .next()
            .expect("intrusive drain should have flushed the live worker's event");
        assert_eq!(batch.event_count(), 1);

        release_tx.send(()).unwrap();
        worker.join().unwrap();
    }

    // Concurrent-stress proptest: the core invariant of the TL buffer
    // drain feature is that no events are lost and none are duplicated,
    // regardless of how `record_encodable_event`, `bump_drain_epoch`, and
    // `drain_all_tl_buffers` interleave across threads. Spawn N writer
    // threads, each recording M events, while a drainer thread
    // concurrently bumps+drains. After joining, a final bump+drain should
    // leave exactly N*M events in the collector.
    proptest::proptest! {
        #![proptest_config(proptest::prelude::ProptestConfig::with_cases(32))]

        #[test]
        fn concurrent_record_and_drain_preserves_event_count(
            num_threads in 1usize..=6,
            events_per_thread in 1u64..=200,
            drain_ticks in 0usize..=10,
        ) {
            let ss = Arc::new(enabled_shared_state());
            let start = Arc::new(std::sync::Barrier::new(num_threads + 1));
            let stop_drainer = Arc::new(AtomicBool::new(false));

            let writers: Vec<_> = (0..num_threads)
                .map(|_| {
                    let ss = ss.clone();
                    let start = start.clone();
                    std::thread::spawn(move || {
                        start.wait();
                        for _ in 0..events_per_thread {
                            ss.record_encodable_event(&sample_event());
                        }
                    })
                })
                .collect();

            let drainer = {
                let ss = ss.clone();
                let stop = stop_drainer.clone();
                std::thread::spawn(move || {
                    let mut ticks = 0;
                    while ticks < drain_ticks && !stop.load(Ordering::Relaxed) {
                        ss.bump_drain_epoch();
                        // Short grace period so any in-flight writer has a
                        // chance to self-flush before the intrusive drain.
                        std::thread::sleep(std::time::Duration::from_micros(50));
                        ss.drain_all_tl_buffers();
                        ticks += 1;
                    }
                })
            };

            start.wait();
            for w in writers {
                w.join().unwrap();
            }
            stop_drainer.store(true, Ordering::Relaxed);
            drainer.join().unwrap();

            // Writer threads have exited, so their TLB `Drop` impls have
            // flushed any remaining events. Do one final bump+drain to
            // prune dead handles (no-op for event capture at this point).
            ss.bump_drain_epoch();
            ss.drain_all_tl_buffers();

            let mut total: u64 = 0;
            while let Some(batch) = ss.collector.next() {
                total += batch.event_count();
            }
            // Sanity: the collector never evicted a batch under these
            // workloads. If it did, the invariant check below would be
            // meaningless.
            proptest::prop_assert_eq!(ss.collector.take_dropped_batches(), 0);
            proptest::prop_assert_eq!(
                total,
                num_threads as u64 * events_per_thread,
                "every recorded event must reach the collector exactly once"
            );
        }
    }

    // Every scenario below shares a caveat: `drain_epoch`/`flush_epoch`/`state`
    // are `Ordering::Relaxed`, but shuttle 0.9.1 treats every ordering as
    // `SeqCst`. A `Relaxed`-specific bug is invisible here no matter how
    // exhaustively a scenario searches.
    #[cfg(shuttle)]
    mod shuttle_tests {
        use super::*;
        use crate::buffer::MemoryBuffer;
        use crate::recording::Recorder;

    // Small values kept: shuttle's search cost grows with interleaving
    // space, unlike a real-thread proptest.
    const WRITERS: usize = 3;
    const EVENTS_PER_WRITER: u64 = 3;

    const LIFECYCLE_WRITERS: usize = 2;

    // Shared with `shuttle_writer_exit_races_drain`: both drive a
    // `Recorder`/flush thread and assert only on decoded, sealed output.

    /// Round-trip event recorded through [`Dial9Handle`](crate::handle::Dial9Handle).
    #[derive(dial9_trace_format::TraceEvent, Clone, Debug, serde::Deserialize)]
    struct LifecycleEvent {
        #[traceevent(timestamp)]
        timestamp_ns: u64,
        id: u64,
    }

    fn decode_lifecycle_events(data: &[u8]) -> Vec<LifecycleEvent> {
        use dial9_trace_format::decoder::Decoder;
        let Some(mut dec) = Decoder::new(data) else {
            assert!(data.is_empty(), "failed to decode a non-empty segment");
            return vec![];
        };
        let mut out = Vec::new();
        dec.for_each_event(|ev| {
            if ev.name == "LifecycleEvent"
                && let Ok(decoded) = ev.deserialize::<LifecycleEvent>()
            {
                out.push(decoded);
            }
        })
        .expect("decode failed");
        out
    }

    /// Drain every sealed segment the in-memory writer has produced so far.
    fn decode_all_lifecycle_events(fs: &crate::fs::Fs) -> Vec<LifecycleEvent> {
        let mut decoded = Vec::new();
        loop {
            let taken = fs.take_files();
            if taken.segments.is_empty() {
                break;
            }
            for seg in taken.segments {
                let (_seg_ref, payload, _accounting) = seg.load().unwrap();
                decoded.extend(decode_lifecycle_events(&payload.into_vec()));
            }
        }
        decoded
    }

    crate::shuttle_test! {
        default;
        // `disable`/`enable` raced against recording, through the
        // `Recorder`/`Dial9Handle`/flush-thread stack. `record_event`
        // doesn't report whether a call actually landed, so this checks
        // decoded output directly: no duplicate ids, no un-attempted
        // ids, and `Stopped` stays terminal under a raced `enable`/
        // `disable`. Also subject to the `Ordering::Relaxed` caveat above.
        fn shuttle_disable_races_record() {
            let _ts_guard = metrique_timesource::set_time_source(
                metrique_timesource::TimeSource::custom(
                    metrique_timesource::fakes::StaticTimeSource::at_time(std::time::UNIX_EPOCH),
                ),
            );

            // Small segments so the writer/drain race actually forces
            // rotation; the total budget is far above this test's data so
            // the ring never evicts before we drain it.
            let writer = MemoryBuffer::builder()
                .max_total_size(100 * 1024 * 1024)
                .max_segment_size(256)
                .build()
                .unwrap();
            let fs = writer.fs_handle().expect("in-memory writer exposes its fs");

            let shared = Arc::new(SharedState::new(0));
            let mut recorder = Recorder::start(shared, writer, None, || || {});
            recorder.handle().enable();
            let handle = recorder.handle().clone();

            let next_id = Arc::new(AtomicU64::new(0));
            let attempted: Arc<Mutex<Vec<u64>>> = Arc::new(Mutex::new(Vec::new()));

            let writers: Vec<_> = (0..LIFECYCLE_WRITERS)
                .map(|_| {
                    let handle = handle.clone();
                    let next_id = next_id.clone();
                    let attempted = attempted.clone();
                    crate::primitives::thread::spawn(move || {
                        for _ in 0..EVENTS_PER_WRITER {
                            let id = next_id.fetch_add(1, Ordering::Relaxed);
                            attempted.lock().unwrap().push(id);
                            handle.record_event(LifecycleEvent {
                                timestamp_ns: id,
                                id,
                            });
                        }
                    })
                })
                .collect();

            let lifecycle = {
                let handle = handle.clone();
                crate::primitives::thread::spawn(move || {
                    handle.disable(); // Enabled -> Disabled
                    handle.enable(); // Disabled -> Enabled (re-arm)
                    handle.disable();
                })
            };

            for w in writers {
                w.join().unwrap();
            }
            lifecycle.join().unwrap();

            // Flushes, seals the final segment, joins the flush thread,
            // then calls `SharedState::mark_stopped`.
            recorder.stop_flush_thread();

            // `Stopped` is terminal regardless of schedule: a `disable()` or
            // `enable()` that lost the race must not have clobbered it back
            // to a live state. Catches a CAS -> plain `store` regression in
            // either.
            assert!(handle.is_stopped(), "Stopped must be terminal");
            handle.enable();
            assert!(!handle.is_enabled(), "enable() must be a no-op once stopped");
            handle.disable();
            assert!(
                handle.is_stopped(),
                "disable() must not un-stop a terminal state"
            );

            let decoded = decode_all_lifecycle_events(&fs);
            let attempted = attempted.lock().unwrap();
            let attempted_ids: std::collections::HashSet<u64> =
                attempted.iter().copied().collect();

            let mut decoded_ids: Vec<u64> = decoded.iter().map(|e| e.id).collect();
            for id in &decoded_ids {
                assert!(
                    attempted_ids.contains(id),
                    "decoded event id {id} was never attempted"
                );
            }
            decoded_ids.sort_unstable();
            let mut deduped = decoded_ids.clone();
            deduped.dedup();
            assert_eq!(
                deduped.len(),
                decoded_ids.len(),
                "an event id was recorded more than once: {decoded_ids:?}"
            );
            assert!(
                decoded_ids.len() <= attempted.len(),
                "more events landed ({}) than were ever attempted ({})",
                decoded_ids.len(),
                attempted.len()
            );
        }
    }

    crate::shuttle_test! {
        default;
        // Writer threads emit one event through `Dial9Handle` and exit
        // immediately, racing `run_flush_loop`'s own ~5ms polling +
        // grace-period cadence. Each writer holds exactly one event, so
        // a plain id-presence check works regardless of which path
        // flushed it. Also subject to the `Ordering::Relaxed` caveat above.
        fn shuttle_writer_exit_races_drain() {
            let _ts_guard = metrique_timesource::set_time_source(
                metrique_timesource::TimeSource::custom(
                    metrique_timesource::fakes::StaticTimeSource::at_time(std::time::UNIX_EPOCH),
                ),
            );

            let writer = MemoryBuffer::builder()
                .max_total_size(100 * 1024 * 1024)
                .max_segment_size(256)
                .build()
                .unwrap();
            let fs = writer.fs_handle().expect("in-memory writer exposes its fs");

            let shared = Arc::new(SharedState::new(0));
            let mut recorder = Recorder::start(shared, writer, None, || || {});
            recorder.handle().enable();
            let handle = recorder.handle().clone();

            let writers: Vec<_> = (0..WRITERS)
                .map(|id| {
                    let handle = handle.clone();
                    crate::primitives::thread::spawn(move || {
                        handle.record_event(LifecycleEvent {
                            timestamp_ns: id as u64,
                            id: id as u64,
                        });
                        // Return immediately: the next scheduling points are
                        // the thread-local destructor, the final `Arc`
                        // decrement, and the `Mutex` drop.
                    })
                })
                .collect();

            for w in writers {
                w.join().unwrap();
            }

            // Flushes, seals the final segment, joins the flush thread.
            recorder.stop_flush_thread();

            let decoded = decode_all_lifecycle_events(&fs);
            let mut ids: Vec<u64> = decoded.iter().map(|e| e.id).collect();
            ids.sort_unstable();
            assert_eq!(
                ids,
                (0..WRITERS as u64).collect::<Vec<_>>(),
                "each writer's event must arrive exactly once, whether flushed by \
                 the flush loop's intrusive drain or by the TLB Drop impl on thread exit"
            );
        }
    }
    }
}