fv_streams_runtime/

dataflow.rs

//! THE DATAFLOW RUNTIME (Phase 2): a pipeline as tasks on OS threads connected by bounded
//! in-memory edges that carry Arrow batches and the two control signals — watermarks and
//! checkpoint barriers — in band (design/2026-09-15-phase2-dataflow-scoping.md).
//!
//! - A **task** owns one operator and runs a loop over its single inbox, into which every
//!   upstream edge sends tagged messages; FIFO per edge is what per-key order rests on.
//! - An **edge** is `sync_channel`-bounded: a full inbox backpressures the sender.
//! - **Shuffle** hashes the routing columns of a batch (`create_hashes`) into vnodes and sends
//!   one sub-batch (`take`) per downstream task by its vnode range; **forward** sends whole batches
//!   to one task; **broadcast** sends the same batch to every downstream task.
//! - **Watermarks**: each source emits `Watermark { source, ts }` (`None` = idle); a task keeps the
//!   latest per (edge, source), fires its operator on the minimum over the active ones, and emits
//!   its own watermark downstream under its own id, so the whole graph shares one time.
//! - **Barriers and epochs**: the coordinator issues epochs — periodically from the runtime's own
//!   ticker, and one last time on `stop` — as a `Control::Barrier` to every source; a source
//!   records its position, reports a `Snapshot`, and forwards `Barrier { epoch }` behind the data
//!   it already sent. A task with several input edges **aligns** — batches arriving on an edge
//!   that already delivered the barrier are held until every edge has — then snapshots its
//!   operator, forwards the barrier, and hands the snapshot to the coordinator. Sinks make their
//!   output durable and `Ack`. When every task has reported an epoch (or finished), the epoch is
//!   complete: the coordinator sends `Control::Commit` to the sources, which commit the positions
//!   they recorded at that barrier. Offsets therefore never run ahead of durable output — the
//!   at-least-once gate, now across task boundaries.
//! - **Eos** flows in order behind the last data and the final barrier; a task finishes when every
//!   input reached it.
//!
//! The runtime is Kafka-free: sources and sinks are traits, tested here with in-memory ones.

use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{channel, sync_channel, Receiver, RecvTimeoutError, Sender, SyncSender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use tokio::runtime::Handle;

use fv_streams_exchange::{LinkReceiver, LinkSender};

use crate::exchange_bridge::{InboundBridge, OutboundBridge};
use crate::placement::WorkerId;

use arrow::array::{RecordBatch, UInt32Array};
use arrow::compute::take;

pub use fv_streams_types::contract::{
    position_record, Acker, Chained, DeferredOperator, DeferredSink, EdgeId, EpochStaging, OpSnapshot, Operator, Out,
    Poll, Sink, Source, StateFile, TaskId,
};

/// A resolved out-route: the sender into each target's inbox, tagged with the edge id and the
/// target task (named when a send blocks).
type Outlet = (EdgeId, TaskId, SyncSender<(EdgeId, Msg)>);

/// What a task is doing right now — the runtime's own answer to "where is it stuck": every task
/// keeps one, the coordinator reads them (`Running::task_states`).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum TaskState {
    Starting = 0,
    /// A source inside `poll`.
    Polling = 1,
    /// A consumer waiting on its inbox.
    Receiving = 2,
    /// Inside the operator's or sink's `on_data`.
    Working = 3,
    /// Blocked sending into `target`'s inbox.
    Sending = 4,
    /// Taking or completing a barrier.
    Barrier = 5,
    /// A source past Eos, waiting for its last epoch's commit.
    CommitWait = 6,
    Done = 7,
}

impl TaskState {
    fn from_u8(v: u8) -> TaskState {
        match v {
            1 => TaskState::Polling,
            2 => TaskState::Receiving,
            3 => TaskState::Working,
            4 => TaskState::Sending,
            5 => TaskState::Barrier,
            6 => TaskState::CommitWait,
            7 => TaskState::Done,
            _ => TaskState::Starting,
        }
    }
}

/// One task's probe: its state, the target of a send, and when it entered the state.
pub struct Probe {
    state: std::sync::atomic::AtomicU8,
    target: std::sync::atomic::AtomicU32,
    since_ms: std::sync::atomic::AtomicI64,
}

impl Probe {
    fn new() -> Self {
        Probe {
            state: std::sync::atomic::AtomicU8::new(0),
            target: std::sync::atomic::AtomicU32::new(u32::MAX),
            since_ms: std::sync::atomic::AtomicI64::new(now_ms()),
        }
    }

    fn set(&self, state: TaskState, target: TaskId) {
        self.target.store(target, Ordering::Relaxed);
        self.state.store(state as u8, Ordering::Relaxed);
        self.since_ms.store(now_ms(), Ordering::Relaxed);
    }

    /// `(state, send target if sending, milliseconds in this state)`.
    pub fn read(&self) -> (TaskState, Option<TaskId>, i64) {
        let state = TaskState::from_u8(self.state.load(Ordering::Relaxed));
        let target = self.target.load(Ordering::Relaxed);
        let since = self.since_ms.load(Ordering::Relaxed);
        (
            state,
            (state == TaskState::Sending && target != u32::MAX).then_some(target),
            (now_ms() - since).max(0),
        )
    }
}
type OutRoute = (Route, Vec<Outlet>);

/// What flows on an edge.
#[derive(Clone, Debug)]
pub enum Msg {
    Data(RecordBatch),
    /// The time of `source` (a source task id) as seen by the sender; `None` = the source is idle.
    Watermark {
        source: TaskId,
        ts: Option<i64>,
    },
    Barrier {
        epoch: u64,
    },
    /// Processing time passes with nothing else to say: a source that polled nothing sends one
    /// (at most ten a second), and every task forwards at most ten a second. Operators with
    /// processing-time work (bounded idleness, ingest-time windows) advance on it.
    Tick {
        now_ms: i64,
    },
    /// The upstream is done (clean shutdown flows downstream in order).
    Eos,
}

/// Test-only structural equality (a `RecordBatch` compares by schema + columns), so seam tests can
/// `assert_eq!` a round-tripped message against what was sent.
#[cfg(test)]
impl PartialEq for Msg {
    fn eq(&self, other: &Msg) -> bool {
        match (self, other) {
            (Msg::Data(a), Msg::Data(b)) => a == b,
            (Msg::Watermark { source: s1, ts: t1 }, Msg::Watermark { source: s2, ts: t2 }) => (s1, t1) == (s2, t2),
            (Msg::Barrier { epoch: a }, Msg::Barrier { epoch: b }) => a == b,
            (Msg::Tick { now_ms: a }, Msg::Tick { now_ms: b }) => a == b,
            (Msg::Eos, Msg::Eos) => true,
            _ => false,
        }
    }
}

/// Hashes the routing columns of a batch into `out` (one `u64` per row) — the shuffle's key hash,
/// supplied by the planner (the engine's DataFusion hash, fixed-seed: a key's vnode is the same on
/// every run, every worker and every restart). The runtime never chooses a hash function.
pub type Hasher = Arc<dyn Fn(&[arrow::array::ArrayRef], &mut [u64]) + Send + Sync>;

/// Where a task's output goes.
#[derive(Clone)]
pub enum Route {
    /// Whole batches to one task.
    Forward(TaskId),
    /// Rows hashed on `columns` into vnodes; each target owns a contiguous vnode range.
    Shuffle {
        columns: Vec<String>,
        vnodes: u32,
        targets: Vec<(TaskId, std::ops::Range<u32>)>,
        hasher: Hasher,
    },
    /// The same batch to every target.
    Broadcast(Vec<TaskId>),
}

impl std::fmt::Debug for Route {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Route::Forward(t) => write!(f, "Forward({t})"),
            Route::Shuffle {
                columns,
                vnodes,
                targets,
                ..
            } => f
                .debug_struct("Shuffle")
                .field("columns", columns)
                .field("vnodes", vnodes)
                .field("targets", targets)
                .finish(),
            Route::Broadcast(ts) => write!(f, "Broadcast({ts:?})"),
        }
    }
}

impl Route {
    pub fn targets(&self) -> Vec<TaskId> {
        match self {
            Route::Forward(t) => vec![*t],
            Route::Shuffle { targets, .. } => targets.iter().map(|(t, _)| *t).collect(),
            Route::Broadcast(ts) => ts.clone(),
        }
    }
}

/// What the coordinator tells a source.
#[derive(Clone, Copy, Debug)]
pub enum Control {
    /// Take a barrier for `epoch`: record the position, forward the barrier.
    Barrier(u64),
    /// Every task reported `epoch`: commit the position recorded at its barrier.
    Commit(u64),
    /// Finish after the current poll: Eos flows downstream in order.
    Stop,
}

/// What the coordinator hears from tasks.
#[derive(Debug)]
pub enum Event {
    /// A source or operator task snapshotted for `epoch`.
    Snapshot { task: TaskId, epoch: u64, snap: OpSnapshot },
    /// A sink made everything before `epoch`'s barrier durable.
    Ack { task: TaskId, epoch: u64 },
    /// A source ran dry: its final position (see [`Source::position`]) — carried into every epoch
    /// completed after it, since a finished source takes no more barriers.
    Drained { task: TaskId, snap: OpSnapshot },
    /// A task finished (its inputs hit Eos, the source ran dry, or the coordinator stopped it),
    /// having used `cpu_ms` of CPU on its thread — where the build's cores went, by task.
    Finished { task: TaskId, cpu_ms: u64 },
    /// A source hit its strict data contract; it finishes right after.
    Failed { task: TaskId, error: String },
}

enum Node {
    Source(Box<dyn Source>),
    Operator(Box<dyn Operator>),
    /// An operator constructed at `start`, once its in-edges are known.
    Deferred(DeferredOperator),
    DeferredSink(DeferredSink),
    Sink(Box<dyn Sink>),
}

/// A sink that takes nothing: the placeholder while a node is swapped.
struct NoSink;
impl Sink for NoSink {
    fn on_data(&mut self, _batch: RecordBatch) {}
    fn on_barrier(&mut self, _epoch: u64) -> bool {
        true
    }
    fn on_eos(&mut self) {}
}

struct TaskDef {
    id: TaskId,
    node: Node,
    routes: Vec<Route>,
}

/// This worker's exchange endpoints for a multi-process run: one [`LinkSender`] per peer worker it
/// sends to, and one [`LinkReceiver`] per peer it receives from (established by the membership
/// handshake before `start`). [`Graph::start`] runs with none of this — a single-worker run where
/// every edge is a local `SyncSender`; the multi-process path is [`Graph::start_worker`], which
/// derives the actual edge wiring from the placement.
///
/// The mechanism is identical either way — an out-edge is a bounded `SyncSender`, drained by the
/// dest task (local) or by an `OutboundBridge` (remote); an `InboundBridge` delivers a decoded
/// `(edge, msg)` into the local dest inbox exactly as a local outlet would. Edge ids agree across
/// workers because every worker builds the same graph in the same order, so the `edge` on the wire
/// names the same logical edge — and thus the same dest task — on both ends.
pub struct WorkerExchange {
    /// The tokio runtime the bridges drive (`block_on`). Must be multi-threaded: its workers run
    /// the links' socket tasks while the bridge threads block on them.
    pub handle: Handle,
    /// A link to each peer worker this worker sends edges to, keyed by that worker's id.
    pub senders: HashMap<WorkerId, LinkSender>,
    /// A link from each peer this worker receives edges from; each delivers its frames as
    /// `(edge, msg)` into local inboxes, routed by the edge's target task.
    pub receivers: Vec<LinkReceiver>,
}

/// The graph under construction: add tasks, connect routes, then `start`.
#[derive(Default)]
pub struct Graph {
    tasks: Vec<TaskDef>,
    capacity: usize,
    barrier_every: Option<Duration>,
    /// The epoch this graph continues from (a restored run continues its checkpoint's numbering,
    /// so the newest manifest is always the newest epoch).
    first_epoch: u64,
}

impl Graph {
    pub fn new(edge_capacity: usize) -> Self {
        Graph {
            tasks: Vec::new(),
            capacity: edge_capacity.max(1),
            barrier_every: None,
            first_epoch: 0,
        }
    }

    /// Continue the epoch numbering from `epoch`: the first epoch issued is `epoch + 1`. A run
    /// restored from a checkpoint passes that checkpoint's epoch, so its own checkpoints always
    /// read as newer than the one it restored from.
    pub fn start_at_epoch(mut self, epoch: u64) -> Self {
        self.first_epoch = epoch;
        self
    }

