//! Host-side client for the scheduler-stats virtio-console bridge.
//!
//! The scx_stats userspace protocol (line-delimited JSON over a Unix
//! socket) is forwarded between host and guest over
//! [`super::virtio_console::VirtioConsole`] port 2:
//!
//! 1. Caller invokes [`SchedStatsClient::request_raw`] (or one of the
//!    typed wrappers [`SchedStatsClient::stats`] /
//!    [`SchedStatsClient::stats_meta`]).
//! 2. The request bytes (a complete `\n`-terminated JSON line) are
//!    pushed onto port 2 RX via
//!    [`super::virtio_console::VirtioConsole::queue_input_port2`].
//! 3. The guest's relay thread reads `/dev/vport0p2` and forwards
//!    the bytes to `/var/run/scx/root/stats`.
//! 4. The scheduler's response — also a `\n`-terminated JSON line —
//!    travels back through the same path: scheduler writes the Unix
//!    socket; relay forwards it to `/dev/vport0p2`; the device
//!    accumulates the bytes in `ports[2].tx_buf` and signals
//!    `stats_tx_evt`.
//! 5. The drainer thread this client owns wakes on `stats_tx_evt`,
//!    drains `ports[2].tx_buf`, and either appends the bytes to the
//!    response buffer (if a request is in flight) or discards them
//!    (logging the discard count).
//! 6. [`SchedStatsClient::request_raw`] blocks on a [`Condvar`]
//!    until a complete `\n`-terminated line accumulates or the
//!    timeout elapses.
//!
//! No extra framing: scx_stats is already line-delimited JSON, so a
//! TLV layer here would just double-frame.
//!
//! # Threading model
//!
//! Each client owns a dedicated drainer thread that polls
//! `stats_tx_evt` via epoll. Cloning the client is cheap (everything
//! is `Arc`-shared) and does NOT spawn additional threads — the
//! single drainer feeds every clone's response buffer because they
//! all share the same `Arc<(Mutex<Vec<u8>>, Condvar)>`.
//!
//! [`SchedStatsClient`] IS safe for concurrent use. scx_stats has
//! no request-id multiplexing, so concurrent host calls would race
//! for the response stream; the client wraps the entire
//! request/response cycle in an internal `Mutex<()>` so callers
//! that issue overlapping requests SERIALISE through the lock
//! rather than corrupting each other's responses. The drainer
//! reads a separate `request_in_flight: AtomicBool` (set inside the
//! lock, cleared on exit) to decide whether to forward drained
//! bytes to the response buffer or discard them — see [`drainer_loop`].
//!
//! # Lifetime
//!
//! Drop on the last clone tears down the drainer thread by writing
//! the `kill_drainer` eventfd. The drainer wakes from `epoll_wait`,
//! observes the kill flag, and exits.

use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex};
// `Duration` is only referenced by the cfg(test) module; `Instant`
// is unused after the move to event-driven (no timeouts) cvar
// waits — keep the production code free of `std::time` so a future
// regression that re-introduces a timeout does not silently
// compile.

use crate::sync::MutexExt;
use serde::{Deserialize, Serialize};
use vmm_sys_util::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
use vmm_sys_util::eventfd::{EFD_NONBLOCK, EventFd};

use super::PiMutex;
use super::virtio_console::VirtioConsole;

/// Synthetic relay-error envelope the drainer injects when the
/// 2x hard cap on the response accumulator is breached. Without
/// this, a runaway guest that bursts > 512 KiB without a newline
/// would silently get its bytes dropped while the host request
/// hung in `cvar.wait` forever. The drainer instead clears the
/// buffer, writes this synthetic envelope, and notifies the cvar
/// so `request_raw` wakes up, parses the envelope, and surfaces
/// [`SchedStatsError::NoScheduler`] with a `cap-overflow` reason.
/// The trailing `\n` matches scx_stats's line-delimited wire
/// format so the existing newline-detection path consumes it
/// without special-casing.
const CAP_OVERFLOW_ERROR_REPLY: &[u8] = b"{\"ktstr_relay_error\":\"response cap overflow\"}\n";

/// Maximum bytes accepted in a single host→guest stats request.
/// scx_stats requests are JSON command lines (`{"req":"stats"}\n`
/// and similar); 256 KiB is far above any legitimate request size.
/// Caps the per-request `queue_input_port2` allocation so a buggy
/// caller can't grow the device's `ports[2].pending_rx` deque without
/// bound on a single push.
pub const MAX_REQUEST_BYTES: usize = 256 * 1024;

/// Maximum bytes the response accumulator may grow to before
/// [`SchedStatsClient::request_raw`] returns
/// [`SchedStatsError::ResponseTooLarge`]. scx_stats responses are
/// JSON objects on a single line; 256 KiB matches the request cap
/// and protects against a hostile guest (or a runaway relay) that
/// sends bytes without ever emitting a newline.
pub const MAX_RESPONSE_BYTES: usize = 256 * 1024;

/// Errors surfaced by [`SchedStatsClient::request_raw`]. Distinct
/// from `anyhow::Error` so callers can branch on the failure mode
/// without string-matching error chains.
#[derive(Debug)]
pub enum SchedStatsError {
    /// The shared response mutex was poisoned by a previous panic.
    /// The client cannot recover; callers should propagate the
    /// error.
    Poisoned,
    /// The caller passed a request larger than [`MAX_REQUEST_BYTES`].
    RequestTooLarge {
        /// Size the caller passed, in bytes.
        size: usize,
        /// Cap configured on the client (`MAX_REQUEST_BYTES`).
        max: usize,
    },
    /// The response accumulator grew past [`MAX_RESPONSE_BYTES`]
    /// without ever emitting a newline. The buffer is cleared
    /// before the error returns so the next request starts from a
    /// clean slate; the partial bytes that triggered the cap are
    /// discarded.
    ResponseTooLarge {
        /// Accumulator size at the moment the cap was breached.
        size: usize,
        /// Cap configured on the client (`MAX_RESPONSE_BYTES`).
        max: usize,
    },
    /// The host-side coordinator marked the run as freezing while
    /// this request was in flight (or about to start). scx_stats
    /// responses are undefined while the scheduler's userspace
    /// thread is paused.
    DuringFreeze,
    /// The run-wide cancel flag was set (the watchdog fired or
    /// the run is shutting down) while this request was in flight
    /// or about to start. The host-side watchdog is the only
    /// "timeout" in the stats path — when it fires, every
    /// outstanding `request_raw` returns immediately with this
    /// variant rather than blocking forever.
    Cancelled,
    /// The guest relay never connected to the scheduler's Unix
    /// socket (no scheduler running, or the scheduler refused the
    /// connection). Delivered as an inline JSON error response by
    /// the in-guest relay; surfaced here so callers can branch on
    /// "no scheduler" without parsing the error JSON.
    NoScheduler {
        /// Reason buffer carried in the inline error response.
        reason: String,
    },
    /// The scheduler returned a non-zero `errno` in the typed
    /// [`StatsResponse`] envelope. The wire payload's verb-specific
    /// `args` field is preserved so the caller can render an
    /// actionable diagnostic.
    SchedulerError {
        /// `errno` field from the scx_stats response envelope.
        errno: i32,
        /// `args` payload from the response. Typically a JSON
        /// object with an error message; opaque to this client.
        args: serde_json::Value,
    },
    /// The typed envelope was successfully decoded but the inner
    /// `args` map did not contain the expected `"resp"` key.
    /// scx_stats schedulers wrap their verb-specific payload under
    /// this key; its absence indicates either a protocol mismatch
    /// or a non-conforming scheduler implementation.
    MissingResp {
        /// Full `args` payload as decoded — useful for diagnostics.
        args: serde_json::Value,
    },
}

