smarm/src/trace.rs

//! Structured per-event tracing for smarm.
//!
//! Enabled by `--features smarm-trace`. Zero cost without the feature.
//!
//! Architecture: MPSC. Every scheduler thread holds a thread-local Sender
//! clone (one mutex acquire per thread, on first use). A dedicated drain
//! thread owns the Receiver, batches records, and writes to a BufWriter.
//! The hot path (record()) is a single channel send — no mutex, no disk I/O.
//!
//! Usage:
//!   cargo test --test runtime <test_name> --features smarm-trace
//!
//! Output: smarm_trace.json in cwd, or $SMARM_TRACE_FILE.
//! View:   https://ui.perfetto.dev  or  chrome://tracing

#[cfg(feature = "smarm-trace")]
#[macro_export]
macro_rules! te {
    ($kind:expr) => { $crate::trace::record($kind) };
}

#[cfg(not(feature = "smarm-trace"))]
#[macro_export]
macro_rules! te {
    ($kind:expr) => { () };
}

#[cfg(feature = "smarm-trace")]
pub use inner::*;

#[cfg(feature = "smarm-trace")]
mod inner {
    use crate::pid::Pid;
    use std::io::Write;
    use std::sync::{mpsc, Mutex};
    use std::time::Instant;

    // -----------------------------------------------------------------------
    // Event kinds
    // -----------------------------------------------------------------------

    #[derive(Clone, Debug)]
    pub enum Event {
        // Actor lifecycle
        Spawn { parent: Pid, child: Pid },
        Resume(Pid),
        Yield(Pid),
        Park(Pid),
        Done(Pid),
        // Wakeup paths
        UnparkDirect(Pid),       // unpark() saw Parked   -> re-queued immediately
        UnparkDeferred(Pid),     // unpark() saw Runnable -> set pending_unpark flag
        UnparkFlagConsumed(Pid), // scheduler saw flag on Park -> re-queued instead
        // Channel
        Send { sender: Pid, receiver: Option<Pid> },
        RecvPark(Pid),
        RecvWake(Pid),
        // Queue
        Enqueue(Pid),
        Dequeue(Pid),
        // RFC 005 wake slot
        SlotPush(Pid), // actor-context wake parked in the waking thread's slot
        SlotPop(Pid),  // scheduler resumed a pid from its own slot
    }

    // -----------------------------------------------------------------------
    // Wire format sent through the channel
    // -----------------------------------------------------------------------

    struct Record {
        nanos: u64,   // ns since open()
        tid:   u64,   // OS thread id
        event: Event,
    }

    // Sentinel: drain thread flushes and exits when it receives this.
    enum Msg {
        Event(Record),
        Flush,
    }

    // -----------------------------------------------------------------------
    // Global sender + start time
    // -----------------------------------------------------------------------

    struct Global {
        sender:  mpsc::Sender<Msg>,
        start:   Instant,
    }

    static GLOBAL: Mutex<Option<Global>> = Mutex::new(None);

    // Per-thread state: cached Sender clone + cached copy of start Instant.
    // The Sender clone is taken once per thread (one mutex hit).
    // The start Instant is copied alongside it — also one mutex hit per thread.
    // record() never touches GLOBAL after that.
    struct LocalState {
        tx:    mpsc::Sender<Msg>,
        start: Instant,
    }

    thread_local! {
        static LOCAL_STATE: std::cell::RefCell<Option<LocalState>> =
            std::cell::RefCell::new(None);
    }

    // -----------------------------------------------------------------------
    // Lifecycle
    // -----------------------------------------------------------------------

    pub fn open() {
        let path = std::env::var("SMARM_TRACE_FILE")
            .unwrap_or_else(|_| "smarm_trace.json".to_owned());

        let (tx, rx) = mpsc::channel::<Msg>();
        let start = Instant::now();

        *GLOBAL.lock().unwrap() = Some(Global { sender: tx, start });

        // Drain thread: owns the Receiver, writes to disk.
        let path_for_thread = path.clone();
        std::thread::Builder::new()
            .name("smarm-trace-drain".into())
            .spawn(move || drain_thread(rx, &path_for_thread))
            .expect("failed to spawn trace drain thread");

        eprintln!("[smarm-trace] writing to {}", path);
    }

    /// Send a Flush sentinel and block until the drain thread finishes writing.
    /// Called by Runtime::run after all scheduler threads have exited.
    pub fn flush() {
        // Drop the global sender so the drain thread's recv() returns Err
        // after the Flush sentinel, signalling clean shutdown.
        let sender = {
            let mut g = GLOBAL.lock().unwrap();
            g.take().map(|g| g.sender)
        };
        if let Some(tx) = sender {
            let _ = tx.send(Msg::Flush);
            // tx drops here — drain thread will see disconnected after Flush.
        }
        // Clear thread-local state.
        LOCAL_STATE.with(|c| *c.borrow_mut() = None);
    }