    /// Issue an epoch (a barrier to every source) this often. Without it, only `stop` issues one.
    pub fn barrier_every(mut self, every: Duration) -> Self {
        self.barrier_every = Some(every).filter(|d| !d.is_zero());
        self
    }

    pub fn source(&mut self, s: Box<dyn Source>) -> TaskId {
        self.add(Node::Source(s))
    }

    pub fn operator(&mut self, o: Box<dyn Operator>) -> TaskId {
        self.add(Node::Operator(o))
    }

    /// An operator that needs to know its input edges (a join's sides): built at `start`.
    pub fn operator_with(&mut self, f: DeferredOperator) -> TaskId {
        self.add(Node::Deferred(f))
    }

    pub fn sink(&mut self, s: Box<dyn Sink>) -> TaskId {
        self.add(Node::Sink(s))
    }

    /// A sink built at start on its owning worker only (see [`DeferredSink`]).
    pub fn sink_with(&mut self, f: DeferredSink) -> TaskId {
        self.add(Node::DeferredSink(f))
    }

    fn add(&mut self, node: Node) -> TaskId {
        let id = self.tasks.len() as TaskId;
        self.tasks.push(TaskDef {
            id,
            node,
            routes: Vec::new(),
        });
        id
    }

    /// Route `from`'s output.
    pub fn route(&mut self, from: TaskId, route: Route) {
        self.tasks[from as usize].routes.push(route);
    }

    /// Chain `sink` into source task `task` (see [`Chained`]); hands the sink back when the task
    /// is not a source.
    pub fn chain_sink(&mut self, task: TaskId, sink: Box<dyn Sink>) -> Result<(), Box<dyn Sink>> {
        let Some(def) = self.tasks.get_mut(task as usize) else {
            return Err(sink);
        };
        if !matches!(def.node, Node::Source(_)) {
            return Err(sink);
        }
        let placeholder = Node::Sink(Box::new(NoSink));
        let Node::Source(source) = std::mem::replace(&mut def.node, placeholder) else {
            unreachable!()
        };
        def.node = Node::Source(Box::new(Chained { source, sink }));
        Ok(())
    }

    /// Spawn every task on this single-process worker; returns the running dataflow. Every edge is
    /// a local `SyncSender`. For one worker of a multi-process run, see [`Graph::start_worker`].
    pub fn start(self, events: SyncSender<Event>) -> Result<Running, String> {
        // A solo placement: every task on the one worker, no exchange — byte-identical to the
        // pre-Phase-4 runtime.
        let solo = vec![0 as WorkerId; self.tasks.len()];
        self.start_worker(events, &solo, 0, None)
    }

    /// Spawn only the tasks this worker owns and wire the graph's edges according to the placement
    /// `of_task` (task id → worker id): an edge between two co-located tasks is a local `SyncSender`;
    /// an edge from a local task to a task on another worker is drained by an `OutboundBridge` onto
    /// the link to that worker; an edge into a local task from a remote one arrives via an
    /// `InboundBridge` on the peer's link and is delivered to the local dest inbox. A task on no
    /// other worker is not run here. **Every worker builds the identical graph in the identical
    /// order**, so edge ids agree across workers and the `edge` on the wire names the same logical
    /// edge — hence the same dest task — on every worker. A solo placement (every task on `me`, no
    /// `exchange`) is exactly [`Graph::start`].
    pub fn start_worker(
        self,
        events: SyncSender<Event>,
        of_task: &[WorkerId],
        me: WorkerId,
        exchange: Option<WorkerExchange>,
    ) -> Result<Running, String> {
        let n = self.tasks.len();
        assert_eq!(of_task.len(), n, "placement must cover every task");
        let is_local = |id: TaskId| of_task[id as usize] == me;
        // An inbox only for a task this worker runs; a remote task has none here.
        let mut inbox_tx: Vec<Option<SyncSender<(EdgeId, Msg)>>> = Vec::with_capacity(n);
        let mut inbox_rx: Vec<Option<Receiver<(EdgeId, Msg)>>> = Vec::with_capacity(n);
        for id in 0..n as TaskId {
            if is_local(id) {
                let (tx, rx) = sync_channel(self.capacity);
                inbox_tx.push(Some(tx));
                inbox_rx.push(Some(rx));
            } else {
                inbox_tx.push(None);
                inbox_rx.push(None);
            }
        }
        // The exchange endpoints, unpacked so the solo case allocates nothing of note.
        let (handle, peer_senders, receivers) = match exchange {
            Some(x) => (Some(x.handle), x.senders, x.receivers),
            None => (None, HashMap::new(), Vec::new()),
        };
        let mut next_edge: EdgeId = 0;
        let mut in_edges: Vec<Vec<(EdgeId, TaskId)>> = vec![Vec::new(); n];
        // edge id -> its target task, so an inbound bridge routes a decoded `(edge, msg)` to the
        // right local inbox — the mirror of a local outlet's `SyncSender`.
        let mut edge_target: Vec<TaskId> = Vec::new();
        // per task: its outgoing routes resolved to (edge id, target inbox sender)
        let mut outs: Vec<Vec<OutRoute>> = vec![Vec::new(); n];
        let mut out_bridges: Vec<OutboundBridge> = Vec::new();
        for t in &self.tasks {
            // Edge ids are assigned for every route on every worker (identical order), so they agree
            // across workers; only a local upstream's outlets are built here (a remote task's routes
            // are wired on its own worker).
            let upstream_local = is_local(t.id);
            for r in &t.routes {
                let mut senders = Vec::new();
                for target in r.targets() {
                    let e = next_edge;
                    next_edge += 1;
                    edge_target.push(target); // index == e
                    in_edges[target as usize].push((e, t.id));
                    if !upstream_local {
                        continue;
                    }
                    // Local target: its inbox. Remote target: a bounded channel (its own
                    // backpressure) drained by an outbound bridge onto the link to that worker.
                    let tx = if is_local(target) {
                        inbox_tx[target as usize].clone().expect("local inbox")
                    } else {
                        let h = handle.as_ref().expect("a remote edge needs an exchange handle");
                        let link = peer_senders
                            .get(&of_task[target as usize])
                            .expect("a link to the target's worker")
                            .clone();
                        let (btx, brx) = sync_channel(self.capacity);
                        out_bridges.push(OutboundBridge::spawn(h.clone(), link, brx));
                        btx
                    };
                    senders.push((e, target, tx));
                }
                if upstream_local {
                    outs[t.id as usize].push((r.clone(), senders));
                }
            }
        }
        // Inbound: one bridge per peer link, delivering `(edge, msg)` into the local dest inbox.
        // The closures hold clones of the local inboxes, so a task with only remote inputs keeps its
        // inbox open until its peer's Eos arrives (then the link closes and the bridge ends).
        let mut in_bridges: Vec<InboundBridge> = Vec::new();
        if let Some(h) = handle.as_ref() {
            for link_rx in receivers {
                let inboxes = inbox_tx.clone();
                let targets = edge_target.clone();
                in_bridges.push(InboundBridge::spawn(h.clone(), link_rx, move |edge, msg| {
                    if let Some(inbox) = targets.get(edge as usize).and_then(|&t| inboxes[t as usize].as_ref()) {
                        let _ = inbox.send((edge, msg));
                    }
                }));
            }
        }
        drop(inbox_tx);
        let local_tasks: Vec<TaskId> = (0..n as TaskId).filter(|&id| is_local(id)).collect();
        let inner = Arc::new(Inner {
            ctl: Mutex::new(Vec::new()),
            epoch: AtomicU64::new(self.first_epoch),
            stopping: AtomicBool::new(false),
            n_tasks: n,
        });
        let probes: Arc<Vec<Probe>> = Arc::new((0..n).map(|_| Probe::new()).collect());
        // Build every local task's node FIRST. A deferred operator's constructor can refuse (a
        // checkpoint it cannot restore), and that must fail the start before any task thread
        // exists — never orphan tasks already spawned with no coordinator to stop them. Same order
        // as the spawn loop below, so task and edge ids agree across workers exactly as before.
        let mut built = Vec::with_capacity(local_tasks.len());
        for t in self.tasks.into_iter() {
            let i = t.id as usize;
            if !is_local(t.id) {
                continue; // a task on another worker: not run in this process.
            }
            let rx = inbox_rx[i].take().expect("inbox rx");
            let out_routes = std::mem::take(&mut outs[i]);
            let ins_with_upstream = std::mem::take(&mut in_edges[i]);
            let ins: Vec<EdgeId> = ins_with_upstream.iter().map(|(e, _)| *e).collect();
            let node = match t.node {
                Node::Deferred(build) => Node::Operator(build(t.id, &ins_with_upstream)?),
                Node::DeferredSink(build) => Node::Sink(build(t.id)?),
                other => other,
            };
            built.push((
                TaskDef {
                    id: t.id,
                    node,
                    routes: t.routes,
                },
                rx,
                out_routes,
                ins,
            ));
        }
        let mut handles = Vec::with_capacity(built.len());
        for (t, rx, out_routes, ins) in built {
            let events = events.clone();
            let task_probes = Arc::clone(&probes);
            let id = t.id;
            // sources and sinks hear the coordinator over their own control channel (sources
            // take barriers and commit positions; sinks commit transactions).
            let ctl = if matches!(t.node, Node::Source(_) | Node::Sink(_)) {
                let (tx, rx) = channel();
                inner.ctl.lock().expect("control senders").push(tx);
                Some(rx)
            } else {
                None
            };
            handles.push(
                std::thread::Builder::new()
                    .name(format!("fv-task-{id}"))
                    .spawn(move || {
                        let mut task = Task {
                            id,
                            node: t.node,
                            rx,
                            ctl,
                            out_routes,
                            in_edges: ins,
                            last_tick: None,
                            probes: task_probes,
                            events,
                            wm: HashMap::new(),
                            last_wm: None,
                            align: Align::default(),
                            eos_seen: HashSet::new(),
                        };
                        task.run();
                    })
                    .expect("spawn task"),
            );
        }
        // the periodic epochs: the runtime's own ticker, until `stop`.
        let ticker = self.barrier_every.map(|every| {
            let inner = Arc::clone(&inner);
            std::thread::Builder::new()
                .name("fv-epochs".into())
                .spawn(move || {
                    let nap = Duration::from_millis(10).min(every);
                    let mut next = Instant::now() + every;
                    while !inner.stopping.load(Ordering::Relaxed) {
                        if Instant::now() >= next {
                            inner.barrier();
                            next = Instant::now() + every;
                        }
                        std::thread::sleep(nap);
                    }
                })
                .expect("spawn epochs")
        });
        Ok(Running {
            inner,
            handles: Mutex::new(handles),
            probes,
            ticker: Mutex::new(ticker),
            out_bridges: Mutex::new(out_bridges),
            in_bridges: Mutex::new(in_bridges),
            local_tasks,
        })
    }
}

/// The coordinator's side of the sources: their control channels and the epoch counter.
struct Inner {
    ctl: Mutex<Vec<Sender<Control>>>,
    epoch: AtomicU64,
    stopping: AtomicBool,
    n_tasks: usize,
}

impl Inner {
    fn send_all(&self, c: Control) {
        for tx in self.ctl.lock().expect("control senders").iter() {
            let _ = tx.send(c); // a finished source dropped its receiver: nothing to tell it
        }
    }

    fn barrier(&self) -> u64 {
        let epoch = self.epoch.fetch_add(1, Ordering::SeqCst) + 1;
        self.send_all(Control::Barrier(epoch));
        epoch
    }
}

/// A started dataflow: issue epochs, commit them, stop it, join it.
pub struct Running {
    inner: Arc<Inner>,
    handles: Mutex<Vec<std::thread::JoinHandle<()>>>,
    ticker: Mutex<Option<std::thread::JoinHandle<()>>>,
    /// Every task's state probe, by task id.
    probes: Arc<Vec<Probe>>,
    /// The exchange bridges (empty for a single-worker run). Joined after the tasks: an outbound
    /// bridge ends when its task drops the edge; an inbound bridge when its peer link closes.
    out_bridges: Mutex<Vec<OutboundBridge>>,
    in_bridges: Mutex<Vec<InboundBridge>>,
    /// The tasks this worker actually runs — `0..n` for a solo run, this worker's own tasks for one
    /// worker of a multi-process run. The tracker waits on exactly these.
    local_tasks: Vec<TaskId>,
}