impl std::fmt::Display for SchedStatsError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Poisoned => write!(f, "scx_stats response buffer mutex was poisoned"),
            Self::RequestTooLarge { size, max } => write!(
                f,
                "scx_stats request size {size} bytes exceeds cap of {max}"
            ),
            Self::ResponseTooLarge { size, max } => write!(
                f,
                "scx_stats response accumulator grew to {size} bytes (cap {max}) \
                 without emitting a newline; partial bytes discarded"
            ),
            Self::DuringFreeze => write!(
                f,
                "scx_stats request rejected: freeze rendezvous active \
                 (scheduler userspace is paused; responses undefined)"
            ),
            Self::Cancelled => write!(
                f,
                "scx_stats request cancelled: run-wide kill flag set \
                 (watchdog fired or shutdown in progress)"
            ),
            Self::NoScheduler { reason } => write!(
                f,
                "scx_stats relay reports no scheduler available: {reason}"
            ),
            Self::SchedulerError { errno, args } => {
                write!(f, "scx_stats scheduler returned errno={errno}: {args}")
            }
            Self::MissingResp { args } => write!(
                f,
                "scx_stats response envelope missing \"resp\" key in args: {args}"
            ),
        }
    }
}

impl std::error::Error for SchedStatsError {}

/// Vendored scx_stats request envelope. Wire shape pinned to
/// scx_stats v1.1.0 `StatsRequest` (server.rs:103-117 of the
/// upstream crate): a JSON object with a `"req"` verb and
/// optional `"args"` map of string-valued keys. We do not depend
/// on the scx_stats crate to keep the dependency surface minimal
/// for a 2-field type.
///
/// `args` matches upstream byte-for-byte as `BTreeMap<String,
/// String>` — non-string argument values would deserialize-fail
/// on the scheduler side (the upstream `BTreeMap<String, String>`
/// rejects e.g. integers). Keep this in lockstep with upstream
/// or update both sides together.
///
/// Wire format examples:
/// * `{"req":"stats"}\n`
/// * `{"req":"stats_meta"}\n`
/// * `{"req":"stats","args":{"target":"foo"}}\n`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StatsRequest {
    /// Verb name. scx_stats schedulers handle `"stats"`,
    /// `"stats_meta"`, and scheduler-specific verbs.
    pub req: String,
    /// Optional argument map. Empty by default; scheduler-specific
    /// verbs may require argument keys. Values are strings on the
    /// wire — see the type-level note above for the upstream
    /// constraint.
    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
    pub args: std::collections::BTreeMap<String, String>,
}

/// Vendored scx_stats response envelope. Wire shape pinned to
/// scx_stats v1.1.0 `StatsResponse` (server.rs:119-123 of the
/// upstream crate). The scheduler emits a JSON object on a
/// single line; the shape is verb-specific. We model the on-wire
/// envelope as `{"errno": Option<i32>, "args": Value}` —
/// present-and-non-zero `errno` flags scheduler-side errors,
/// otherwise `args` carries the verb's payload.
///
/// `args` is intentionally `serde_json::Value` (more permissive
/// than upstream's `BTreeMap<String, Value>`): every scx_stats
/// scheduler emits an object today, so [`extract_resp`] checks
/// `Value::Object` before unwrapping `args["resp"]` and the
/// permissiveness costs nothing on the inbound path.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StatsResponse {
    /// scx_stats error code. `0` (or absent) signals success.
    /// Non-zero matches the scx_stats schedulers' error-mapping
    /// convention (`Err(...)` → libc errno).
    #[serde(default)]
    pub errno: i32,
    /// Verb-specific payload. For `"stats"` this is the scheduler's
    /// stats object; for `"stats_meta"` it is the schema map; for
    /// scheduler-specific verbs it is whatever the scheduler
    /// emits.
    #[serde(default = "default_args_value")]
    pub args: serde_json::Value,
}

fn default_args_value() -> serde_json::Value {
    serde_json::Value::Null
}

/// Inline error payload the in-guest relay emits when it cannot
/// connect to the scheduler's Unix socket. Wire format is
/// `{"ktstr_relay_error":"<reason>"}\n`. Callers receive
/// [`SchedStatsError::NoScheduler`] when this is detected; the
/// JSON shape lets the host translate the error without
/// scheduler-side cooperation.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct RelayError {
    ktstr_relay_error: String,
}

/// Shared response accumulator written by the drainer thread and
/// drained by [`SchedStatsClient::request_raw`].
///
/// `Mutex<Vec<u8>>` carries the bytes; the [`Condvar`] is
/// signalled after every append so blocked readers wake immediately
/// rather than polling.
type StatsResponseBuf = Arc<(Mutex<Vec<u8>>, Condvar)>;

/// Internal state shared between every clone of a
/// [`SchedStatsClient`] AND the drainer thread.
///
/// `Drop` on the last `Arc` clone fires the kill eventfd so the
/// drainer thread exits.
struct ClientShared {
    /// Device handle. Used by `request_raw` to push request bytes
    /// via `queue_input_port2` AND by the drainer to call
    /// `drain_port2_bulk`. Both are short critical sections.
    virtio_con: Arc<PiMutex<VirtioConsole>>,
    /// Response accumulator + Condvar for blocked readers.
    response_buf: StatsResponseBuf,
    /// Optional shared freeze flag. Set by the freeze coordinator
    /// during failure-dump rendezvous; gates `request_raw` calls
    /// (returns [`SchedStatsError::DuringFreeze`]).
    freeze: Option<Arc<AtomicBool>>,
    /// Optional run-wide cancel flag. When `Some(flag)` and the
    /// flag reads `true` (Acquire), `request_raw` returns
    /// [`SchedStatsError::Cancelled`] without ever blocking. The
    /// drainer also watches a paired eventfd (`cancel_evt`)
    /// and notifies the response Condvar when it fires so a
    /// blocked `request_raw` wakes immediately on shutdown — no
    /// timeout needed. The host watchdog is the only "timeout"
    /// in the system.
    cancel: Option<Arc<AtomicBool>>,
    /// Internal serialisation lock held for the entire
    /// request/response cycle. scx_stats has no req-id multiplexing,
    /// so concurrent host calls must serialise to avoid mixing
    /// responses; held across `queue_input_port2` and the
    /// `Condvar::wait` wait.
    request_lock: Mutex<()>,
    /// `true` while a `request_raw` call is mid-flight on any
    /// clone of this client. Set inside `request_lock`, cleared
    /// on exit. The drainer reads this to decide whether to
    /// append drained bytes (true) or discard them (false).
    request_in_flight: Arc<AtomicBool>,
    /// Cumulative count of bytes the drainer has discarded because
    /// no request was in flight when they arrived. Logged by the
    /// drainer at every discard event.
    discarded_bytes: Arc<AtomicU64>,
    /// Eventfd written by `Drop` on the last `Arc` clone. The
    /// drainer's epoll wakes on this and exits.
    kill_drainer: EventFd,
}