    // -----------------------------------------------------------------------
    // Hot path
    // -----------------------------------------------------------------------

    pub fn record(event: Event) {
        // Disable preemption for the entire duration of record(). Any
        // allocation here (mutex internals, channel send, lazy init) would
        // trigger PreemptingAllocator -> maybe_preempt -> switch_to_scheduler,
        // which would try to re-acquire inner.shared (already held at many
        // te!() call sites) -> deadlock. Guard at the very top, before any
        // allocation-capable call.
        let was_enabled = crate::preempt::PREEMPTION_ENABLED
            .with(|e| { let v = e.get(); e.set(false); v });

        LOCAL_STATE.with(|cell| {
            let mut opt = cell.borrow_mut();
            // Lazily initialise: one mutex hit per thread, ever.
            if opt.is_none() {
                if let Some(g) = GLOBAL.lock().unwrap().as_ref() {
                    let tx = g.sender.clone();
                    *opt = Some(LocalState { tx, start: g.start });
                }
            }
            if let Some(ls) = opt.as_ref() {
                let nanos = ls.start.elapsed().as_nanos() as u64;
                let tid   = os_tid();
                let _ = ls.tx.send(Msg::Event(Record { nanos, tid, event }));
            }
        });

        crate::preempt::PREEMPTION_ENABLED.with(|e| e.set(was_enabled));
    }

    // -----------------------------------------------------------------------
    // Drain thread
    // -----------------------------------------------------------------------

    fn drain_thread(rx: mpsc::Receiver<Msg>, path: &str) {
        let f = match std::fs::File::create(path) {
            Ok(f) => f,
            Err(e) => { eprintln!("[smarm-trace] create failed: {}", e); return; }
        };
        let mut w = std::io::BufWriter::new(f);
        let _ = writeln!(w, "{{\"traceEvents\":[");

        let mut count: u64 = 0;
        let mut first = true;

        loop {
            match rx.recv() {
                Ok(Msg::Event(r)) => {
                    let (name, actor_idx) = chrome_fields(&r.event);
                    let ts_us = r.nanos as f64 / 1000.0;
                    if !first { let _ = w.write_all(b",\n"); }
                    first = false;
                    let _ = write!(w,
                        "{{\"ph\":\"i\",\"ts\":{:.3},\"pid\":{},\"tid\":{},\"name\":{:?},\"s\":\"g\"}}",
                        ts_us, actor_idx, r.tid, name);
                    count += 1;
                }
                Ok(Msg::Flush) | Err(_) => {
                    // Clean close.
                    let _ = writeln!(w, "\n]}}");
                    let _ = w.flush();
                    eprintln!("[smarm-trace] {} events written", count);
                    return;
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Chrome trace helpers
    // -----------------------------------------------------------------------

    fn chrome_fields(ev: &Event) -> (String, u32) {
        match ev {
            Event::Spawn { parent, child } =>
                (format!("spawn c={}", child.index()), parent.index()),
            Event::Resume(p)             => ("resume".into(),               p.index()),
            Event::Yield(p)              => ("yield".into(),                p.index()),
            Event::Park(p)               => ("park".into(),                 p.index()),
            Event::Done(p)               => ("done".into(),                 p.index()),
            Event::UnparkDirect(p)       => ("unpark_direct".into(),        p.index()),
            Event::UnparkDeferred(p)     => ("unpark_deferred".into(),      p.index()),
            Event::UnparkFlagConsumed(p) => ("unpark_flag_consumed".into(), p.index()),
            Event::Send { sender, receiver } => (
                format!("send rx={}", receiver
                    .map(|p| p.index().to_string())
                    .unwrap_or_else(|| "none".into())),
                sender.index(),
            ),
            Event::RecvPark(p) => ("recv_park".into(), p.index()),
            Event::RecvWake(p) => ("recv_wake".into(), p.index()),
            Event::Enqueue(p)  => ("enqueue".into(),   p.index()),
            Event::Dequeue(p)  => ("dequeue".into(),   p.index()),
            Event::SlotPush(p) => ("slot_push".into(), p.index()),
            Event::SlotPop(p)  => ("slot_pop".into(),  p.index()),
        }
    }

    fn os_tid() -> u64 {
        unsafe { libc::syscall(libc::SYS_gettid) as u64 }
    }
}