impl Running {
    /// Every task's current state: `(task, state, send target, milliseconds in that state)`.
    pub fn task_states(&self) -> Vec<(TaskId, TaskState, Option<TaskId>, i64)> {
        self.probes
            .iter()
            .enumerate()
            .map(|(i, p)| {
                let (s, t, ms) = p.read();
                (i as TaskId, s, t, ms)
            })
            .collect()
    }

    pub fn n_tasks(&self) -> usize {
        self.inner.n_tasks
    }

    /// A tracker for the tasks this worker runs (every task for a solo run).
    pub fn tracker(&self) -> EpochTracker {
        EpochTracker::new_for(self.local_tasks.clone())
    }

    /// Issue the next epoch now (a barrier to every source); returns it.
    pub fn barrier(&self) -> u64 {
        self.inner.barrier()
    }

    /// Issue a barrier for a *specific* epoch to this worker's sources, setting the local epoch to
    /// it. A joined worker calls this when the leader injects epoch `e`, so every worker's sources
    /// stamp the same epoch — the barrier then aligns across a cross-worker shuffle. (The leader
    /// uses [`Running::barrier`] to pick `e`, then tells the workers to inject it.)
    pub fn inject_epoch(&self, epoch: u64) {
        self.inner.epoch.store(epoch, Ordering::SeqCst);
        self.inner.send_all(Control::Barrier(epoch));
    }

    /// The last epoch issued.
    pub fn epoch(&self) -> u64 {
        self.inner.epoch.load(Ordering::SeqCst)
    }

    /// Every task reported `epoch`: the sources may commit the positions they recorded at it.
    pub fn commit(&self, epoch: u64) {
        self.inner.send_all(Control::Commit(epoch));
    }

    /// Tell the sources to stop (Eos flows behind the last data) — WITHOUT issuing a barrier, unlike
    /// [`Running::stop`]. A clustered run uses this at shutdown: the final epoch was already issued
    /// and committed by the leader-coordinated round, so a fresh uncommitted barrier here would leave
    /// the sources waiting for a commit that never comes.
    pub fn signal_stop(&self) {
        self.inner.stopping.store(true, Ordering::SeqCst);
        self.inner.send_all(Control::Stop);
    }

    /// Stop: no more periodic epochs, one final barrier, then `Stop` to every source (Eos flows
    /// downstream in order behind the barrier). Returns the final epoch. Idempotent.
    pub fn stop(&self) -> u64 {
        if self.inner.stopping.swap(true, Ordering::SeqCst) {
            return self.epoch();
        }
        if let Some(t) = self.ticker.lock().expect("ticker").take() {
            let _ = t.join();
        }
        let epoch = self.inner.barrier();
        self.inner.send_all(Control::Stop);
        epoch
    }

    /// Drive `events` until every task finished: a completed epoch is committed to the sources
    /// and handed to `on_epoch` with its snapshots; a task's failure stops the graph and is the
    /// result. Call after `stop` (or for a graph whose sources are bounded).
    pub fn finish(
        &self,
        events: &Receiver<Event>,
        tracker: &mut EpochTracker,
        mut on_epoch: impl FnMut(Completed) -> Result<(), String>,
    ) -> Result<(), String> {
        let mut failure: Option<String> = None;
        while !tracker.all_finished() {
            match events.recv_timeout(Duration::from_millis(50)) {
                Ok(e) => {
                    if let Event::Failed { error, .. } = &e {
                        failure.get_or_insert_with(|| error.clone());
                        self.stop();
                        // a sink waiting for its last epoch's commit must not wait forever: the
                        // commit is never coming (uncommitted = replay), so close the control
                        // channels and let it finish.
                        self.inner.ctl.lock().expect("control senders").clear();
                    }
                    // the coordinator makes the epoch durable in `on_epoch` (its checkpoint)
                    // BEFORE anyone commits on it; an error there ends the run.
                    for c in tracker.on_event(&e) {
                        let epoch = c.epoch;
                        match on_epoch(c) {
                            Ok(()) => self.commit(epoch),
                            Err(err) => {
                                // a checkpoint that failed is the run's failure, said now: the
                                // tasks waiting for its commit are released (uncommitted = replay).
                                eprintln!("dataflow: epoch {} not committed — {err}", epoch);
                                failure.get_or_insert(err);
                                self.stop();
                                self.inner.ctl.lock().expect("control senders").clear();
                            }
                        }
                    }
                }
                Err(RecvTimeoutError::Timeout) => {}
                Err(RecvTimeoutError::Disconnected) => break,
            }
        }
        match failure {
            Some(f) => Err(f),
            None => Ok(()),
        }
    }

    /// Join every task. Closes the control channels first, so a source waiting for a final commit
    /// that will never come (the coordinator gave up) still exits — uncommitted means replay.
    pub fn join(&self) {
        self.inner.stopping.store(true, Ordering::SeqCst);
        if let Some(t) = self.ticker.lock().expect("ticker").take() {
            let _ = t.join();
        }
        self.inner.ctl.lock().expect("control senders").clear();
        let handles: Vec<_> = self.handles.lock().expect("handles").drain(..).collect();
        for h in handles {
            let _ = h.join();
        }
        // Tasks are done: an outbound bridge's edge is now dropped (so it drains and ends, closing
        // its link). Join those first.
        for b in self.out_bridges.lock().expect("out bridges").drain(..) {
            b.join();
        }
        // Every local task has finished, so all data and Eos are delivered: signal the inbound
        // bridges to stop rather than wait for the peer to close its link. Waiting would deadlock —
        // two workers would each hold their sockets open until the other closed first.
        let in_bridges: Vec<InboundBridge> = self.in_bridges.lock().expect("in bridges").drain(..).collect();
        for b in &in_bridges {
            b.signal_stop();
        }
        for b in in_bridges {
            b.join();
        }
    }
}

/// An epoch every task reported (or had finished before): its snapshots, by task.
#[derive(Debug, Clone)]
pub struct Completed {
    pub epoch: u64,
    /// `(task, snapshot)` from every source and operator that snapshotted, in task order.
    pub snapshots: Vec<(TaskId, OpSnapshot)>,
}

/// The coordinator's book: which tasks reported which epoch, which tasks finished. Epochs
/// complete in order. An epoch completes when every *expected* task has reported it (or finished
/// before it) — for a single-process run that is every task `0..n`; for one worker of a multi-
/// process run it is only the tasks placed on this worker (the leader aggregates across workers).
#[derive(Debug)]
pub struct EpochTracker {
    /// The tasks this coordinator waits on — `0..n` for a solo run, a worker's own tasks otherwise.
    expected: Vec<TaskId>,
    finished: HashSet<TaskId>,
    /// A drained source's final position, added to every epoch it no longer reports.
    drained: HashMap<TaskId, OpSnapshot>,
    pending: BTreeMap<u64, HashMap<TaskId, Option<OpSnapshot>>>,
    cpu_ms: HashMap<TaskId, u64>,
}

impl EpochTracker {
    /// A tracker for a contiguous task set `0..n_tasks` (the single-process run).
    pub fn new(n_tasks: usize) -> Self {
        EpochTracker::new_for((0..n_tasks as TaskId).collect())
    }

    /// A tracker for an explicit set of task ids (one worker's own tasks).
    pub fn new_for(expected: Vec<TaskId>) -> Self {
        EpochTracker {
            expected,
            finished: HashSet::new(),
            drained: HashMap::new(),
            pending: BTreeMap::new(),
            cpu_ms: HashMap::new(),
        }
    }

    /// Book one event; returns the epochs it completed, in order.
    pub fn on_event(&mut self, e: &Event) -> Vec<Completed> {
        match e {
            Event::Snapshot { task, epoch, snap } => {
                self.pending
                    .entry(*epoch)
                    .or_default()
                    .insert(*task, Some(snap.clone()));
            }
            Event::Ack { task, epoch } => {
                self.pending.entry(*epoch).or_default().insert(*task, None);
            }
            Event::Finished { task, cpu_ms } => {
                self.finished.insert(*task);
                self.cpu_ms.insert(*task, *cpu_ms);
            }
            Event::Drained { task, snap } => {
                self.drained.insert(*task, snap.clone());
            }
            Event::Failed { .. } => {}
        }
        let mut done = Vec::new();
        while let Some((&epoch, reported)) = self.pending.iter().next() {
            let complete = self
                .expected
                .iter()
                .all(|t| reported.contains_key(t) || self.finished.contains(t));
            if !complete {
                break;
            }
            let reported = self.pending.remove(&epoch).expect("present");
            // a drained source reported nothing for this epoch: its final position stands in
            let drained: Vec<(TaskId, OpSnapshot)> = self
                .drained
                .iter()
                .filter(|(t, _)| !reported.contains_key(t))
                .map(|(t, snap)| (*t, snap.clone()))
                .collect();
            let mut snapshots: Vec<(TaskId, OpSnapshot)> = reported
                .into_iter()
                .filter_map(|(t, b)| b.map(|b| (t, b)))
                .chain(drained)
                .collect();
            snapshots.sort_by_key(|(t, _)| *t);
            done.push(Completed { epoch, snapshots });
        }
        done
    }

    pub fn all_finished(&self) -> bool {
        self.expected.iter().all(|t| self.finished.contains(t))
    }

    pub fn finished(&self) -> usize {
        self.finished.len()
    }

    /// The CPU each finished task used on its thread, in milliseconds.
    pub fn cpu_ms(&self) -> &HashMap<TaskId, u64> {
        &self.cpu_ms
    }
}

/// The CPU time of every thread in this process so far, summed by thread name (milliseconds):
/// the engine's tasks by their `fv-task-N` names, librdkafka's `rdk:*` threads, the rest. Empty
/// where `/proc` is absent. This is where a build's cores go, at a glance.
pub fn process_cpu_by_thread() -> BTreeMap<String, u64> {
    let mut out = BTreeMap::new();
    let Ok(tasks) = std::fs::read_dir("/proc/self/task") else {
        return out;
    };
    for t in tasks.flatten() {
        let dir = t.path();
        let Ok(comm) = std::fs::read_to_string(dir.join("comm")) else {
            continue;
        };
        let Ok(stat) = std::fs::read_to_string(dir.join("stat")) else {
            continue;
        };
        if let Some(ms) = cpu_ms_of_stat(&stat) {
            *out.entry(comm.trim().to_string()).or_default() += ms;
        }
    }
    out
}

/// utime + stime of a `/proc/<pid>/task/<tid>/stat` line, in milliseconds.
fn cpu_ms_of_stat(stat: &str) -> Option<u64> {
    // the comm field may contain spaces and parentheses: split after the last ')'.
    let rest = &stat[stat.rfind(')')? + 1..];
    let fields: Vec<&str> = rest.split_whitespace().collect();
    // fields after the comm: state(0) ppid(1) … utime(11) stime(12)
    let utime: u64 = fields.get(11)?.parse().ok()?;
    let stime: u64 = fields.get(12)?.parse().ok()?;
    Some((utime + stime) * 10)
}

/// The CPU time this thread has used so far (user + system), from `/proc/thread-self/stat`;
/// `None` where that file does not exist. Linux reports it in USER_HZ ticks, which is 100 on
/// every architecture regardless of the kernel's HZ.
pub fn thread_cpu_ms() -> Option<u64> {
    let stat = std::fs::read_to_string("/proc/thread-self/stat").ok()?;
    // the comm field may contain spaces and parentheses: split after the last ')'.
    let rest = &stat[stat.rfind(')')? + 1..];
    let fields: Vec<&str> = rest.split_whitespace().collect();
    // fields after the comm: state(0) ppid(1) … utime(11) stime(12)
    let utime: u64 = fields.get(11)?.parse().ok()?;
    let stime: u64 = fields.get(12)?.parse().ok()?;
    Some((utime + stime) * 10)
}

/// Alignment state: which input edges delivered the current barrier, and the data held back on
/// them until the rest catch up.
#[derive(Default)]
struct Align {
    epoch: Option<u64>,
    delivered: HashSet<EdgeId>,
    held: VecDeque<(EdgeId, Msg)>,
}

struct Task {
    id: TaskId,
    node: Node,
    rx: Receiver<(EdgeId, Msg)>,
    /// The coordinator's channel (sources only).
    ctl: Option<Receiver<Control>>,
    out_routes: Vec<OutRoute>,
    in_edges: Vec<EdgeId>,
    events: SyncSender<Event>,
    /// When this task last forwarded a tick (at most ten a second, whatever arrives).
    last_tick: Option<Instant>,
    /// Every task's probe; this task writes its own (`probes[id]`).
    probes: Arc<Vec<Probe>>,
    /// latest watermark per (input edge, source): `None` = that source is idle.
    wm: HashMap<(EdgeId, TaskId), Option<i64>>,
    last_wm: Option<i64>,
    align: Align,
    eos_seen: HashSet<EdgeId>,
}