impl Drop for ClientShared {
    fn drop(&mut self) {
        // Wake the drainer so it observes the implicit "shared
        // state is being torn down" (the EventFd write triggers
        // the `kill_drainer` epoll fd; the drainer thread's
        // epoll_wait returns and the thread exits). Failure to
        // write is harmless — the drainer's next epoll cycle on
        // stats_tx_evt will re-park, and the kernel will reap the
        // thread when the process exits.
        let _ = self.kill_drainer.write(1);
    }
}

/// Host-side client for the scheduler-stats virtio-console bridge.
///
/// Cloning is cheap: every field is `Arc`-shared and no additional
/// drainer thread is spawned. Drop on the LAST clone tears down
/// the drainer.
///
/// # Freeze semantics
///
/// scx_stats responses are undefined during a host-side freeze
/// rendezvous: the freeze coordinator suspends every vCPU thread,
/// the in-guest scheduler userspace stops advancing, and the
/// relay's blocking read on `/var/run/scx/root/stats` does not
/// return until the scheduler resumes. The optional `freeze` flag
/// gates requests on this state — when the flag is `true`,
/// `request_raw` returns [`SchedStatsError::DuringFreeze`] without
/// touching the device.
#[derive(Clone)]
pub struct SchedStatsClient {
    shared: Arc<ClientShared>,
}

impl SchedStatsClient {
    /// Construct a new client and spawn its drainer thread.
    ///
    /// Scope is `pub(crate)` because [`PiMutex<VirtioConsole>`] is
    /// `pub(crate)` — only the in-tree `KtstrVm::run_vm` wiring
    /// constructs clients; user code receives a fully constructed
    /// [`SchedStatsClient`] via
    /// [`crate::vmm::VmResult::stats_client`] and calls only the
    /// `pub` request methods on it.
    pub(crate) fn new(
        virtio_con: Arc<PiMutex<VirtioConsole>>,
        freeze: Option<Arc<AtomicBool>>,
        cancel: Option<Arc<AtomicBool>>,
        cancel_evt: Option<Arc<EventFd>>,
    ) -> std::io::Result<Self> {
        let response_buf: StatsResponseBuf = Arc::new((Mutex::new(Vec::new()), Condvar::new()));
        let request_in_flight = Arc::new(AtomicBool::new(false));
        let discarded_bytes = Arc::new(AtomicU64::new(0));
        let kill_drainer = EventFd::new(EFD_NONBLOCK)?;
        let kill_drainer_for_thread = kill_drainer.try_clone()?;
        let stats_tx_evt = virtio_con.lock().stats_tx_evt().try_clone()?;

        let shared = Arc::new(ClientShared {
            virtio_con: virtio_con.clone(),
            response_buf: Arc::clone(&response_buf),
            freeze,
            cancel,
            request_lock: Mutex::new(()),
            request_in_flight: Arc::clone(&request_in_flight),
            discarded_bytes: Arc::clone(&discarded_bytes),
            kill_drainer,
        });

        // Spawn the drainer thread. It owns the stats_tx_evt clone
        // (parked in epoll), the kill_drainer clone, and an
        // optional cancel_evt clone. Lifetimes are tied to the
        // thread itself so dropping the last `Arc<ClientShared>`
        // writes kill_drainer, the epoll wakes, and the thread
        // exits — and a host-wide cancel edge wakes the drainer
        // first so any blocked `request_raw` cvar wait gets a
        // notify_all + cancel-flag check before its caller's
        // watchdog deadline elapses.
        let drain_virtio_con = virtio_con;
        let drain_response_buf = response_buf;
        let drain_request_in_flight = request_in_flight;
        let drain_discarded_bytes = discarded_bytes;
        std::thread::Builder::new()
            .name("ktstr-sched-stats-drain".into())
            .spawn(move || {
                drainer_loop(
                    stats_tx_evt,
                    kill_drainer_for_thread,
                    cancel_evt,
                    drain_virtio_con,
                    drain_response_buf,
                    drain_request_in_flight,
                    drain_discarded_bytes,
                );
            })?;

        Ok(Self { shared })
    }

    /// Send a raw request line (without trailing `\n`) and return
    /// the raw response line (also without `\n`). The trailing
    /// newline scx_stats expects on the wire is appended internally
    /// — callers pass JSON or any other line-shaped payload as a
    /// `&str` and never have to worry about framing.
    pub fn request_raw(&self, line: &str) -> Result<Vec<u8>, SchedStatsError> {
        // Phase A: on-wire byte length (line + trailing '\n') cap +
        // freeze + cancel pre-checks. Reject before touching the
        // device or the lock so a misuse that exceeds
        // MAX_REQUEST_BYTES does not stall a concurrent caller.
        let on_wire_len = line.len().saturating_add(1);
        if on_wire_len > MAX_REQUEST_BYTES {
            return Err(SchedStatsError::RequestTooLarge {
                size: on_wire_len,
                max: MAX_REQUEST_BYTES,
            });
        }
        if let Some(flag) = self.shared.freeze.as_ref()
            && flag.load(Ordering::Acquire)
        {
            return Err(SchedStatsError::DuringFreeze);
        }
        if let Some(flag) = self.shared.cancel.as_ref()
            && flag.load(Ordering::Acquire)
        {
            return Err(SchedStatsError::Cancelled);
        }

        // Serialise the entire request/response cycle. The lock is
        // held until this method returns so the drainer's
        // `request_in_flight` flag (set inside the lock, cleared on
        // exit) cannot race with a concurrent caller. Mutex
        // poisoning recovers via `into_inner` — we only ever access
        // the unit `()` payload, which has no invariant a panicked
        // thread could leave broken.
        let _request_guard = self.shared.request_lock.lock_unpoisoned();

        self.predrain_stale_port2_tx();

        // Mark in-flight so the drainer routes drained bytes to the
        // response buffer instead of discarding them. RAII clear on
        // scope exit covers every error-return path below. The store
        // and guard stay HERE (not in a helper): in_flight is set
        // after the pre-drain and before the clear/wait, and the
        // guard's Drop spans the wait and fires on return.
        self.shared.request_in_flight.store(true, Ordering::Release);
        let _in_flight_guard = InFlightGuard {
            flag: &self.shared.request_in_flight,
        };

        self.clear_stale_response_buf()?;
        self.push_request_to_port2(line);
        self.wait_for_response_line()
    }