impl Task {
    fn run(&mut self) {
        // a panic inside an operator, source or sink ends the BUILD, loudly: the coordinator
        // hears `Failed` and stops the graph. (Left uncaught, the thread dies, its downstream sees
        // a closed edge and finishes, its upstream keeps polling into nothing, and the pipeline
        // sits there consuming — the runtime's stall report found exactly that.)
        let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match &self.node {
            Node::Source(_) => self.run_source(),
            _ => self.run_consumer(),
        }));
        if let Err(payload) = outcome {
            let what = payload
                .downcast_ref::<String>()
                .cloned()
                .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
                .unwrap_or_else(|| "panic".into());
            let _ = self.events.send(Event::Failed {
                task: self.id,
                error: format!("task {} panicked: {what}", self.id),
            });
            // No Eos: Eos means a clean end. The edges out of this task close with it, and a
            // downstream task whose input closes with this edge still open reports the loss
            // (`input_lost`) — the run fails rather than finishing on partial input.
        }
        self.state(TaskState::Done, u32::MAX);
        let _ = self.events.send(Event::Finished {
            task: self.id,
            cpu_ms: thread_cpu_ms().unwrap_or(0),
        });
    }

    fn state(&self, s: TaskState, target: TaskId) {
        self.probes[self.id as usize].set(s, target);
    }

    fn run_source(&mut self) {
        let ctl = self.ctl.take().expect("a source has a control channel");
        // the last barrier taken, and the last epoch committed (a commit can arrive while the
        // source is still polling — it must not be waited for twice).
        let mut last_epoch: u64 = 0;
        let mut committed: u64 = 0;
        let mut last_tick = Instant::now();
        let mut stopping = false;
        let mut failed = false;
        let mut drained = false;
        loop {
            // the coordinator first: barriers and commits, in the order issued.
            while let Ok(c) = ctl.try_recv() {
                match c {
                    Control::Barrier(epoch) => {
                        last_epoch = epoch;
                        self.source_barrier(epoch);
                    }
                    Control::Commit(epoch) => {
                        committed = committed.max(epoch);
                        self.source_commit(epoch);
                    }
                    Control::Stop => stopping = true,
                }
            }
            if stopping {
                break;
            }
            let Node::Source(s) = &mut self.node else {
                unreachable!()
            };
            let mut out = Out::default();
            self.probes[self.id as usize].set(TaskState::Polling, u32::MAX);
            let poll = s.poll(&mut out);
            let quiet = out.batches.is_empty();
            self.emit(out);
            if quiet && last_tick.elapsed() >= Duration::from_millis(100) {
                last_tick = Instant::now();
                self.broadcast(Msg::Tick { now_ms: now_ms() });
            }
            match poll {
                Poll::Again => {}
                Poll::Idle => std::thread::sleep(Duration::from_millis(5)),
                Poll::Done => {
                    drained = true;
                    break;
                }
                Poll::Failed(error) => {
                    let _ = self.events.send(Event::Failed { task: self.id, error });
                    failed = true;
                    break;
                }
            }
        }
        if let Node::Source(s) = &mut self.node {
            if drained {
                // the final position, for every checkpoint completed after this source ended
                let pos = s.position();
                if !pos.is_empty() {
                    let _ = self.events.send(Event::Drained {
                        task: self.id,
                        snap: OpSnapshot::whole(pos),
                    });
                }
            }
            s.on_stop();
        }
        if !failed {
            self.broadcast(Msg::Eos); // a clean end only: a failed source's edges just close
        }
        // the last barrier taken commits once every sink acked it: wait for that commit, unless it
        // already came. A closed control channel means the coordinator is done with us
        // (uncommitted = replay). Barriers issued after Eos are not taken: downstream treats this
        // edge as aligned, and the epoch completes once this task finishes.
        if !failed && last_epoch > committed {
            self.state(TaskState::CommitWait, u32::MAX);
            while let Ok(c) = ctl.recv() {
                if let Control::Commit(epoch) = c {
                    self.source_commit(epoch);
                    if epoch >= last_epoch {
                        break;
                    }
                }
            }
        }
    }

    fn source_barrier(&mut self, epoch: u64) {
        let Node::Source(s) = &mut self.node else {
            unreachable!()
        };
        let snap = OpSnapshot::whole(s.on_barrier(epoch));
        let _ = self.events.send(Event::Snapshot {
            task: self.id,
            epoch,
            snap,
        });
        self.broadcast(Msg::Barrier { epoch });
    }

    fn source_commit(&mut self, epoch: u64) {
        let Node::Source(s) = &mut self.node else {
            unreachable!()
        };
        s.on_commit(epoch);
    }

    /// The inbox closed before every in-edge delivered its Eos: an upstream task — or the worker it
    /// runs on — is gone without ending its stream. A clean end always arrives as Eos on every edge
    /// (a task broadcasts Eos before its senders drop; a peer's outbound bridge sends it before the
    /// link closes), so a closed inbox with an edge still open is a loss, never a drain: the run
    /// fails instead of finishing on partial input. Without this a worker that died during a
    /// bounded drain looked like "all upstreams done" and the leader stopped clean, its share of
    /// the output silently missing. Reached at most once per task — never on the data path.
    #[cold]
    fn input_lost(&mut self) {
        let open: Vec<EdgeId> = self
            .in_edges
            .iter()
            .copied()
            .filter(|e| !self.eos_seen.contains(e))
            .collect();
        if open.is_empty() {
            return; // every edge ended cleanly; the senders simply dropped after their Eos.
        }
        let _ = self.events.send(Event::Failed {
            task: self.id,
            error: format!(
                "task {}: input closed before end-of-stream on edge(s) {open:?} — an upstream task, or the \
                 worker running it, is gone",
                self.id
            ),
        });
    }

    fn run_consumer(&mut self) {
        let ctl = self.ctl.take();
        let mut last_epoch: u64 = 0;
        let mut committed: u64 = 0;
        loop {
            // the coordinator's commits reach a sink between messages (ticks keep the inbox
            // moving at least ten times a second, so a commit waits at most that long).
            if let Some(ctl) = &ctl {
                while let Ok(c) = ctl.try_recv() {
                    if let Control::Commit(epoch) = c {
                        committed = committed.max(epoch);
                        if let Node::Sink(s) = &mut self.node {
                            s.on_commit(epoch);
                        }
                    }
                }
            }
            self.state(TaskState::Receiving, u32::MAX);
            let Ok((edge, msg)) = self.rx.recv() else {
                self.input_lost();
                break;
            };
            if let Msg::Barrier { epoch } = &msg {
                last_epoch = last_epoch.max(*epoch);
            }
            if self.receive(edge, msg) {
                break;
            }
        }
        // a sink's last epoch commits once every task reported it: wait for that commit, so an
        // exactly-once transaction is never left open at a clean stop (a closed control channel
        // means the coordinator is done with us: uncommitted = replay).
        if let (Some(ctl), Node::Sink(_)) = (&ctl, &self.node) {
            if last_epoch > committed {
                self.state(TaskState::CommitWait, u32::MAX);
                while let Ok(c) = ctl.recv() {
                    if let Control::Commit(epoch) = c {
                        if let Node::Sink(s) = &mut self.node {
                            s.on_commit(epoch);
                        }
                        if epoch >= last_epoch {
                            break;
                        }
                    }
                }
            }
        }
    }

    /// Hold or handle: an edge that already delivered the pending barrier holds EVERYTHING it
    /// sends after it — data, watermarks, its next barrier, Eos — in order, until the other edges
    /// catch up. (A later barrier let through early would restart alignment and lose the pending
    /// epoch: the source behind it then waits forever for a commit that cannot come.)
    fn receive(&mut self, edge: EdgeId, msg: Msg) -> bool {
        if self.align.epoch.is_some() && self.align.delivered.contains(&edge) {
            self.align.held.push_back((edge, msg));
            return false;
        }
        self.handle(edge, msg)
    }

    /// Handle one message; `true` when the task is done (every input reached Eos).
    fn handle(&mut self, edge: EdgeId, msg: Msg) -> bool {
        match msg {
            Msg::Data(b) => {
                self.state(TaskState::Working, u32::MAX);
                let mut out = Out::default();
                match &mut self.node {
                    Node::Operator(o) => o.on_data(edge, b, &mut out),
                    Node::Sink(s) => s.on_data(b),
                    Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
                }
                self.emit(out);
            }
            Msg::Watermark { source, ts } => {
                self.wm.insert((edge, source), ts);
                // the combined time is the minimum over the active inputs — and it holds until
                // EVERY input edge has spoken once (a watermark or the idle marker): the minimum
                // over the inputs heard so far would let a fast upstream fire windows the slow
                // one still has rows for.
                let every_edge_spoke = self.in_edges.iter().all(|e| self.wm.keys().any(|(edge, _)| edge == e));
                let active: Vec<i64> = self.wm.values().filter_map(|t| *t).collect();
                let combined = if active.is_empty() || !every_edge_spoke {
                    None
                } else {
                    active.iter().copied().min()
                };
                if let Some(wm) = combined {
                    if self.last_wm.is_none_or(|l| wm > l) {
                        self.last_wm = Some(wm);
                        let mut out = Out::default();
                        if let Node::Operator(o) = &mut self.node {
                            o.on_watermark(wm, now_ms(), &mut out);
                            // an operator's own time is the combined input time it has processed.
                            if out.watermark.is_none() {
                                out.watermark = Some(Some(wm));
                            }
                        }
                        self.emit(out);
                    }
                } else if active.is_empty() && (self.last_wm.is_some() || every_edge_spoke) {
                    // every upstream source is idle: say so downstream.
                    self.broadcast(Msg::Watermark {
                        source: self.id,
                        ts: None,
                    });
                }
            }
            Msg::Barrier { epoch } => {
                let a = &mut self.align;
                if a.epoch != Some(epoch) {
                    a.epoch = Some(epoch);
                    a.delivered.clear();
                }
                a.delivered.insert(edge);
                if self.aligned() {
                    return self.complete_barrier(epoch);
                }
            }
            Msg::Tick { now_ms } => {
                let mut out = Out::default();
                if let Node::Operator(o) = &mut self.node {
                    o.on_tick(now_ms, &mut out);
                }
                self.emit(out);
                if self.last_tick.is_none_or(|t| t.elapsed() >= Duration::from_millis(100)) {
                    self.last_tick = Some(Instant::now());
                    self.broadcast(Msg::Tick { now_ms });
                }
            }
            Msg::Eos => {
                self.eos_seen.insert(edge);
                // an edge at Eos never delivers another barrier: it counts as aligned — and a
                // barrier the other edges already delivered completes NOW (a source that finished
                // between two epochs never took the later one; the pending epoch must not wait
                // for it forever).
                if let Some(epoch) = self.align.epoch {
                    if self.aligned() && self.complete_barrier(epoch) {
                        return true;
                    }
                }
                if self.in_edges.iter().all(|e| self.eos_seen.contains(e)) {
                    let mut out = Out::default();
                    match &mut self.node {
                        Node::Operator(o) => o.on_eos(&mut out),
                        Node::Sink(s) => s.on_eos(),
                        Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
                    }
                    self.emit(out);
                    self.broadcast(Msg::Eos);
                    return true;
                }
            }
        }
        false
    }

    /// Every input edge delivered the pending barrier, or is at Eos.
    fn aligned(&self) -> bool {
        self.in_edges
            .iter()
            .all(|e| self.align.delivered.contains(e) || self.eos_seen.contains(e))
    }

    /// The pending barrier aligned: snapshot (or flush and ack), forward it, then replay what
    /// alignment held — in arrival order and through the hold check, so a held barrier for the
    /// next epoch starts the next alignment and what followed it on that edge is held again.
    /// `true` when the replay reached the end of every input.
    fn complete_barrier(&mut self, epoch: u64) -> bool {
        self.state(TaskState::Barrier, u32::MAX);
        let mut out = Out::default();
        match &mut self.node {
            Node::Operator(o) => match o.on_barrier(epoch, &mut out) {
                Ok(snap) => {
                    let _ = self.events.send(Event::Snapshot {
                        task: self.id,
                        epoch,
                        snap,
                    });
                }
                Err(error) => {
                    // The checkpoint could not be taken: the build fails with the operator's own
                    // diagnostic — the end state a panic in the operator reaches (see `run`),
                    // reached without one. The barrier is not forwarded; this task is done.
                    let _ = self.events.send(Event::Failed {
                        task: self.id,
                        error: format!("task {}: checkpoint at epoch {epoch}: {error}", self.id),
                    });
                    return true; // no Eos (a clean end); downstream sees the edge close
                }
            },
            Node::Sink(s) => {
                if s.on_barrier(epoch) {
                    let _ = self.events.send(Event::Ack { task: self.id, epoch });
                }
            }
            Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
        }
        self.emit(out);
        self.broadcast(Msg::Barrier { epoch });
        self.align.epoch = None;
        self.align.delivered.clear();
        let held: Vec<(EdgeId, Msg)> = self.align.held.drain(..).collect();
        for (e, m) in held {
            if self.receive(e, m) {
                return true;
            }
        }
        false
    }

    /// Route an operator's outputs along every out-route.
    fn emit(&mut self, out: Out) {
        for b in out.batches {
            for (route, senders) in &self.out_routes {
                match route {
                    Route::Forward(_) | Route::Broadcast(_) => {
                        for (e, target, tx) in senders {
                            self.probes[self.id as usize].set(TaskState::Sending, *target);
                            let _ = tx.send((*e, Msg::Data(b.clone())));
                        }
                    }
                    Route::Shuffle {
                        columns,
                        vnodes,
                        targets,
                        hasher,
                    } => {
                        let parts = shuffle(&b, columns, *vnodes, targets, hasher.as_ref());
                        for ((e, target, tx), part) in senders.iter().zip(parts) {
                            if let Some(p) = part {
                                self.probes[self.id as usize].set(TaskState::Sending, *target);
                                let _ = tx.send((*e, Msg::Data(p)));
                            }
                        }
                    }
                }
            }
        }
        if let Some(ts) = out.watermark {
            self.broadcast(Msg::Watermark { source: self.id, ts });
        }
        self.state(TaskState::Working, u32::MAX);
    }

    fn broadcast(&self, msg: Msg) {
        for (_, senders) in &self.out_routes {
            for (e, target, tx) in senders {
                self.probes[self.id as usize].set(TaskState::Sending, *target);
                let _ = tx.send((*e, msg.clone()));
            }
        }
        self.state(TaskState::Working, u32::MAX);
    }
}