    /// Drain any stale guest→host bytes from port2_tx_buf BEFORE
    /// `request_raw` flips in_flight=true. This NARROWS the
    /// stale-bytes race window but does not fully close it.
    ///
    /// Closed portion — bytes already in port2_tx_buf at this
    /// point: pre-drain empties `port[2].tx_buf` under the same
    /// virtio_con lock the drainer's TOKEN_DATA arm
    /// (drainer_loop) takes, so the drainer's next wake sees an
    /// empty buffer and routes nothing to response_buf.
    /// discarded_bytes counter bumps so stuck-stats post-mortem
    /// can see this path firing.
    ///
    /// Residual portion — bytes that arrive after pre-drain release:
    /// `VirtioConsole::process_tx` (port-id 2, defined in
    /// virtio_console/device/mod.rs) runs only when the vCPU services a
    /// QUEUE_NOTIFY MMIO write on port-2 TX (q7 per the multiport
    /// queue numbering in virtio_console/mod.rs). The guest can
    /// publish a chain to the avail ring + issue NOTIFY BEFORE
    /// our pre-drain, but the vCPU MMIO handler may not have run
    /// yet (e.g. coming out of a freeze rendezvous that paused
    /// every vCPU via SIGRTMIN). After our pre-drain releases
    /// virtio_con, the vCPU services the queued NOTIFY, process_tx
    /// pushes stale guest bytes into `port[2].tx_buf`, kicks
    /// stats_tx_evt. If this lands AFTER our in_flight=true store
    /// (in `request_raw`), the drainer wakes, sees in_flight=true,
    /// forwards the stale bytes as if they were the current
    /// request's response.
    ///
    /// The realistic trigger is cancellation-followed-by-immediate-
    /// new-request (freeze-rendezvous cancel, watchdog cancel)
    /// where the guest is still flushing a prior abandoned
    /// request's response. A hostile guest can also exploit this
    /// window deliberately.
    ///
    /// Fully closing this is NOT possible from host-side alone.
    /// Two protocol-level paths were investigated:
    ///   (A) per-request token echoed by the guest scheduler —
    ///       requires patching scx_stats (StatsResponse hardcodes
    ///       `{resp: ...}` with no echo of request fields) AND
    ///       every scheduler binary that uses it. Out of ktstr's
    ///       scope.
    ///   (B) generation counter on response_buf chunks — provably
    ///       fails: generation tracks DRAIN time (when the drainer
    ///       processes bytes), not PRODUCTION time (when the guest
    ///       emitted them). Stale-but-late bytes from a freeze /
    ///       cancel window get stamped with the current generation
    ///       at drain time, so the wait-loop accepts them.
    /// The residual class is a documentation note about an
    /// upstream-protocol limitation, not an actionable host-side
    /// follow-up.
    ///
    /// Race partner: the drainer's TOKEN_DATA arm in drainer_loop
    /// (`fn drainer_loop` later in this file) reads in_flight
    /// (Acquire) after the drain-under-virtio_con-lock completes
    /// and the lock is released — that's where the racing append
    /// happens.
    fn predrain_stale_port2_tx(&self) {
        let stale = {
            let mut g = self.shared.virtio_con.lock();
            g.drain_port2_bulk()
        };
        if !stale.is_empty() {
            let stale_len = stale.len();
            let total = self
                .shared
                .discarded_bytes
                .fetch_add(stale_len as u64, Ordering::Relaxed)
                .saturating_add(stale_len as u64);
            tracing::debug!(
                stale_bytes = stale_len,
                total_discarded = total,
                "scx_stats request_raw: drained stale port2_tx_buf bytes \
                 before flipping in_flight=true (stale-bytes race-close)"
            );
        }
    }

    /// Drain any stale bytes left over from a prior incomplete
    /// call. scx_stats is strict request/response — anything
    /// sitting in the buffer when a fresh request starts is
    /// either a torn previous response or relay-restart
    /// residue. Log the discard so a stuck-relay diagnosis can
    /// see how many bytes were thrown away on each request,
    /// and bump `discarded_bytes` so the counter exposed via
    /// [`Self::discarded_bytes`] reflects this path too — not
    /// just the drainer's no-request-in-flight discards.
    fn clear_stale_response_buf(&self) -> Result<(), SchedStatsError> {
        let (lock, _cvar) = &*self.shared.response_buf;
        {
            let mut buf = lock.lock().map_err(|_| SchedStatsError::Poisoned)?;
            if !buf.is_empty() {
                let stale = buf.len();
                let total = self
                    .shared
                    .discarded_bytes
                    .fetch_add(stale as u64, Ordering::Relaxed)
                    .saturating_add(stale as u64);
                tracing::debug!(
                    stale_bytes = stale,
                    total_discarded = total,
                    "scx_stats request_raw: clearing stale response bytes from prior call"
                );
                buf.clear();
            }
        }
        Ok(())
    }

    /// Push the request bytes onto port 2 RX. Append the
    /// trailing newline scx_stats expects. Bound the device
    /// mutex critical section to the queue_input_port2 call.
    ///
    /// Drop any host→guest bytes that are still sitting in
    /// `ports[2].pending_rx` from a prior request that was abandoned
    /// mid-push (e.g. a freeze rendezvous landed before the guest
    /// read those bytes). Without this clear, the new request
    /// would be concatenated onto the dead tail of the previous
    /// one and the guest relay would forward torn JSON to the
    /// scheduler. Account for the discard via `discarded_bytes`
    /// so a stuck-stats post-mortem can see it.
    fn push_request_to_port2(&self, line: &str) {
        let mut g = self.shared.virtio_con.lock();
        let stale_in = g.clear_port2_pending_rx();
        if stale_in > 0 {
            let total = self
                .shared
                .discarded_bytes
                .fetch_add(stale_in as u64, Ordering::Relaxed)
                .saturating_add(stale_in as u64);
            tracing::debug!(
                stale_pending_rx = stale_in,
                total_discarded = total,
                "scx_stats request_raw: clearing stale port2_pending_rx \
                 (prior request abandoned mid-push)"
            );
        }
        // Two pushes are equivalent to one combined push from
        // the device's perspective — the bytes land in the
        // pending-RX deque in order.
        g.queue_input_port2(line.as_bytes());
        g.queue_input_port2(b"\n");
    }

    /// Wait (no timeout) for a complete line, freeze, or cancel.
    /// The drainer notifies the cvar on every appended byte AND
    /// when its kill_drainer eventfd fires (cancel edge or
    /// last-Arc drop), so a blocked `cvar.wait` always wakes
    /// promptly without polling. The host watchdog is the only
    /// backstop "timeout" — when it fires it sets `cancel`,
    /// writes the cancel_evt, and the drainer's notify_all
    /// wakes us into the cancel-check below.
    fn wait_for_response_line(&self) -> Result<Vec<u8>, SchedStatsError> {
        let (lock, cvar) = &*self.shared.response_buf;
        let mut buf = lock.lock().map_err(|_| SchedStatsError::Poisoned)?;
        loop {
            if let Some(idx) = buf.iter().position(|&b| b == b'\n') {
                let mut response = buf.split_off(idx + 1);
                std::mem::swap(&mut *buf, &mut response);
                response.pop();
                // Inspect for the relay-emitted error envelope. If
                // the line parses as `{"ktstr_relay_error":"..."}`
                // surface it as the typed NoScheduler error so the
                // caller can branch without parsing the JSON.
                if let Ok(err) = serde_json::from_slice::<RelayError>(&response) {
                    return Err(SchedStatsError::NoScheduler {
                        reason: err.ktstr_relay_error,
                    });
                }
                return Ok(response);
            }
            if buf.len() > MAX_RESPONSE_BYTES {
                let size = buf.len();
                buf.clear();
                return Err(SchedStatsError::ResponseTooLarge {
                    size,
                    max: MAX_RESPONSE_BYTES,
                });
            }
            if let Some(flag) = self.shared.freeze.as_ref()
                && flag.load(Ordering::Acquire)
            {
                buf.clear();
                return Err(SchedStatsError::DuringFreeze);
            }
            if let Some(flag) = self.shared.cancel.as_ref()
                && flag.load(Ordering::Acquire)
            {
                buf.clear();
                return Err(SchedStatsError::Cancelled);
            }
            buf = cvar.wait(buf).map_err(|_| SchedStatsError::Poisoned)?;
        }
    }