fn now_ms() -> i64 {
    chrono::Utc::now().timestamp_millis()
}

/// Split a batch by vnode: `hash(columns) % vnodes` → the target whose range holds it. Returns one
/// optional sub-batch per target, in target order. Missing routing columns hash as null.
pub fn shuffle(
    b: &RecordBatch,
    columns: &[String],
    vnodes: u32,
    targets: &[(TaskId, std::ops::Range<u32>)],
    hasher: &(dyn Fn(&[arrow::array::ArrayRef], &mut [u64]) + Send + Sync),
) -> Vec<Option<RecordBatch>> {
    let n = b.num_rows();
    let vnode_of: Vec<u32> = if columns.is_empty() {
        vec![0; n]
    } else {
        let cols: Vec<arrow::array::ArrayRef> = columns
            .iter()
            .map(|c| match b.column_by_name(c) {
                Some(a) => Arc::clone(a),
                None => arrow::array::new_null_array(&arrow::datatypes::DataType::Null, n),
            })
            .collect();
        let mut hashes = vec![0u64; n];
        hasher(&cols, &mut hashes);
        hashes.iter().map(|h| (h % vnodes as u64) as u32).collect()
    };
    targets
        .iter()
        .map(|(_, range)| {
            let idx: Vec<u32> = (0..n as u32)
                .filter(|i| range.contains(&vnode_of[*i as usize]))
                .collect();
            if idx.is_empty() {
                None
            } else if idx.len() == n {
                Some(b.clone())
            } else {
                let idx = UInt32Array::from(idx);
                let cols: Vec<arrow::array::ArrayRef> = b
                    .columns()
                    .iter()
                    .map(|c| take(c.as_ref(), &idx, None).expect("take"))
                    .collect();
                Some(RecordBatch::try_new(b.schema(), cols).expect("shuffled batch"))
            }
        })
        .collect()
}

/// Contiguous vnode ranges for `tasks` tasks over `vnodes` vnodes.
pub fn vnode_ranges(vnodes: u32, tasks: u32) -> Vec<std::ops::Range<u32>> {
    let tasks = tasks.max(1).min(vnodes.max(1));
    (0..tasks)
        .map(|t| {
            let start = vnodes * t / tasks;
            let end = vnodes * (t + 1) / tasks;
            start..end
        })
        .collect()
}

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

    /// A deterministic hash of each row's routing cells (their display strings), for the tests.
    pub(super) fn test_hasher() -> Hasher {
        Arc::new(|cols: &[arrow::array::ArrayRef], out: &mut [u64]| {
            use std::hash::{Hash, Hasher as _};
            for (i, slot) in out.iter_mut().enumerate() {
                let mut h = std::collections::hash_map::DefaultHasher::new();
                for c in cols {
                    arrow::util::display::array_value_to_string(c, i)
                        .unwrap_or_default()
                        .hash(&mut h);
                }
                *slot = h.finish();
            }
        })
    }
    use arrow::array::Array;
    use arrow::array::{Float64Array, Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use std::sync::atomic::AtomicUsize;

    fn batch(keys: &[&str], ts: &[i64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("k", DataType::Utf8, true),
            Field::new("ts", DataType::Int64, true),
            Field::new("v", DataType::Float64, true),
        ]));
        let k: StringArray = keys.to_vec().into();
        let t: Int64Array = ts.to_vec().into();
        let v: Float64Array = ts.iter().map(|x| *x as f64).collect::<Vec<_>>().into();
        RecordBatch::try_new(schema, vec![Arc::new(k), Arc::new(t), Arc::new(v)]).unwrap()
    }

    /// A bounded source: emits its batches one per poll (with a watermark), then Done.
    struct VecSource {
        batches: VecDeque<RecordBatch>,
    }
    impl Source for VecSource {
        fn poll(&mut self, out: &mut Out) -> Poll {
            match self.batches.pop_front() {
                Some(b) => {
                    let ts = b
                        .column_by_name("ts")
                        .and_then(|c| arrow::compute::max(c.as_any().downcast_ref::<Int64Array>().unwrap()));
                    out.push(b);
                    out.watermark = Some(ts);
                    Poll::Again
                }
                None => Poll::Done,
            }
        }
        fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
            format!("src@{epoch}").into_bytes()
        }
        fn position(&mut self) -> Vec<u8> {
            b"src@done".to_vec()
        }
    }

    /// Counts rows per key and emits the count on every watermark; snapshots its counts.
    #[derive(Default)]
    struct CountOp {
        counts: HashMap<String, i64>,
        fired: Vec<i64>,
    }
    impl Operator for CountOp {
        fn on_data(&mut self, _edge: EdgeId, batch: RecordBatch, _out: &mut Out) {
            let k = batch.column_by_name("k").unwrap();
            let k = k.as_any().downcast_ref::<StringArray>().unwrap();
            for i in 0..k.len() {
                *self.counts.entry(k.value(i).to_string()).or_default() += 1;
            }
        }
        fn on_watermark(&mut self, wm: i64, _now: i64, _out: &mut Out) {
            self.fired.push(wm);
        }
        fn on_barrier(&mut self, epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
            Ok(OpSnapshot::whole(
                format!("op@{epoch}:{}", self.counts.values().sum::<i64>()).into_bytes(),
            ))
        }
        fn on_eos(&mut self, out: &mut Out) {
            // emit the counts as a batch: one row per key.
            let mut keys: Vec<&String> = self.counts.keys().collect();
            keys.sort();
            let ks: StringArray = keys.iter().map(|k| k.as_str()).collect::<Vec<_>>().into();
            let ns: Int64Array = keys.iter().map(|k| self.counts[*k]).collect::<Vec<_>>().into();
            let schema = Arc::new(Schema::new(vec![
                Field::new("k", DataType::Utf8, true),
                Field::new("n", DataType::Int64, true),
            ]));
            out.push(RecordBatch::try_new(schema, vec![Arc::new(ks), Arc::new(ns)]).unwrap());
        }
    }

    /// Collects everything it receives.
    struct CollectSink {
        rows: Arc<Mutex<Vec<(String, i64)>>>,
        barriers: Arc<AtomicUsize>,
    }
    impl Sink for CollectSink {
        fn on_data(&mut self, batch: RecordBatch) {
            let k = batch.column_by_name("k").unwrap();
            let k = k.as_any().downcast_ref::<StringArray>().unwrap();
            let n = batch.column_by_name("n").unwrap();
            let n = n.as_any().downcast_ref::<Int64Array>().unwrap();
            let mut rows = self.rows.lock().unwrap();
            for i in 0..k.len() {
                rows.push((k.value(i).to_string(), n.value(i)));
            }
        }
        fn on_barrier(&mut self, _epoch: u64) -> bool {
            self.barriers.fetch_add(1, Ordering::Relaxed);
            true
        }
        fn on_eos(&mut self) {}
    }

    /// An unbounded source (idle until stopped) that remembers what it was told to commit.
    struct IdleSource {
        polls: usize,
        fail_at: Option<usize>,
        commits: Arc<Mutex<Vec<u64>>>,
        barriers: Arc<Mutex<Vec<u64>>>,
    }
    impl Source for IdleSource {
        fn poll(&mut self, out: &mut Out) -> Poll {
            self.polls += 1;
            if self.fail_at == Some(self.polls) {
                return Poll::Failed("boom".into());
            }
            out.push(batch(&["a"], &[self.polls as i64]));
            Poll::Idle
        }
        fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
            self.barriers.lock().unwrap().push(epoch);
            epoch.to_le_bytes().to_vec()
        }
        fn on_commit(&mut self, epoch: u64) {
            self.commits.lock().unwrap().push(epoch);
        }
    }

    /// A sink that records the order of what it sees: `b<epoch>` for a barrier, `e` for Eos, `d`
    /// for data.
    struct TraceSink {
        trace: Arc<Mutex<Vec<String>>>,
    }
    impl Sink for TraceSink {
        fn on_data(&mut self, _b: RecordBatch) {
            self.trace.lock().unwrap().push("d".into());
        }
        fn on_barrier(&mut self, epoch: u64) -> bool {
            self.trace.lock().unwrap().push(format!("b{epoch}"));
            true
        }
        fn on_eos(&mut self) {
            self.trace.lock().unwrap().push("e".into());
        }
    }

    #[test]
    fn a_shuffled_pipeline_counts_every_key_exactly_once_across_two_operator_tasks() {
        let mut g = Graph::new(16);
        let data: Vec<RecordBatch> = (0..20)
            .map(|i| batch(&["a", "b", "c", "d"], &[i * 10, i * 10 + 1, i * 10 + 2, i * 10 + 3]))
            .collect();
        let src = g.source(Box::new(VecSource { batches: data.into() }));
        let op0 = g.operator(Box::new(CountOp::default()));
        let op1 = g.operator(Box::new(CountOp::default()));
        let rows = Arc::new(Mutex::new(Vec::new()));
        let barriers = Arc::new(AtomicUsize::new(0));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::clone(&rows),
            barriers: Arc::clone(&barriers),
        }));
        let ranges = vnode_ranges(64, 2);
        g.route(
            src,
            Route::Shuffle {
                columns: vec!["k".into()],
                vnodes: 64,
                targets: vec![(op0, ranges[0].clone()), (op1, ranges[1].clone())],
                hasher: test_hasher(),
            },
        );
        g.route(op0, Route::Forward(sink));
        g.route(op1, Route::Forward(sink));
        let (etx, erx) = sync_channel(1024);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
        running.join();
        let mut got: Vec<(String, i64)> = rows.lock().unwrap().clone();
        got.sort();
        assert_eq!(
            got,
            vec![("a".into(), 20), ("b".into(), 20), ("c".into(), 20), ("d".into(), 20)],
            "every key counted exactly once, on exactly one task"
        );
        assert_eq!(tracker.finished(), 4);
    }

    #[test]
    fn a_peer_that_dies_before_eos_fails_the_worker_instead_of_finishing_it() {
        // Worker 1 runs only the sink, fed by a remote edge from a peer. The peer dies (its runtime
        // — and so its sockets — go away) without ever sending Eos. The link closes, the inbound
        // bridge ends, the sink's inbox closes with the edge still open: that must surface as the
        // run's FAILURE naming the edge, not as a clean finish on partial input (the leader used to
        // stop clean and the dead worker's share of the output was silently missing).
        let rt_peer = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap();
        let rt = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap();
        let listener = rt.block_on(tokio::net::TcpListener::bind("127.0.0.1:0")).unwrap();
        let addr = listener.local_addr().unwrap();
        let srv = rt.spawn(async move { fv_streams_exchange::Link::accept(&listener, 1, 1 << 20).await.unwrap() });
        let peer = rt_peer
            .block_on(fv_streams_exchange::Link::connect(addr, 1, 1 << 20))
            .unwrap();
        let (_w1_send, w1_from_peer) = rt.block_on(srv).unwrap().into_split();

        let rows = Arc::new(Mutex::new(Vec::new()));
        let mut g = Graph::new(16);
        let src = g.source(Box::new(VecSource {
            batches: Vec::new().into(),
        }));
        let op = g.operator(Box::new(CountOp::default()));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::clone(&rows),
            barriers: Arc::new(AtomicUsize::new(0)),
        }));
        g.route(src, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let of_task = vec![0u32, 0, 1]; // src + op on the (dead) peer, the sink here
        let (etx, erx) = sync_channel(1024);
        let r1 = g
            .start_worker(
                etx,
                &of_task,
                1,
                Some(WorkerExchange {
                    handle: rt.handle().clone(),
                    senders: HashMap::new(),
                    receivers: vec![w1_from_peer],
                }),
            )
            .expect("start worker 1");
        // the peer dies mid-stream: nothing sent, no Eos.
        drop(peer);
        drop(rt_peer);
        let mut t1 = r1.tracker();
        let err = r1
            .finish(&erx, &mut t1, |_| Ok(()))
            .expect_err("a lost upstream fails the run");
        assert!(
            err.contains("input closed before end-of-stream") && err.contains("is gone"),
            "names the loss: {err}"
        );
        assert!(t1.all_finished(), "the sink task still ended (no hang)");
        r1.join();
        assert!(
            rows.lock().unwrap().is_empty(),
            "no partial output was presented as final"
        );
        drop(rt);
    }

    #[test]
    fn two_workers_split_by_placement_produce_the_same_result_as_one_process() {
        // The per-worker gate: the SAME `src -> op -> sink` graph, built identically on two workers
        // and split by placement — src+op on worker 0, sink on worker 1 — running as two independent
        // coordinators in one process, connected by a loopback Link. The op->sink edge is the only
        // cross-worker edge; the aggregate, watermarks and the final Eos travel as Arrow-IPC/postcard
        // frames over TCP and land in the sink exactly as a local edge would. Each worker instantiates
        // ONLY its own tasks (proven by each tracker's finished count) and the result is identical to
        // the single-process run — placement is transparent to correctness.
        let rt = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap();
        let handle = rt.handle().clone();

        // One link between worker 0 and worker 1: w0 sends the op->sink edge, w1 receives it.
        let (w0_to_w1, w1_from_w0) = handle.block_on(async {
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();
            let srv =
                tokio::spawn(async move { fv_streams_exchange::Link::accept(&listener, 1, 1 << 20).await.unwrap() });
            let a = fv_streams_exchange::Link::connect(addr, 1, 1 << 20).await.unwrap();
            let b = srv.await.unwrap();
            let (a_send, _a_recv) = a.into_split();
            let (_b_send, b_recv) = b.into_split();
            (a_send, b_recv)
        });

        // Every worker builds the identical graph (deterministic task/edge ids). The sink collects
        // into a shared buffer so the test can read worker 1's output.
        let rows = Arc::new(Mutex::new(Vec::new()));
        let build = || {
            let mut g = Graph::new(16);
            let data: Vec<RecordBatch> = (0..20)
                .map(|i| batch(&["a", "b", "c", "d"], &[i * 10, i * 10 + 1, i * 10 + 2, i * 10 + 3]))
                .collect();
            let src = g.source(Box::new(VecSource { batches: data.into() }));
            let op = g.operator(Box::new(CountOp::default()));
            let sink = g.sink(Box::new(CollectSink {
                rows: Arc::clone(&rows),
                barriers: Arc::new(AtomicUsize::new(0)),
            }));
            g.route(src, Route::Forward(op));
            g.route(op, Route::Forward(sink));
            (g, src, op, sink)
        };
        let (g0, _s, _o, _k) = build();
        let (g1, _s1, _o1, _k1) = build();
        // Placement: src (0) and op (1) on worker 0; sink (2) on worker 1.
        let of_task = vec![0u32, 0, 1];

        // Worker 1 (the sink): receives the op->sink edge from worker 0; sends nothing.
        let (etx1, erx1) = sync_channel(1024);
        let r1 = g1
            .start_worker(
                etx1,
                &of_task,
                1,
                Some(WorkerExchange {
                    handle: handle.clone(),
                    senders: HashMap::new(),
                    receivers: vec![w1_from_w0],
                }),
            )
            .expect("start worker 1");
        let w1 = std::thread::spawn(move || {
            let mut t1 = r1.tracker();
            r1.finish(&erx1, &mut t1, |_| Ok(())).unwrap();
            r1.join();
            t1.finished()
        });

        // Worker 0 (source + operator): sends the op->sink edge to worker 1.
        let (etx0, erx0) = sync_channel(1024);
        let mut senders = HashMap::new();
        senders.insert(1u32, w0_to_w1);
        let r0 = g0
            .start_worker(
                etx0,
                &of_task,
                0,
                Some(WorkerExchange {
                    handle: handle.clone(),
                    senders,
                    receivers: vec![],
                }),
            )
            .expect("start worker 0");
        let mut t0 = r0.tracker();
        r0.finish(&erx0, &mut t0, |_| Ok(())).unwrap();
        r0.join();
        let w1_finished = w1.join().unwrap();

        // Each worker ran only its own tasks.
        assert_eq!(t0.finished(), 2, "worker 0 ran exactly src + op");
        assert_eq!(w1_finished, 1, "worker 1 ran exactly the sink");
        // The output is identical to the single-process run.
        let mut got: Vec<(String, i64)> = rows.lock().unwrap().clone();
        got.sort();
        assert_eq!(
            got,
            vec![("a".into(), 20), ("b".into(), 20), ("c".into(), 20), ("d".into(), 20)],
            "the aggregate crossed the worker boundary intact"
        );
        drop(rt);
    }

    #[test]
    fn the_combined_watermark_holds_until_every_input_has_spoken() {
        // REGRESSION (the crash e2e): a fast source's watermark reached the operator before the
        // slow source's first one; the minimum over the inputs heard so far fired windows the
        // slow source still had rows for. Every input edge must speak before time moves.
        struct Src {
            wm: i64,
            delay_ms: u64,
            emitted: bool,
        }
        impl Source for Src {
            fn poll(&mut self, out: &mut Out) -> Poll {
                if self.emitted {
                    return Poll::Done;
                }
                std::thread::sleep(Duration::from_millis(self.delay_ms));
                self.emitted = true;
                out.push(batch(&["x"], &[self.wm]));
                out.watermark = Some(Some(self.wm));
                Poll::Again
            }
            fn on_barrier(&mut self, _epoch: u64) -> Vec<u8> {
                Vec::new()
            }
        }
        struct WmOp(Arc<Mutex<Vec<i64>>>);
        impl Operator for WmOp {
            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
            fn on_watermark(&mut self, wm: i64, _n: i64, _o: &mut Out) {
                self.0.lock().unwrap().push(wm);
            }
            fn on_barrier(&mut self, _epoch: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
                Ok(OpSnapshot::whole(Vec::new()))
            }
            fn on_eos(&mut self, _o: &mut Out) {}
        }
        let mut g = Graph::new(4);
        let fast = g.source(Box::new(Src {
            wm: 1_000,
            delay_ms: 0,
            emitted: false,
        }));
        let slow = g.source(Box::new(Src {
            wm: 5,
            delay_ms: 150,
            emitted: false,
        }));
        let seen = Arc::new(Mutex::new(Vec::new()));
        let op = g.operator(Box::new(WmOp(Arc::clone(&seen))));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(AtomicUsize::new(0)),
        }));
        g.route(fast, Route::Forward(op));
        g.route(slow, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, erx) = sync_channel(64);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
        running.join();
        let seen = seen.lock().unwrap().clone();
        assert_eq!(
            seen,
            vec![5],
            "time moved only once both inputs had spoken, to their minimum"
        );
    }

    #[test]
    fn watermarks_take_the_minimum_over_two_sources_and_barriers_align() {
        // two sources → one operator: the operator's watermark firings must never exceed the
        // slower source's time, and a barrier from both must align before the snapshot.
        struct SlowSource {
            n: usize,
            step: i64,
            emitted: usize,
        }
        impl Source for SlowSource {
            fn poll(&mut self, out: &mut Out) -> Poll {
                if self.emitted >= self.n {
                    return Poll::Done;
                }
                self.emitted += 1;
                let t = self.emitted as i64 * self.step;
                out.push(batch(&["x"], &[t]));
                out.watermark = Some(Some(t));
                std::thread::sleep(Duration::from_millis(2));
                Poll::Again
            }
            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
                format!("{epoch}").into_bytes()
            }
        }
        struct WmOp {
            seen: Arc<Mutex<Vec<i64>>>,
            aligned: Arc<AtomicUsize>,
        }
        impl Operator for WmOp {
            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
            fn on_watermark(&mut self, wm: i64, _n: i64, _o: &mut Out) {
                self.seen.lock().unwrap().push(wm);
            }
            fn on_barrier(&mut self, _epoch: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
                self.aligned.fetch_add(1, Ordering::Relaxed);
                Ok(OpSnapshot::whole(Vec::new()))
            }
            fn on_eos(&mut self, _o: &mut Out) {}
        }
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(20));
        let fast = g.source(Box::new(SlowSource {
            n: 200,
            step: 10,
            emitted: 0,
        }));
        let slow = g.source(Box::new(SlowSource {
            n: 200,
            step: 1,
            emitted: 0,
        }));
        let seen = Arc::new(Mutex::new(Vec::new()));
        let aligned = Arc::new(AtomicUsize::new(0));
        let op = g.operator(Box::new(WmOp {
            seen: Arc::clone(&seen),
            aligned: Arc::clone(&aligned),
        }));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(AtomicUsize::new(0)),
        }));
        g.route(fast, Route::Forward(op));
        g.route(slow, Route::Forward(op));
        g.route(op, Route::Broadcast(vec![sink]));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let mut completed: Vec<Completed> = Vec::new();
        running
            .finish(&erx, &mut tracker, |c| {
                completed.push(c);
                Ok(())
            })
            .unwrap();
        running.join();
        let seen = seen.lock().unwrap().clone();
        assert!(!seen.is_empty());
        assert!(seen.windows(2).all(|w| w[1] > w[0]), "monotonic");
        assert!(
            *seen.last().unwrap() <= 200,
            "never past the slow source's time: {:?}",
            seen.last()
        );
        // the operator snapshotted once per aligned barrier, and every completed epoch carries
        // the operator's snapshot plus one per source still running at the time.
        assert!(!completed.is_empty(), "the ticker issued epochs");
        assert_eq!(
            completed
                .iter()
                .filter(|c| c.snapshots.iter().any(|(t, _)| *t == op))
                .count(),
            aligned.load(Ordering::Relaxed)
        );
        let epochs: Vec<u64> = completed.iter().map(|c| c.epoch).collect();
        assert!(
            epochs.windows(2).all(|w| w[1] == w[0] + 1),
            "epochs complete in order: {epochs:?}"
        );
        assert!(completed.iter().all(|c| c.snapshots.len() <= 3));
    }

    #[test]
    fn a_fast_edge_delivering_several_barriers_ahead_of_a_slow_one_still_aligns_every_epoch() {
        // REGRESSION: the fast source takes many barriers while the slow source sits in one poll;
        // the operator's fast edge therefore delivers barriers e, e+1, e+2 before the slow edge
        // delivers e. Each must wait its turn: every epoch aligns, snapshots, and commits, and the
        // fast source (waiting for its last commit) exits — the run must not hang.
        struct PacedSource {
            polls: usize,
            left: usize,
            nap: Duration,
        }
        impl Source for PacedSource {
            fn poll(&mut self, out: &mut Out) -> Poll {
                if self.left == 0 {
                    return Poll::Done;
                }
                self.left -= 1;
                self.polls += 1;
                std::thread::sleep(self.nap);
                out.push(batch(&["p"], &[self.polls as i64]));
                Poll::Again
            }
            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
                epoch.to_le_bytes().to_vec()
            }
        }
        let mut g = Graph::new(2).barrier_every(Duration::from_millis(5));
        let fast = g.source(Box::new(PacedSource {
            polls: 0,
            left: 400,
            nap: Duration::from_millis(1),
        }));
        let slow = g.source(Box::new(PacedSource {
            polls: 0,
            left: 6,
            nap: Duration::from_millis(60),
        }));
        let op = g.operator(Box::new(CountOp::default()));
        let rows = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::clone(&rows),
            barriers: Arc::new(AtomicUsize::new(0)),
        }));
        g.route(fast, Route::Forward(op));
        g.route(slow, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let mut completed: Vec<Completed> = Vec::new();
        // a watchdog: past the deadline it clears the sources' control channels, which releases a
        // source stuck waiting for a commit — `finish` then returns and the test reports the hang
        // instead of blocking the whole suite forever.
        let inner = Arc::clone(&running.inner);
        let done = Arc::new(AtomicBool::new(false));
        let hung = Arc::new(AtomicBool::new(false));
        let watchdog = std::thread::spawn({
            let done = Arc::clone(&done);
            let hung = Arc::clone(&hung);
            move || {
                let deadline = Instant::now() + Duration::from_secs(20);
                while !done.load(Ordering::Relaxed) {
                    if Instant::now() > deadline {
                        hung.store(true, Ordering::Relaxed);
                        inner.ctl.lock().unwrap().clear();
                        return;
                    }
                    std::thread::sleep(Duration::from_millis(20));
                }
            }
        });
        running
            .finish(&erx, &mut tracker, |c| {
                completed.push(c);
                Ok(())
            })
            .unwrap();
        running.join();
        done.store(true, Ordering::Relaxed);
        watchdog.join().unwrap();
        assert!(
            !hung.load(Ordering::Relaxed),
            "the dataflow hung: alignment lost an epoch"
        );
        assert_eq!(rows.lock().unwrap().as_slice(), &[("p".to_string(), 406)]);
        let epochs: Vec<u64> = completed.iter().map(|c| c.epoch).collect();
        assert!(
            epochs.len() >= 5,
            "several epochs completed while the slow source lagged: {epochs:?}"
        );
        assert!(epochs.windows(2).all(|w| w[1] == w[0] + 1), "in order: {epochs:?}");
        assert!(
            completed.iter().all(|c| c.snapshots.iter().any(|(t, _)| *t == op)),
            "every completed epoch carries the operator's aligned snapshot"
        );
    }

    #[test]
    fn a_source_finishing_between_epochs_does_not_stall_the_pending_alignment() {
        // REGRESSION: the bounded source takes no barrier for the epoch issued during its last,
        // slow poll and finishes; the live source did take it, so the operator holds that epoch
        // pending on the finished source's edge. Its Eos must complete the alignment, or every
        // later epoch waits behind it and the live source can never commit.
        struct SlowThenDone {
            polls: usize,
        }
        impl Source for SlowThenDone {
            fn poll(&mut self, out: &mut Out) -> Poll {
                self.polls += 1;
                if self.polls > 2 {
                    return Poll::Done;
                }
                std::thread::sleep(Duration::from_millis(40));
                out.push(batch(&["s"], &[self.polls as i64]));
                Poll::Again
            }
            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
                epoch.to_le_bytes().to_vec()
            }
        }
        let commits = Arc::new(Mutex::new(Vec::new()));
        let barriers = Arc::new(Mutex::new(Vec::new()));
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(5));
        let live = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::clone(&commits),
            barriers: Arc::clone(&barriers),
        }));
        let bounded = g.source(Box::new(SlowThenDone { polls: 0 }));
        let op = g.operator(Box::new(CountOp::default()));
        let rows = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(CollectSink {
            rows: Arc::clone(&rows),
            barriers: Arc::new(AtomicUsize::new(0)),
        }));
        g.route(live, Route::Forward(op));
        g.route(bounded, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let mut completed: Vec<u64> = Vec::new();
        let inner = Arc::clone(&running.inner);
        let done = Arc::new(AtomicBool::new(false));
        let hung = Arc::new(AtomicBool::new(false));
        let watchdog = std::thread::spawn({
            let done = Arc::clone(&done);
            let hung = Arc::clone(&hung);
            move || {
                let deadline = Instant::now() + Duration::from_secs(20);
                while !done.load(Ordering::Relaxed) {
                    if Instant::now() > deadline {
                        hung.store(true, Ordering::Relaxed);
                        inner.ctl.lock().unwrap().clear();
                        return;
                    }
                    std::thread::sleep(Duration::from_millis(20));
                }
            }
        });
        // let the bounded source finish and a few more epochs pass, then stop the live one.
        let until = Instant::now() + Duration::from_millis(250);
        while Instant::now() < until {
            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
                for c in tracker.on_event(&e) {
                    running.commit(c.epoch);
                    completed.push(c.epoch);
                }
            }
        }
        let last = running.stop();
        running
            .finish(&erx, &mut tracker, |c| {
                completed.push(c.epoch);
                Ok(())
            })
            .unwrap();
        running.join();
        done.store(true, Ordering::Relaxed);
        watchdog.join().unwrap();
        assert!(
            !hung.load(Ordering::Relaxed),
            "the pending epoch stalled on the finished source"
        );
        assert!(
            completed.windows(2).all(|w| w[1] == w[0] + 1),
            "in order: {completed:?}"
        );
        assert_eq!(completed.last(), Some(&last));
        assert_eq!(
            commits.lock().unwrap().last(),
            Some(&last),
            "the live source committed the final epoch"
        );
        let taken = barriers.lock().unwrap().len();
        assert!(
            taken > 10,
            "the live source kept taking epochs after the other finished: {taken}"
        );
        let got = rows.lock().unwrap().clone();
        assert_eq!(got.iter().find(|(k, _)| k == "s").map(|(_, n)| *n), Some(2));
    }

    #[test]
    fn a_graph_started_at_an_epoch_issues_the_next_ones() {
        // a restored run continues its checkpoint's numbering: restored from epoch 7, its first
        // barrier is 8, so its checkpoints always read as newer than the one it restored from.
        let barriers = Arc::new(Mutex::new(Vec::new()));
        let mut g = Graph::new(4).start_at_epoch(7);
        let src = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::clone(&barriers),
        }));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::new(Mutex::new(Vec::new())),
        }));
        g.route(src, Route::Forward(sink));
        let (etx, erx) = sync_channel(64);
        let running = g.start(etx).expect("start");
        assert_eq!(running.epoch(), 7);
        assert_eq!(running.barrier(), 8);
        let last = running.stop();
        assert_eq!(last, 9);
        let mut tracker = running.tracker();
        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
        running.join();
        assert_eq!(barriers.lock().unwrap().as_slice(), &[8, 9]);
    }

    #[test]
    fn epochs_complete_when_every_sink_acked_and_the_sources_hear_the_commit() {
        let commits = Arc::new(Mutex::new(Vec::new()));
        let barriers = Arc::new(Mutex::new(Vec::new()));
        let trace = Arc::new(Mutex::new(Vec::new()));
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let src = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::clone(&commits),
            barriers: Arc::clone(&barriers),
        }));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::clone(&trace),
        }));
        g.route(src, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let mut completed = Vec::new();
        // commit epochs while it runs, as the build loop does.
        let deadline = Instant::now() + Duration::from_millis(120);
        while Instant::now() < deadline {
            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
                for c in tracker.on_event(&e) {
                    running.commit(c.epoch);
                    completed.push(c.epoch);
                }
            }
        }
        let last = running.stop();
        running
            .finish(&erx, &mut tracker, |c| {
                completed.push(c.epoch);
                Ok(())
            })
            .unwrap();
        running.join();
        let commits = commits.lock().unwrap().clone();
        let barriers = barriers.lock().unwrap().clone();
        assert!(barriers.len() >= 3, "the ticker issued epochs: {barriers:?}");
        assert_eq!(barriers.last(), Some(&last), "stop issued the final barrier");
        assert_eq!(
            commits, completed,
            "the source hears exactly the completed epochs, in order"
        );
        assert_eq!(
            commits.last(),
            Some(&last),
            "the final epoch committed before the source exited"
        );
        let trace = trace.lock().unwrap().clone();
        let eos = trace.iter().position(|t| t == "e").expect("eos reached the sink");
        assert_eq!(
            trace[eos - 1],
            format!("b{last}"),
            "the final barrier arrives right before Eos: {trace:?}"
        );
        assert!(trace[..eos].iter().any(|t| t == "d"));
    }

    #[test]
    fn an_operator_panic_ends_the_run_failed_and_every_task_finishes() {
        // the third batch panics inside the operator: the coordinator hears the panic as the
        // build's failure, stops the sources, and every task still finishes.
        struct PanicOp {
            seen: usize,
        }
        impl Operator for PanicOp {
            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {
                self.seen += 1;
                if self.seen == 3 {
                    panic!("interleave: types differ");
                }
            }
            fn on_watermark(&mut self, _wm: i64, _n: i64, _o: &mut Out) {}
            fn on_barrier(&mut self, _e: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
                Ok(OpSnapshot::whole(Vec::new()))
            }
            fn on_eos(&mut self, _o: &mut Out) {}
        }
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let src = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(Mutex::new(Vec::new())),
        }));
        let op = g.operator(Box::new(PanicOp { seen: 0 }));
        let trace = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::clone(&trace),
        }));
        g.route(src, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
        running.join();
        assert!(err.contains("task 1 panicked") && err.contains("interleave"), "{err}");
        assert!(
            tracker.all_finished(),
            "the panic stopped the source and the sink finished"
        );
        let (_, state, _, _) = running.task_states()[op as usize];
        assert_eq!(state, TaskState::Done);
    }

    #[test]
    fn an_operator_that_cannot_checkpoint_fails_the_run_with_its_diagnostic_and_no_panic() {
        // the operator's barrier returns an error (a full disk under its spill dir): the run fails
        // with that message — the end state a panic reaches, without one, and the diagnostic intact.
        struct NoDiskOp;
        impl Operator for NoDiskOp {
            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
            fn on_watermark(&mut self, _wm: i64, _n: i64, _o: &mut Out) {}
            fn on_barrier(&mut self, _e: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
                Err("state io: write checkpoint file /spill/.ckpt/s1-hot.arrow: No space left on device".into())
            }
            fn on_eos(&mut self, _o: &mut Out) {}
        }
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let src = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(Mutex::new(Vec::new())),
        }));
        let op = g.operator(Box::new(NoDiskOp));
        let trace = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::clone(&trace),
        }));
        g.route(src, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
        running.join();
        assert!(
            err.contains("checkpoint at epoch") && err.contains("No space left"),
            "{err}"
        );
        assert!(!err.contains("panicked"), "a returned error, not a caught panic: {err}");
        assert!(
            tracker.all_finished(),
            "the failure stopped the source and the sink finished"
        );
        let (_, state, _, _) = running.task_states()[op as usize];
        assert_eq!(state, TaskState::Done);
    }

    #[test]
    fn a_deferred_operator_that_cannot_be_built_fails_start_before_any_task_runs() {
        // a restore the operator refuses is reported by `start` itself: no thread is spawned, so
        // nothing is left running with no coordinator to stop it.
        use std::sync::atomic::AtomicUsize;
        struct CountingSource(Arc<AtomicUsize>);
        impl Source for CountingSource {
            fn poll(&mut self, _out: &mut Out) -> Poll {
                self.0.fetch_add(1, Ordering::Relaxed);
                Poll::Idle
            }
            fn on_barrier(&mut self, _epoch: u64) -> Vec<u8> {
                Vec::new()
            }
        }
        let polls = Arc::new(AtomicUsize::new(0));
        let mut g = Graph::new(4);
        let src = g.source(Box::new(CountingSource(Arc::clone(&polls))));
        let op = g.operator_with(Box::new(|task, _in_edges| {
            Err::<Box<dyn Operator>, String>(format!(
                "stage 1 streamJoin task {task}: cannot restore its checkpoint — foreign manifest"
            ))
        }));
        let trace = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::clone(&trace),
        }));
        g.route(src, Route::Forward(op));
        g.route(op, Route::Forward(sink));
        let (etx, _erx) = sync_channel(64);
        let err = g.start(etx).err().expect("start should have refused the operator");
        assert!(
            err.contains("cannot restore its checkpoint") && err.contains("foreign manifest"),
            "{err}"
        );
        std::thread::sleep(Duration::from_millis(50));
        assert_eq!(
            polls.load(Ordering::Relaxed),
            0,
            "no task ran: the source was never polled"
        );
        assert!(trace.lock().unwrap().is_empty(), "the sink saw nothing");
    }

    #[test]
    fn a_failed_source_ends_the_run_with_its_error() {
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let src = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: Some(3),
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(Mutex::new(Vec::new())),
        }));
        let other = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(Mutex::new(Vec::new())),
        }));
        let trace = Arc::new(Mutex::new(Vec::new()));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::clone(&trace),
        }));
        g.route(src, Route::Forward(sink));
        g.route(other, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
        running.join();
        assert_eq!(err, "boom");
        assert!(tracker.all_finished(), "the failure stopped the other source too");
        // Eos means a clean end: the failed source's edge just closed, so the sink ended on a lost
        // input — it never saw `on_eos`, and so never treated uncommitted output as final.
        assert_ne!(trace.lock().unwrap().last().map(String::as_str), Some("e"));
    }

    #[test]
    fn a_drained_sources_final_position_rides_every_later_epoch() {
        // A source that ran dry takes no more barriers, so the epochs completed after it would
        // name no position for it — and a restore would replay it from the start. Its final
        // position (`Drained`) stands in for every such epoch.
        let mut t = EpochTracker::new(3); // 0 source, 1 operator, 2 sink
        let pos = OpSnapshot::whole(b"src@done".to_vec());
        assert!(t
            .on_event(&Event::Drained {
                task: 0,
                snap: pos.clone()
            })
            .is_empty());
        assert!(t.on_event(&Event::Finished { task: 0, cpu_ms: 0 }).is_empty());
        for epoch in 1..=2 {
            assert!(t
                .on_event(&Event::Snapshot {
                    task: 1,
                    epoch,
                    snap: OpSnapshot::whole(vec![1]),
                })
                .is_empty());
            let done = t.on_event(&Event::Ack { task: 2, epoch });
            assert_eq!(done.len(), 1);
            let tasks: Vec<TaskId> = done[0].snapshots.iter().map(|(t, _)| *t).collect();
            assert_eq!(tasks, vec![0, 1], "epoch {epoch}: the drained source's position is in");
            assert_eq!(done[0].snapshots[0].1.head, pos.head);
        }
        // a source that did report an epoch keeps its own snapshot (no double entry)
        let mut t = EpochTracker::new(2);
        t.on_event(&Event::Snapshot {
            task: 0,
            epoch: 1,
            snap: OpSnapshot::whole(b"at@1".to_vec()),
        });
        t.on_event(&Event::Drained {
            task: 0,
            snap: pos.clone(),
        });
        t.on_event(&Event::Finished { task: 0, cpu_ms: 0 });
        let done = t.on_event(&Event::Ack { task: 1, epoch: 1 });
        assert_eq!(done[0].snapshots.len(), 1);
        assert_eq!(done[0].snapshots[0].1.head, b"at@1".to_vec());
    }

    #[test]
    fn a_source_that_drains_before_the_first_barrier_still_has_its_position_checkpointed() {
        // Two sources into one sink: one drains at once (a bounded input that ran out before any
        // checkpoint — the cluster crash e2e's leader), the other keeps the run alive so barriers
        // keep coming. Every completed epoch must carry the drained source's final position.
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let src = g.source(Box::new(VecSource {
            batches: vec![batch(&["a"], &[1])].into(),
        }));
        let other = g.source(Box::new(IdleSource {
            polls: 0,
            fail_at: None,
            commits: Arc::new(Mutex::new(Vec::new())),
            barriers: Arc::new(Mutex::new(Vec::new())),
        }));
        let sink = g.sink(Box::new(TraceSink {
            trace: Arc::new(Mutex::new(Vec::new())),
        }));
        g.route(src, Route::Forward(sink));
        g.route(other, Route::Forward(sink));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        std::thread::sleep(Duration::from_millis(80));
        running.stop();
        let completed = Arc::new(Mutex::new(Vec::new()));
        let seen = Arc::clone(&completed);
        running
            .finish(&erx, &mut tracker, move |c| {
                seen.lock().unwrap().push(c);
                Ok(())
            })
            .unwrap();
        running.join();
        let completed = completed.lock().unwrap();
        assert!(!completed.is_empty(), "epochs completed after the source drained");
        for c in completed.iter() {
            let pos = c.snapshots.iter().find(|(t, _)| *t == src).map(|(_, s)| s.head.clone());
            assert_eq!(
                pos,
                Some(b"src@done".to_vec()),
                "epoch {}: the drained source's final position",
                c.epoch
            );
        }
    }

    #[test]
    fn epoch_tracker_completes_in_order_and_counts_finished_tasks_as_reported() {
        let mut t = EpochTracker::new(3); // 0 source, 1 operator, 2 sink
        let snap = |task, epoch| Event::Snapshot {
            task,
            epoch,
            snap: OpSnapshot::whole(vec![task as u8]),
        };
        assert!(t.on_event(&snap(0, 1)).is_empty());
        assert!(t.on_event(&snap(1, 1)).is_empty());
        // epoch 2 reported by the source before epoch 1's ack: nothing completes yet.
        assert!(t.on_event(&snap(0, 2)).is_empty());
        let done = t.on_event(&Event::Ack { task: 2, epoch: 1 });
        assert_eq!(done.len(), 1);
        assert_eq!(done[0].epoch, 1);
        assert_eq!(
            done[0].snapshots,
            vec![(0, OpSnapshot::whole(vec![0])), (1, OpSnapshot::whole(vec![1]))]
        );
        // the operator finished (Eos) before epoch 2's barrier reached it: it counts as reported.
        assert!(t.on_event(&Event::Finished { task: 1, cpu_ms: 0 }).is_empty());
        let done = t.on_event(&Event::Ack { task: 2, epoch: 2 });
        assert_eq!(done.iter().map(|c| c.epoch).collect::<Vec<_>>(), vec![2]);
        assert!(!t.all_finished());
        t.on_event(&Event::Finished { task: 0, cpu_ms: 0 });
        t.on_event(&Event::Finished { task: 2, cpu_ms: 0 });
        assert!(t.all_finished());
    }

    #[test]
    fn a_chained_sink_runs_in_the_source_task_and_flushes_before_the_position_is_taken() {
        // one task, no edge: data reaches the sink in order; at a barrier the sink's flush comes
        // before the source records its position; a refused flush leaves no position; Eos reaches
        // the sink when the task stops.
        struct TracingSource {
            trace: Arc<Mutex<Vec<String>>>,
            polls: usize,
        }
        impl Source for TracingSource {
            fn poll(&mut self, out: &mut Out) -> Poll {
                self.polls += 1;
                out.push(batch(&["c"], &[self.polls as i64]));
                Poll::Idle
            }
            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
                self.trace.lock().unwrap().push(format!("position {epoch}"));
                epoch.to_le_bytes().to_vec()
            }
            fn on_commit(&mut self, epoch: u64) {
                self.trace.lock().unwrap().push(format!("commit {epoch}"));
            }
        }
        struct TracingSink {
            trace: Arc<Mutex<Vec<String>>>,
            refuse: u64,
        }
        impl Sink for TracingSink {
            fn on_data(&mut self, b: RecordBatch) {
                self.trace.lock().unwrap().push(format!("data {}", b.num_rows()));
            }
            fn on_barrier(&mut self, epoch: u64) -> bool {
                self.trace.lock().unwrap().push(format!("flush {epoch}"));
                epoch != self.refuse
            }
            fn on_eos(&mut self) {
                self.trace.lock().unwrap().push("eos".into());
            }
        }
        let trace = Arc::new(Mutex::new(Vec::new()));
        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
        let task = g.source(Box::new(Chained {
            source: TracingSource {
                trace: Arc::clone(&trace),
                polls: 0,
            },
            sink: TracingSink {
                trace: Arc::clone(&trace),
                refuse: 2,
            },
        }));
        let (etx, erx) = sync_channel(4096);
        let running = g.start(etx).expect("start");
        let mut tracker = running.tracker();
        let mut completed = Vec::new();
        let until = Instant::now() + Duration::from_millis(80);
        while Instant::now() < until {
            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
                for c in tracker.on_event(&e) {
                    running.commit(c.epoch);
                    completed.push(c);
                }
            }
        }
        let last = running.stop();
        running
            .finish(&erx, &mut tracker, |c| {
                completed.push(c);
                Ok(())
            })
            .unwrap();
        running.join();
        assert_eq!(
            running.n_tasks(),
            1,
            "no sink task: the sink lives in the source's task"
        );
        let trace = trace.lock().unwrap().clone();
        let at = |s: &str| trace.iter().position(|t| t == s);
        assert!(
            at("data 1").unwrap() < at("flush 1").unwrap(),
            "data before the first barrier"
        );
        assert!(
            at("flush 1").unwrap() < at("position 1").unwrap(),
            "durable before the position: {trace:?}"
        );
        assert!(
            at("position 2").is_none(),
            "a refused flush records no position for that epoch"
        );
        assert!(at("flush 3").unwrap() < at("position 3").unwrap());
        // the last epoch commits after its flush and position, and the sink sees Eos when the
        // task stops. Which of the two lands last depends on whether the ticker or the stop
        // issued that epoch (its commit can reach the task before the stop does); both orders
        // keep offsets behind durable output.
        let final_commit = at(&format!("commit {last}")).expect("the final epoch commits");
        assert!(
            at(&format!("flush {last}")).unwrap() < final_commit,
            "the final commit follows its flush: {trace:?}"
        );
        let eos = at("eos").expect("the sink saw Eos when the task stopped");
        assert!(
            trace.len() - 1 == final_commit || trace.len() - 1 == eos,
            "the final commit or Eos lands last: {trace:?}"
        );
        assert!(
            completed.iter().any(|c| c.epoch == 2),
            "epoch 2 still completes (one task, one snapshot)"
        );
        let _ = task;
    }

    #[test]
    fn thread_cpu_ms_reads_this_threads_own_time() {
        // Linux only: the ledger every task reports on Finished. Burn a little CPU and see it.
        let Some(before) = thread_cpu_ms() else {
            return; // not Linux: the ledger is absent, never wrong
        };
        // burn until the ledger moved (a loaded CI runner may give this thread little CPU per
        // wall second; the ledger's resolution is 10 ms).
        let mut x: u64 = 0;
        let start = Instant::now();
        let mut after = before;
        while after < before + 20 && start.elapsed() < Duration::from_secs(5) {
            for _ in 0..100_000 {
                x = x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            }
            after = thread_cpu_ms().unwrap();
        }
        assert!(x != 1, "keep the loop");
        assert!(after >= before + 20, "cpu time advanced: {before} → {after}");
        // another thread's time is its own: a fresh thread starts near zero.
        let fresh = std::thread::spawn(|| thread_cpu_ms().unwrap()).join().unwrap();
        assert!(
            fresh < after,
            "a new thread's ledger is not this one's: {fresh} vs {after}"
        );
    }

    #[test]
    fn shuffle_splits_rows_by_vnode_range_and_keeps_all_rows() {
        let b = batch(&["a", "b", "c", "d", "e", "f"], &[1, 2, 3, 4, 5, 6]);
        let ranges = vnode_ranges(8, 3);
        let targets: Vec<(TaskId, std::ops::Range<u32>)> =
            ranges.into_iter().enumerate().map(|(i, r)| (i as TaskId, r)).collect();
        let parts = shuffle(&b, &["k".to_string()], 8, &targets, test_hasher().as_ref());
        let total: usize = parts.iter().flatten().map(|p| p.num_rows()).sum();
        assert_eq!(total, 6);
        // the same key always lands on the same target.
        let again = shuffle(&b, &["k".to_string()], 8, &targets, test_hasher().as_ref());
        for (p, q) in parts.iter().zip(again.iter()) {
            assert_eq!(p.as_ref().map(|x| x.num_rows()), q.as_ref().map(|x| x.num_rows()));
        }
        // no routing columns: everything to the first target.
        let global = shuffle(&b, &[], 8, &targets, test_hasher().as_ref());
        assert_eq!(global[0].as_ref().map(|x| x.num_rows()), Some(6));
        assert!(global[1].is_none() && global[2].is_none());
    }

    #[test]
    fn vnode_ranges_cover_every_vnode_exactly_once() {
        let r = vnode_ranges(64, 5);
        assert_eq!(r.len(), 5);
        let covered: Vec<u32> = r.iter().flat_map(|x| x.clone()).collect();
        assert_eq!(covered, (0..64).collect::<Vec<_>>());
        assert_eq!(vnode_ranges(4, 9).len(), 4, "never more tasks than vnodes");
    }
}