    /// Send a typed scx_stats request and parse the typed response
    /// envelope. The envelope's `errno` field is NOT checked here
    /// — callers that want errno-as-error semantics should use
    /// [`Self::stats`] / [`Self::stats_meta`] (which surface
    /// non-zero errno as [`SchedStatsError::SchedulerError`]) or
    /// inspect `StatsResponse.errno` themselves.
    pub fn request(&self, request: &StatsRequest) -> Result<StatsResponse, anyhow::Error> {
        // Encode as JSON without the trailing newline — request_raw
        // appends it. `to_string` is preferable to `to_vec` here
        // because request_raw takes a `&str`.
        let line = serde_json::to_string(request)?;
        let raw = self.request_raw(&line).map_err(|e| anyhow::anyhow!(e))?;
        Ok(serde_json::from_slice::<StatsResponse>(&raw)?)
    }

    /// Convenience: send a `"stats"` request and return the inner
    /// `args["resp"]` payload. `args` is an ordered list of
    /// scheduler-specific argument key/value pairs (e.g.
    /// `&[("target", "top")]` to request the scheduler's "top"
    /// dataset). Pass `&[]` for the default `"stats"` payload.
    ///
    /// Errors:
    /// * [`SchedStatsError::SchedulerError`] when the scheduler's
    ///   response envelope carries `errno != 0`.
    /// * [`SchedStatsError::MissingResp`] when the envelope
    ///   succeeded but `args["resp"]` is absent (protocol
    ///   mismatch).
    /// * Other variants surface from the underlying
    ///   [`Self::request`] / [`Self::request_raw`] flow.
    pub fn stats(&self, args: &[(&str, &str)]) -> Result<serde_json::Value, anyhow::Error> {
        let mut req = StatsRequest {
            req: "stats".to_string(),
            args: std::collections::BTreeMap::new(),
        };
        for (k, v) in args {
            req.args.insert((*k).to_string(), (*v).to_string());
        }
        let resp = self.request(&req)?;
        extract_resp(resp)
    }

    /// Convenience: send a `"stats_meta"` request and return the
    /// inner `args["resp"]` payload (the metadata schema for the
    /// scheduler's `"stats"` payload).
    pub fn stats_meta(&self) -> Result<serde_json::Value, anyhow::Error> {
        let req = StatsRequest {
            req: "stats_meta".to_string(),
            args: std::collections::BTreeMap::new(),
        };
        let resp = self.request(&req)?;
        extract_resp(resp)
    }

    /// Cumulative number of bytes the client has discarded over its
    /// lifetime. Four paths feed this counter:
    /// 1. The drainer thread receives port-2 bytes when no request
    ///    is in flight (a stale response from a torn prior call,
    ///    or relay-emitted bytes the host never asked for).
    /// 2. [`Self::request_raw`] clears stale bytes from the
    ///    response buffer at request start (residue left in
    ///    `response_buf` from prior request's partial-newline
    ///    consumption — the post-newline tail per `split_off`
    ///    semantics in the wait-loop's success arm).
    /// 3. [`Self::request_raw`] clears stale `ports[2].pending_rx`
    ///    (host→guest) at request start — bytes from a prior
    ///    request abandoned mid-push would otherwise concatenate
    ///    onto the new request line.
    /// 4. [`Self::request_raw`] pre-drains `ports[2].tx_buf`
    ///    (guest→host) BEFORE flipping `in_flight=true`
    ///    (stale-bytes race-close — bytes already in the buffer
    ///    that the drainer would otherwise misattribute to the
    ///    current request).
    ///
    /// All four paths bump `discarded_bytes`, so this accessor lets
    /// test code (and operators inspecting a stuck stats path)
    /// see whether bytes are leaking past the request/response
    /// envelope. The per-site `tracing::debug!` log lines carry
    /// the per-path context (search for "scx_stats" in stats
    /// logs to disambiguate which path bumped the counter).
    pub fn discarded_bytes(&self) -> u64 {
        self.shared.discarded_bytes.load(Ordering::Relaxed)
    }
}

/// Convert a [`StatsResponse`] envelope into the verb-specific
/// inner payload: rejects non-zero `errno`, then extracts the
/// `"resp"` key from `args`. Centralised so [`SchedStatsClient::stats`]
/// and [`SchedStatsClient::stats_meta`] share identical error
/// surfaces.
fn extract_resp(resp: StatsResponse) -> Result<serde_json::Value, anyhow::Error> {
    if resp.errno != 0 {
        return Err(anyhow::anyhow!(SchedStatsError::SchedulerError {
            errno: resp.errno,
            args: resp.args,
        }));
    }
    let serde_json::Value::Object(mut map) = resp.args else {
        return Err(anyhow::anyhow!(SchedStatsError::MissingResp {
            args: resp.args
        }));
    };
    match map.remove("resp") {
        Some(v) => Ok(v),
        None => Err(anyhow::anyhow!(SchedStatsError::MissingResp {
            args: serde_json::Value::Object(map),
        })),
    }
}

impl std::fmt::Debug for SchedStatsClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SchedStatsClient")
            .field(
                "response_buf_len",
                &self
                    .shared
                    .response_buf
                    .0
                    .lock()
                    .map(|b| b.len())
                    .unwrap_or(0),
            )
            .field(
                "discarded_bytes",
                &self.shared.discarded_bytes.load(Ordering::Relaxed),
            )
            .field(
                "request_in_flight",
                &self.shared.request_in_flight.load(Ordering::Relaxed),
            )
            .finish()
    }
}

/// RAII guard: clear `request_in_flight` on drop. Prevents a
/// stuck-true flag if `request_raw` returns early via
/// `?` from any of the lock / mutex calls.
struct InFlightGuard<'a> {
    flag: &'a Arc<AtomicBool>,
}

impl<'a> Drop for InFlightGuard<'a> {
    fn drop(&mut self) {
        self.flag.store(false, Ordering::Release);
    }
}

/// Drainer-thread main loop. Runs until `kill_drainer` fires (the
/// last `Arc<ClientShared>` clone has dropped) OR `cancel_evt`
/// fires (the run-wide kill flag was set, e.g. watchdog timeout).
/// Drains `ports[2].tx_buf` on every `stats_tx_evt` wake; appends to
/// the response buffer when a request is in flight, discards
/// otherwise. On either kill edge, performs a final
/// `cvar.notify_all()` so any blocked `request_raw` wakes,
/// observes the cancel flag, and returns
/// [`SchedStatsError::Cancelled`].
fn drainer_loop(
    stats_tx_evt: EventFd,
    kill_drainer: EventFd,
    cancel_evt: Option<Arc<EventFd>>,
    virtio_con: Arc<PiMutex<VirtioConsole>>,
    response_buf: StatsResponseBuf,
    request_in_flight: Arc<AtomicBool>,
    discarded_bytes: Arc<AtomicU64>,
) {
    const TOKEN_DATA: u64 = 0;
    const TOKEN_KILL: u64 = 1;
    const TOKEN_CANCEL: u64 = 2;

    let epoll = match Epoll::new() {
        Ok(e) => e,
        Err(e) => {
            tracing::error!(error = %e, "stats drainer: epoll_create1 failed; aborting drainer");
            return;
        }
    };
    for (fd, token, name) in [
        (stats_tx_evt.as_raw_fd(), TOKEN_DATA, "stats_tx_evt"),
        (kill_drainer.as_raw_fd(), TOKEN_KILL, "kill_drainer"),
    ] {
        if let Err(e) = epoll.ctl(
            ControlOperation::Add,
            fd,
            EpollEvent::new(EventSet::IN, token),
        ) {
            tracing::error!(
                error = %e,
                fd_name = name,
                "stats drainer: epoll_ctl ADD failed; aborting drainer"
            );
            return;
        }
    }
    // Optional cancel_evt: the host's run-wide kill_evt clone.
    // When the watchdog or the BSP shutdown path writes it, our
    // epoll wakes; we notify the response cvar (so blocked
    // `request_raw` waiters re-check the cancel flag) and exit.
    if let Some(cancel) = cancel_evt.as_ref()
        && let Err(e) = epoll.ctl(
            ControlOperation::Add,
            cancel.as_raw_fd(),
            EpollEvent::new(EventSet::IN, TOKEN_CANCEL),
        )
    {
        tracing::warn!(
            error = %e,
            "stats drainer: epoll_ctl ADD on cancel_evt failed; \
             cancel edge will not wake blocked requests promptly"
        );
    }

    let mut events_buf = [EpollEvent::default(); 3];
    loop {
        let event_count = match epoll.wait(-1, &mut events_buf) {
            Ok(n) => n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => {
                tracing::error!(error = %e, "stats drainer: epoll_wait failed; exiting");
                return;
            }
        };
        for ev in &events_buf[..event_count] {
            match ev.data() {
                TOKEN_KILL => {
                    let _ = kill_drainer.read();
                    // Acquire the response_buf lock BEFORE
                    // notifying the cvar. Without the lock the
                    // notify_all can fire after the request thread
                    // has dropped the lock-guard but BEFORE it has
                    // entered cvar.wait — that wake is then lost
                    // and the request hangs until the next
                    // legitimate notify (which on the kill path
                    // will never come). Holding the lock across
                    // notify_all forces the drainer to interleave
                    // with the request's check-and-park sequence:
                    // if the request has already entered cvar.wait,
                    // notify_all wakes it; if the request is still
                    // inside the response-buf check, the drainer
                    // blocks on the lock until the request enters
                    // cvar.wait, then notifies. Either way the
                    // wake is delivered.
                    //
                    // `lock()` returns LockResult; both Ok and Err
                    // (poisoned) variants contain a MutexGuard
                    // that unlocks on Drop, so binding the
                    // LockResult to `_guard` is sufficient — we
                    // don't read the guarded payload, only need
                    // exclusive ownership for the notify.
                    let (lock, cvar) = &*response_buf;
                    let _guard = lock.lock();
                    cvar.notify_all();
                    return;
                }
                TOKEN_CANCEL => {
                    // Drain the cancel eventfd counter (the writer
                    // set the run-wide cancel AtomicBool first;
                    // request_raw rechecks it after wakeup).
                    if let Some(c) = cancel_evt.as_ref() {
                        let _ = c.read();
                    }
                    // Lock-then-notify, same TOCTOU rationale
                    // as the TOKEN_KILL arm above.
                    let (lock, cvar) = &*response_buf;
                    let _guard = lock.lock();
                    cvar.notify_all();
                    return;
                }
                TOKEN_DATA => {
                    let _ = stats_tx_evt.read();
                    // Drain regardless of in-flight state so a
                    // hostile or runaway guest can't grow
                    // port2_tx_buf without bound when no host
                    // request is outstanding.
                    let bytes = {
                        let mut g = virtio_con.lock();
                        g.drain_port2_bulk()
                    };
                    if bytes.is_empty() {
                        continue;
                    }
                    if request_in_flight.load(Ordering::Acquire) {
                        let (lock, cvar) = &*response_buf;
                        if let Ok(mut guard) = lock.lock() {
                            // Hard 2x cap to prevent unbounded growth
                            // even when a request is mid-flight (a
                            // hostile guest could send bytes faster
                            // than the request can complete).
                            // Inject a synthetic error envelope so
                            // request_raw wakes from cvar.wait,
                            // observes a complete `\n`-terminated
                            // line, and surfaces NoScheduler — the
                            // prior code dropped bytes silently and
                            // the request blocked forever.
                            let new_total = guard.len().saturating_add(bytes.len());
                            if new_total > MAX_RESPONSE_BYTES.saturating_mul(2) {
                                tracing::warn!(
                                    current = guard.len(),
                                    incoming = bytes.len(),
                                    cap = MAX_RESPONSE_BYTES * 2,
                                    "stats drainer: hard cap reached; injecting cap-overflow error envelope"
                                );
                                guard.clear();
                                guard.extend_from_slice(CAP_OVERFLOW_ERROR_REPLY);
                                cvar.notify_all();
                                continue;
                            }
                            guard.extend_from_slice(&bytes);
                            cvar.notify_all();
                        }
                    } else {
                        // Discard bytes when no request is
                        // outstanding, log the discard count.
                        let total = discarded_bytes
                            .fetch_add(bytes.len() as u64, Ordering::Relaxed)
                            .saturating_add(bytes.len() as u64);
                        tracing::debug!(
                            this_drop = bytes.len(),
                            total_discarded = total,
                            "stats drainer: discarding port-2 bytes (no request in flight)"
                        );
                    }
                }
                _ => {}
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    fn make_client_full(
        freeze: Option<Arc<AtomicBool>>,
        cancel: Option<Arc<AtomicBool>>,
        cancel_evt: Option<Arc<EventFd>>,
    ) -> SchedStatsClient {
        let virtio_con = Arc::new(PiMutex::new(VirtioConsole::new()));
        SchedStatsClient::new(virtio_con, freeze, cancel, cancel_evt).expect("client construct")
    }

    fn make_client() -> SchedStatsClient {
        make_client_full(None, None, None)
    }

    /// Helper: simulate the drainer appending a complete response
    /// line. The drainer normally sets `request_in_flight` from the
    /// device side; tests stand it up manually. The thread waits
    /// briefly for `request_in_flight` to flip true (the request
    /// thread sets it after grabbing the request_lock) before
    /// appending so the test races the request rather than racing
    /// the writer.
    fn pre_populate(client: &SchedStatsClient, bytes: &[u8]) -> std::thread::JoinHandle<()> {
        let buf = Arc::clone(&client.shared.response_buf);
        let in_flight = Arc::clone(&client.shared.request_in_flight);
        let bytes = bytes.to_vec();
        std::thread::spawn(move || {
            // Wait for the request thread to mark in-flight. Bound
            // the wait with a finite spin count so a test bug
            // doesn't hang the suite.
            for _ in 0..200 {
                if in_flight.load(Ordering::Acquire) {
                    break;
                }
                std::thread::sleep(Duration::from_millis(1));
            }
            let (lock, cvar) = &*buf;
            let mut guard = lock.lock().unwrap();
            guard.extend_from_slice(&bytes);
            cvar.notify_all();
        })
    }

    /// Pre-populating the response buffer wakes a blocked
    /// `request_raw` and returns the bytes before the first newline.
    #[test]
    fn drainer_append_then_request_returns_first_line() {
        let client = make_client();
        let writer = pre_populate(&client, b"hello\n");
        let resp = client
            .request_raw("x")
            .expect("must wake when bytes arrive");
        assert_eq!(resp, b"hello");
        writer.join().unwrap();
    }

    /// A request larger than [`MAX_REQUEST_BYTES`] (after appending
    /// the trailing newline) is rejected before any device
    /// interaction.
    #[test]
    fn oversize_request_rejected() {
        let client = make_client();
        // request_raw appends '\n', so a line of MAX_REQUEST_BYTES
        // bytes triggers the cap (MAX + 1 on wire).
        let big = "x".repeat(MAX_REQUEST_BYTES);
        let err = client
            .request_raw(&big)
            .expect_err("must reject oversize request");
        match err {
            SchedStatsError::RequestTooLarge { size, max } => {
                assert_eq!(size, MAX_REQUEST_BYTES + 1);
                assert_eq!(max, MAX_REQUEST_BYTES);
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
    }

    /// Pre-existing bytes in port2_tx_buf at request start are
    /// drained + discarded BEFORE in_flight=true flips, so the
    /// drainer's next wake cannot route them to response_buf as
    /// the current request's response. Pins the stale-bytes
    /// race-close at the head of `request_raw`.
    ///
    /// Mechanism: inject stale bytes directly into port[2].tx_buf
    /// via the cfg(test) helper, then call request_raw with a
    /// pre_populate writer staged to deliver the legitimate
    /// response. The pre-drain at request_raw entry must catch
    /// the stale bytes (verified via `discarded_bytes` counter
    /// increment); the wait-loop must return the legitimate
    /// pre_populate response, NOT the stale bytes.
    ///
    /// The test bypasses stats_tx_evt (production drainer wake
    /// signal) — injection is silent, so the drainer never
    /// processes the stale bytes. The pre-drain at request_raw
    /// entry is the load-bearing path. A regression that removes
    /// the pre-drain would leave port[2].tx_buf populated when
    /// in_flight=true flips; any subsequent drainer wake would
    /// route the stale to response_buf. The test's
    /// discarded_bytes assertion fails-loud if the pre-drain
    /// stops firing.
    ///
    /// Pins the closed portion of the race-close (see the
    /// closed-portion section of the race-close comment at the head
    /// of `request_raw`). The residual portion — bytes arriving after
    /// pre-drain release via lagging vCPU MMIO QUEUE_NOTIFY service —
    /// is NOT exercised by this test; it is fundamentally unclosable
    /// from host-side alone.
    #[test]
    fn request_raw_drains_stale_port2_tx_buf_before_in_flight() {
        let client = make_client();
        let initial_discarded = client.discarded_bytes();
        let stale: &[u8] = b"prior-request-tail-garbage";

        // Inject pre-existing stale bytes into port[2].tx_buf
        // (simulates bytes from a prior request's tail still
        // sitting in the device when the new request starts).
        {
            let mut g = client.shared.virtio_con.lock();
            g.inject_port2_tx_for_test(stale);
        }

        let expected_response = b"new-request-response\n";
        let writer = pre_populate(&client, expected_response);

        let resp = client
            .request_raw("dummy")
            .expect("request_raw must return the legitimate response");

        assert_eq!(
            &resp, b"new-request-response",
            "request_raw must return the legitimate pre_populate \
             response, NOT the stale port[2].tx_buf bytes — the \
             pre-drain at request_raw entry failed to catch the \
             stale residue"
        );

        let discarded_delta = client.discarded_bytes() - initial_discarded;
        assert!(
            discarded_delta >= stale.len() as u64,
            "discarded_bytes counter did not increment by the \
             expected pre-drain amount: delta={discarded_delta} \
             stale.len()={} (the pre-drain path did not fire)",
            stale.len()
        );

        writer.join().unwrap();
    }

    /// A response that grows past [`MAX_RESPONSE_BYTES`] without a
    /// newline returns [`SchedStatsError::ResponseTooLarge`].
    #[test]
    fn oversize_response_returns_response_too_large() {
        let client = make_client();
        let payload: Vec<u8> = std::iter::repeat_n(b'A', MAX_RESPONSE_BYTES + 1).collect();
        let writer = pre_populate(&client, &payload);
        let err = client
            .request_raw("x")
            .expect_err("must reject oversize response");
        match err {
            SchedStatsError::ResponseTooLarge { size, max } => {
                assert!(size > MAX_RESPONSE_BYTES);
                assert_eq!(max, MAX_RESPONSE_BYTES);
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
        writer.join().unwrap();
    }

    /// A request issued while the freeze flag is set returns
    /// [`SchedStatsError::DuringFreeze`] without ever blocking.
    #[test]
    fn during_freeze_rejects_request() {
        let freeze = Arc::new(AtomicBool::new(true));
        let client = make_client_full(Some(freeze.clone()), None, None);

        let err = client
            .request_raw("x")
            .expect_err("must reject during freeze");
        assert!(matches!(err, SchedStatsError::DuringFreeze));

        freeze.store(false, Ordering::Release);
        let writer = pre_populate(&client, b"ok\n");
        let resp = client.request_raw("x").expect("must succeed after thaw");
        assert_eq!(resp, b"ok");
        writer.join().unwrap();
    }

    /// A request issued while the cancel flag is set returns
    /// [`SchedStatsError::Cancelled`] without ever blocking.
    #[test]
    fn cancel_flag_set_before_request_rejects() {
        let cancel = Arc::new(AtomicBool::new(true));
        let cancel_evt = Arc::new(EventFd::new(EFD_NONBLOCK).unwrap());
        let client = make_client_full(None, Some(cancel), Some(cancel_evt));
        let err = client
            .request_raw("x")
            .expect_err("must reject when cancel pre-set");
        assert!(matches!(err, SchedStatsError::Cancelled));
    }

    /// Cancel flag flipped + cancel_evt fired DURING a blocked
    /// request wakes the cvar wait via the drainer, surfaces
    /// [`SchedStatsError::Cancelled`], and tears down cleanly.
    #[test]
    fn cancel_during_blocked_request_wakes() {
        let cancel = Arc::new(AtomicBool::new(false));
        let cancel_evt = Arc::new(EventFd::new(EFD_NONBLOCK).unwrap());
        let client = make_client_full(None, Some(cancel.clone()), Some(cancel_evt.clone()));

        let in_flight = Arc::clone(&client.shared.request_in_flight);
        let waker = std::thread::spawn(move || {
            // Wait until the request actually parks on the cvar
            // (in_flight goes true after grabbing the request_lock
            // and before the wait).
            for _ in 0..200 {
                if in_flight.load(Ordering::Acquire) {
                    break;
                }
                std::thread::sleep(Duration::from_millis(1));
            }
            cancel.store(true, Ordering::Release);
            let _ = cancel_evt.write(1);
        });

        let err = client
            .request_raw("x")
            .expect_err("must wake on cancel edge");
        assert!(matches!(err, SchedStatsError::Cancelled));
        waker.join().unwrap();
    }

    /// Two concurrent `request_raw` calls SERIALISE through the
    /// internal request_lock — the second waits for the first to
    /// finish rather than racing for the response buffer. Both
    /// observe their own response.
    #[test]
    fn concurrent_requests_serialise_via_lock() {
        let client = make_client();
        let buf = Arc::clone(&client.shared.response_buf);
        let in_flight = Arc::clone(&client.shared.request_in_flight);

        // Writer feeds two responses in sequence: each fires once
        // the request thread has marked in_flight (so it's
        // already parked on the cvar). Because the request_lock
        // serialises the two requests, in_flight goes
        // true→false→true across the gap between the two requests.
        let writer = std::thread::spawn(move || {
            // First response.
            for _ in 0..200 {
                if in_flight.load(Ordering::Acquire) {
                    let (lock, cvar) = &*buf;
                    let mut guard = lock.lock().unwrap();
                    guard.extend_from_slice(b"first\n");
                    cvar.notify_all();
                    break;
                }
                std::thread::sleep(Duration::from_millis(1));
            }
            // Wait for the first request to finish (in_flight
            // goes false), then for the second request to grab
            // the lock (in_flight goes true again).
            for _ in 0..200 {
                if !in_flight.load(Ordering::Acquire) {
                    break;
                }
                std::thread::sleep(Duration::from_millis(1));
            }
            for _ in 0..200 {
                if in_flight.load(Ordering::Acquire) {
                    let (lock, cvar) = &*buf;
                    let mut guard = lock.lock().unwrap();
                    if guard.is_empty() {
                        guard.extend_from_slice(b"second\n");
                        cvar.notify_all();
                        break;
                    }
                }
                std::thread::sleep(Duration::from_millis(1));
            }
        });

        let c2 = client.clone();
        let h = std::thread::spawn(move || c2.request_raw("a").expect("first request succeeds"));
        // Brief delay so the first request has a chance to take
        // the lock before the second arrives. This is a test-only
        // synchronisation, not a production timeout.
        std::thread::sleep(Duration::from_millis(20));
        let resp_b = client
            .request_raw("b")
            .expect("second request succeeds (after lock release)");
        let resp_a = h.join().expect("first thread joins");
        // Order may flip depending on which thread acquired the
        // lock first — assert the set of responses received,
        // matching the writer's output.
        let mut got = vec![resp_a, resp_b];
        got.sort();
        assert_eq!(got, vec![b"first".to_vec(), b"second".to_vec()]);
        writer.join().unwrap();
    }

    /// A relay-emitted error envelope (`{"ktstr_relay_error":"..."}`)
    /// surfaces as [`SchedStatsError::NoScheduler`].
    #[test]
    fn relay_error_envelope_surfaces_as_no_scheduler() {
        let client = make_client();
        let writer = pre_populate(
            &client,
            br#"{"ktstr_relay_error":"no scheduler running"}
"#,
        );
        let err = client
            .request_raw("x")
            .expect_err("must surface NoScheduler");
        match err {
            SchedStatsError::NoScheduler { reason } => {
                assert_eq!(reason, "no scheduler running");
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
        writer.join().unwrap();
    }

    /// Typed `request` round-trips through serde — request
    /// serializes to the scx_stats wire format and response
    /// deserializes from it.
    #[test]
    fn typed_request_round_trip() {
        let req = StatsRequest {
            req: "stats".to_string(),
            args: std::collections::BTreeMap::new(),
        };
        let bytes = serde_json::to_vec(&req).unwrap();
        let decoded: StatsRequest = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(decoded.req, "stats");
        assert!(decoded.args.is_empty());

        // Args are serialized as plain JSON strings (matches upstream
        // scx_stats v1.1.0 server.rs:107 `BTreeMap<String, String>`),
        // not wrapped in any other type. A regression to
        // `BTreeMap<String, serde_json::Value>` would emit
        // `{"target":"top"}` either way for string values, so this
        // also pins decoded round-trip equality and rejects integer
        // / object values at parse time on the upstream side.
        let mut req_args = StatsRequest {
            req: "stats".to_string(),
            args: std::collections::BTreeMap::new(),
        };
        req_args
            .args
            .insert("target".to_string(), "top".to_string());
        let bytes_args = serde_json::to_vec(&req_args).unwrap();
        assert_eq!(
            std::str::from_utf8(&bytes_args).unwrap(),
            r#"{"req":"stats","args":{"target":"top"}}"#,
        );
        let decoded_args: StatsRequest = serde_json::from_slice(&bytes_args).unwrap();
        assert_eq!(decoded_args.args.get("target"), Some(&"top".to_string()));

        let resp_wire = br#"{"errno":0,"args":{"resp":{"foo":42}}}"#;
        let resp: StatsResponse = serde_json::from_slice(resp_wire).unwrap();
        assert_eq!(resp.errno, 0);
        assert_eq!(resp.args["resp"]["foo"], 42);
    }

    /// `extract_resp` returns the inner `args["resp"]` payload on
    /// success and surfaces [`SchedStatsError::SchedulerError`] on
    /// non-zero errno.
    #[test]
    fn extract_resp_happy_path_and_errno() {
        let resp_ok = StatsResponse {
            errno: 0,
            args: serde_json::json!({"resp": {"counter": 7}}),
        };
        let payload = extract_resp(resp_ok).unwrap();
        assert_eq!(payload["counter"], 7);

        let resp_err = StatsResponse {
            errno: 22,
            args: serde_json::json!({"message": "EINVAL"}),
        };
        let err = extract_resp(resp_err).unwrap_err();
        let downcast = err.downcast_ref::<SchedStatsError>().expect("downcast");
        match downcast {
            SchedStatsError::SchedulerError { errno, args } => {
                assert_eq!(*errno, 22);
                assert_eq!(args["message"], "EINVAL");
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
    }

    /// `extract_resp` surfaces [`SchedStatsError::MissingResp`]
    /// when `args["resp"]` is absent.
    #[test]
    fn extract_resp_missing_resp_key() {
        let resp = StatsResponse {
            errno: 0,
            args: serde_json::json!({"other": 1}),
        };
        let err = extract_resp(resp).unwrap_err();
        let downcast = err.downcast_ref::<SchedStatsError>().expect("downcast");
        assert!(matches!(downcast, SchedStatsError::MissingResp { .. }));
    }
}
