Skip to main content

fv_streams_runtime/
dataflow.rs

1//! THE DATAFLOW RUNTIME (Phase 2): a pipeline as tasks on OS threads connected by bounded
2//! in-memory edges that carry Arrow batches and the two control signals — watermarks and
3//! checkpoint barriers — in band (design/2026-09-15-phase2-dataflow-scoping.md).
4//!
5//! - A **task** owns one operator and runs a loop over its single inbox, into which every
6//!   upstream edge sends tagged messages; FIFO per edge is what per-key order rests on.
7//! - An **edge** is `sync_channel`-bounded: a full inbox backpressures the sender.
8//! - **Shuffle** hashes the routing columns of a batch (`create_hashes`) into vnodes and sends
9//!   one sub-batch (`take`) per downstream task by its vnode range; **forward** sends whole batches
10//!   to one task; **broadcast** sends the same batch to every downstream task.
11//! - **Watermarks**: each source emits `Watermark { source, ts }` (`None` = idle); a task keeps the
12//!   latest per (edge, source), fires its operator on the minimum over the active ones, and emits
13//!   its own watermark downstream under its own id, so the whole graph shares one time.
14//! - **Barriers and epochs**: the coordinator issues epochs — periodically from the runtime's own
15//!   ticker, and one last time on `stop` — as a `Control::Barrier` to every source; a source
16//!   records its position, reports a `Snapshot`, and forwards `Barrier { epoch }` behind the data
17//!   it already sent. A task with several input edges **aligns** — batches arriving on an edge
18//!   that already delivered the barrier are held until every edge has — then snapshots its
19//!   operator, forwards the barrier, and hands the snapshot to the coordinator. Sinks make their
20//!   output durable and `Ack`. When every task has reported an epoch (or finished), the epoch is
21//!   complete: the coordinator sends `Control::Commit` to the sources, which commit the positions
22//!   they recorded at that barrier. Offsets therefore never run ahead of durable output — the
23//!   at-least-once gate, now across task boundaries.
24//! - **Eos** flows in order behind the last data and the final barrier; a task finishes when every
25//!   input reached it.
26//!
27//! The runtime is Kafka-free: sources and sinks are traits, tested here with in-memory ones.
28
29use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
30use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
31use std::sync::mpsc::{channel, sync_channel, Receiver, RecvTimeoutError, Sender, SyncSender};
32use std::sync::{Arc, Mutex};
33use std::time::{Duration, Instant};
34
35use tokio::runtime::Handle;
36
37use fv_streams_exchange::{LinkReceiver, LinkSender};
38
39use crate::exchange_bridge::{InboundBridge, OutboundBridge};
40use crate::placement::WorkerId;
41
42use arrow::array::{RecordBatch, UInt32Array};
43use arrow::compute::take;
44
45pub use fv_streams_types::contract::{
46    position_record, Acker, Chained, DeferredOperator, DeferredSink, EdgeId, EpochStaging, OpSnapshot, Operator, Out,
47    Poll, Sink, Source, StateFile, TaskId,
48};
49
50/// A resolved out-route: the sender into each target's inbox, tagged with the edge id and the
51/// target task (named when a send blocks).
52type Outlet = (EdgeId, TaskId, SyncSender<(EdgeId, Msg)>);
53
54/// What a task is doing right now — the runtime's own answer to "where is it stuck": every task
55/// keeps one, the coordinator reads them (`Running::task_states`).
56#[derive(Clone, Copy, Debug, PartialEq, Eq)]
57#[repr(u8)]
58pub enum TaskState {
59    Starting = 0,
60    /// A source inside `poll`.
61    Polling = 1,
62    /// A consumer waiting on its inbox.
63    Receiving = 2,
64    /// Inside the operator's or sink's `on_data`.
65    Working = 3,
66    /// Blocked sending into `target`'s inbox.
67    Sending = 4,
68    /// Taking or completing a barrier.
69    Barrier = 5,
70    /// A source past Eos, waiting for its last epoch's commit.
71    CommitWait = 6,
72    Done = 7,
73}
74
75impl TaskState {
76    fn from_u8(v: u8) -> TaskState {
77        match v {
78            1 => TaskState::Polling,
79            2 => TaskState::Receiving,
80            3 => TaskState::Working,
81            4 => TaskState::Sending,
82            5 => TaskState::Barrier,
83            6 => TaskState::CommitWait,
84            7 => TaskState::Done,
85            _ => TaskState::Starting,
86        }
87    }
88}
89
90/// One task's probe: its state, the target of a send, and when it entered the state.
91pub struct Probe {
92    state: std::sync::atomic::AtomicU8,
93    target: std::sync::atomic::AtomicU32,
94    since_ms: std::sync::atomic::AtomicI64,
95}
96
97impl Probe {
98    fn new() -> Self {
99        Probe {
100            state: std::sync::atomic::AtomicU8::new(0),
101            target: std::sync::atomic::AtomicU32::new(u32::MAX),
102            since_ms: std::sync::atomic::AtomicI64::new(now_ms()),
103        }
104    }
105
106    fn set(&self, state: TaskState, target: TaskId) {
107        self.target.store(target, Ordering::Relaxed);
108        self.state.store(state as u8, Ordering::Relaxed);
109        self.since_ms.store(now_ms(), Ordering::Relaxed);
110    }
111
112    /// `(state, send target if sending, milliseconds in this state)`.
113    pub fn read(&self) -> (TaskState, Option<TaskId>, i64) {
114        let state = TaskState::from_u8(self.state.load(Ordering::Relaxed));
115        let target = self.target.load(Ordering::Relaxed);
116        let since = self.since_ms.load(Ordering::Relaxed);
117        (
118            state,
119            (state == TaskState::Sending && target != u32::MAX).then_some(target),
120            (now_ms() - since).max(0),
121        )
122    }
123}
124type OutRoute = (Route, Vec<Outlet>);
125
126/// What flows on an edge.
127#[derive(Clone, Debug)]
128pub enum Msg {
129    Data(RecordBatch),
130    /// The time of `source` (a source task id) as seen by the sender; `None` = the source is idle.
131    Watermark {
132        source: TaskId,
133        ts: Option<i64>,
134    },
135    Barrier {
136        epoch: u64,
137    },
138    /// Processing time passes with nothing else to say: a source that polled nothing sends one
139    /// (at most ten a second), and every task forwards at most ten a second. Operators with
140    /// processing-time work (bounded idleness, ingest-time windows) advance on it.
141    Tick {
142        now_ms: i64,
143    },
144    /// The upstream is done (clean shutdown flows downstream in order).
145    Eos,
146}
147
148/// Test-only structural equality (a `RecordBatch` compares by schema + columns), so seam tests can
149/// `assert_eq!` a round-tripped message against what was sent.
150#[cfg(test)]
151impl PartialEq for Msg {
152    fn eq(&self, other: &Msg) -> bool {
153        match (self, other) {
154            (Msg::Data(a), Msg::Data(b)) => a == b,
155            (Msg::Watermark { source: s1, ts: t1 }, Msg::Watermark { source: s2, ts: t2 }) => (s1, t1) == (s2, t2),
156            (Msg::Barrier { epoch: a }, Msg::Barrier { epoch: b }) => a == b,
157            (Msg::Tick { now_ms: a }, Msg::Tick { now_ms: b }) => a == b,
158            (Msg::Eos, Msg::Eos) => true,
159            _ => false,
160        }
161    }
162}
163
164/// Hashes the routing columns of a batch into `out` (one `u64` per row) — the shuffle's key hash,
165/// supplied by the planner (the engine's DataFusion hash, fixed-seed: a key's vnode is the same on
166/// every run, every worker and every restart). The runtime never chooses a hash function.
167pub type Hasher = Arc<dyn Fn(&[arrow::array::ArrayRef], &mut [u64]) + Send + Sync>;
168
169/// Where a task's output goes.
170#[derive(Clone)]
171pub enum Route {
172    /// Whole batches to one task.
173    Forward(TaskId),
174    /// Rows hashed on `columns` into vnodes; each target owns a contiguous vnode range.
175    Shuffle {
176        columns: Vec<String>,
177        vnodes: u32,
178        targets: Vec<(TaskId, std::ops::Range<u32>)>,
179        hasher: Hasher,
180    },
181    /// The same batch to every target.
182    Broadcast(Vec<TaskId>),
183}
184
185impl std::fmt::Debug for Route {
186    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
187        match self {
188            Route::Forward(t) => write!(f, "Forward({t})"),
189            Route::Shuffle {
190                columns,
191                vnodes,
192                targets,
193                ..
194            } => f
195                .debug_struct("Shuffle")
196                .field("columns", columns)
197                .field("vnodes", vnodes)
198                .field("targets", targets)
199                .finish(),
200            Route::Broadcast(ts) => write!(f, "Broadcast({ts:?})"),
201        }
202    }
203}
204
205impl Route {
206    pub fn targets(&self) -> Vec<TaskId> {
207        match self {
208            Route::Forward(t) => vec![*t],
209            Route::Shuffle { targets, .. } => targets.iter().map(|(t, _)| *t).collect(),
210            Route::Broadcast(ts) => ts.clone(),
211        }
212    }
213}
214
215/// What the coordinator tells a source.
216#[derive(Clone, Copy, Debug)]
217pub enum Control {
218    /// Take a barrier for `epoch`: record the position, forward the barrier.
219    Barrier(u64),
220    /// Every task reported `epoch`: commit the position recorded at its barrier.
221    Commit(u64),
222    /// Finish after the current poll: Eos flows downstream in order.
223    Stop,
224}
225
226/// What the coordinator hears from tasks.
227#[derive(Debug)]
228pub enum Event {
229    /// A source or operator task snapshotted for `epoch`.
230    Snapshot { task: TaskId, epoch: u64, snap: OpSnapshot },
231    /// A sink made everything before `epoch`'s barrier durable.
232    Ack { task: TaskId, epoch: u64 },
233    /// A source ran dry: its final position (see [`Source::position`]) — carried into every epoch
234    /// completed after it, since a finished source takes no more barriers.
235    Drained { task: TaskId, snap: OpSnapshot },
236    /// A task finished (its inputs hit Eos, the source ran dry, or the coordinator stopped it),
237    /// having used `cpu_ms` of CPU on its thread — where the build's cores went, by task.
238    Finished { task: TaskId, cpu_ms: u64 },
239    /// A source hit its strict data contract; it finishes right after.
240    Failed { task: TaskId, error: String },
241}
242
243enum Node {
244    Source(Box<dyn Source>),
245    Operator(Box<dyn Operator>),
246    /// An operator constructed at `start`, once its in-edges are known.
247    Deferred(DeferredOperator),
248    DeferredSink(DeferredSink),
249    Sink(Box<dyn Sink>),
250}
251
252/// A sink that takes nothing: the placeholder while a node is swapped.
253struct NoSink;
254impl Sink for NoSink {
255    fn on_data(&mut self, _batch: RecordBatch) {}
256    fn on_barrier(&mut self, _epoch: u64) -> bool {
257        true
258    }
259    fn on_eos(&mut self) {}
260}
261
262struct TaskDef {
263    id: TaskId,
264    node: Node,
265    routes: Vec<Route>,
266}
267
268/// This worker's exchange endpoints for a multi-process run: one [`LinkSender`] per peer worker it
269/// sends to, and one [`LinkReceiver`] per peer it receives from (established by the membership
270/// handshake before `start`). [`Graph::start`] runs with none of this — a single-worker run where
271/// every edge is a local `SyncSender`; the multi-process path is [`Graph::start_worker`], which
272/// derives the actual edge wiring from the placement.
273///
274/// The mechanism is identical either way — an out-edge is a bounded `SyncSender`, drained by the
275/// dest task (local) or by an `OutboundBridge` (remote); an `InboundBridge` delivers a decoded
276/// `(edge, msg)` into the local dest inbox exactly as a local outlet would. Edge ids agree across
277/// workers because every worker builds the same graph in the same order, so the `edge` on the wire
278/// names the same logical edge — and thus the same dest task — on both ends.
279pub struct WorkerExchange {
280    /// The tokio runtime the bridges drive (`block_on`). Must be multi-threaded: its workers run
281    /// the links' socket tasks while the bridge threads block on them.
282    pub handle: Handle,
283    /// A link to each peer worker this worker sends edges to, keyed by that worker's id.
284    pub senders: HashMap<WorkerId, LinkSender>,
285    /// A link from each peer this worker receives edges from; each delivers its frames as
286    /// `(edge, msg)` into local inboxes, routed by the edge's target task.
287    pub receivers: Vec<LinkReceiver>,
288}
289
290/// The graph under construction: add tasks, connect routes, then `start`.
291#[derive(Default)]
292pub struct Graph {
293    tasks: Vec<TaskDef>,
294    capacity: usize,
295    barrier_every: Option<Duration>,
296    /// The epoch this graph continues from (a restored run continues its checkpoint's numbering,
297    /// so the newest manifest is always the newest epoch).
298    first_epoch: u64,
299}
300
301impl Graph {
302    pub fn new(edge_capacity: usize) -> Self {
303        Graph {
304            tasks: Vec::new(),
305            capacity: edge_capacity.max(1),
306            barrier_every: None,
307            first_epoch: 0,
308        }
309    }
310
311    /// Continue the epoch numbering from `epoch`: the first epoch issued is `epoch + 1`. A run
312    /// restored from a checkpoint passes that checkpoint's epoch, so its own checkpoints always
313    /// read as newer than the one it restored from.
314    pub fn start_at_epoch(mut self, epoch: u64) -> Self {
315        self.first_epoch = epoch;
316        self
317    }
318
319    /// Issue an epoch (a barrier to every source) this often. Without it, only `stop` issues one.
320    pub fn barrier_every(mut self, every: Duration) -> Self {
321        self.barrier_every = Some(every).filter(|d| !d.is_zero());
322        self
323    }
324
325    pub fn source(&mut self, s: Box<dyn Source>) -> TaskId {
326        self.add(Node::Source(s))
327    }
328
329    pub fn operator(&mut self, o: Box<dyn Operator>) -> TaskId {
330        self.add(Node::Operator(o))
331    }
332
333    /// An operator that needs to know its input edges (a join's sides): built at `start`.
334    pub fn operator_with(&mut self, f: DeferredOperator) -> TaskId {
335        self.add(Node::Deferred(f))
336    }
337
338    pub fn sink(&mut self, s: Box<dyn Sink>) -> TaskId {
339        self.add(Node::Sink(s))
340    }
341
342    /// A sink built at start on its owning worker only (see [`DeferredSink`]).
343    pub fn sink_with(&mut self, f: DeferredSink) -> TaskId {
344        self.add(Node::DeferredSink(f))
345    }
346
347    fn add(&mut self, node: Node) -> TaskId {
348        let id = self.tasks.len() as TaskId;
349        self.tasks.push(TaskDef {
350            id,
351            node,
352            routes: Vec::new(),
353        });
354        id
355    }
356
357    /// Route `from`'s output.
358    pub fn route(&mut self, from: TaskId, route: Route) {
359        self.tasks[from as usize].routes.push(route);
360    }
361
362    /// Chain `sink` into source task `task` (see [`Chained`]); hands the sink back when the task
363    /// is not a source.
364    pub fn chain_sink(&mut self, task: TaskId, sink: Box<dyn Sink>) -> Result<(), Box<dyn Sink>> {
365        let Some(def) = self.tasks.get_mut(task as usize) else {
366            return Err(sink);
367        };
368        if !matches!(def.node, Node::Source(_)) {
369            return Err(sink);
370        }
371        let placeholder = Node::Sink(Box::new(NoSink));
372        let Node::Source(source) = std::mem::replace(&mut def.node, placeholder) else {
373            unreachable!()
374        };
375        def.node = Node::Source(Box::new(Chained { source, sink }));
376        Ok(())
377    }
378
379    /// Spawn every task on this single-process worker; returns the running dataflow. Every edge is
380    /// a local `SyncSender`. For one worker of a multi-process run, see [`Graph::start_worker`].
381    pub fn start(self, events: SyncSender<Event>) -> Result<Running, String> {
382        // A solo placement: every task on the one worker, no exchange — byte-identical to the
383        // pre-Phase-4 runtime.
384        let solo = vec![0 as WorkerId; self.tasks.len()];
385        self.start_worker(events, &solo, 0, None)
386    }
387
388    /// Spawn only the tasks this worker owns and wire the graph's edges according to the placement
389    /// `of_task` (task id → worker id): an edge between two co-located tasks is a local `SyncSender`;
390    /// an edge from a local task to a task on another worker is drained by an `OutboundBridge` onto
391    /// the link to that worker; an edge into a local task from a remote one arrives via an
392    /// `InboundBridge` on the peer's link and is delivered to the local dest inbox. A task on no
393    /// other worker is not run here. **Every worker builds the identical graph in the identical
394    /// order**, so edge ids agree across workers and the `edge` on the wire names the same logical
395    /// edge — hence the same dest task — on every worker. A solo placement (every task on `me`, no
396    /// `exchange`) is exactly [`Graph::start`].
397    pub fn start_worker(
398        self,
399        events: SyncSender<Event>,
400        of_task: &[WorkerId],
401        me: WorkerId,
402        exchange: Option<WorkerExchange>,
403    ) -> Result<Running, String> {
404        let n = self.tasks.len();
405        assert_eq!(of_task.len(), n, "placement must cover every task");
406        let is_local = |id: TaskId| of_task[id as usize] == me;
407        // An inbox only for a task this worker runs; a remote task has none here.
408        let mut inbox_tx: Vec<Option<SyncSender<(EdgeId, Msg)>>> = Vec::with_capacity(n);
409        let mut inbox_rx: Vec<Option<Receiver<(EdgeId, Msg)>>> = Vec::with_capacity(n);
410        for id in 0..n as TaskId {
411            if is_local(id) {
412                let (tx, rx) = sync_channel(self.capacity);
413                inbox_tx.push(Some(tx));
414                inbox_rx.push(Some(rx));
415            } else {
416                inbox_tx.push(None);
417                inbox_rx.push(None);
418            }
419        }
420        // The exchange endpoints, unpacked so the solo case allocates nothing of note.
421        let (handle, peer_senders, receivers) = match exchange {
422            Some(x) => (Some(x.handle), x.senders, x.receivers),
423            None => (None, HashMap::new(), Vec::new()),
424        };
425        let mut next_edge: EdgeId = 0;
426        let mut in_edges: Vec<Vec<(EdgeId, TaskId)>> = vec![Vec::new(); n];
427        // edge id -> its target task, so an inbound bridge routes a decoded `(edge, msg)` to the
428        // right local inbox — the mirror of a local outlet's `SyncSender`.
429        let mut edge_target: Vec<TaskId> = Vec::new();
430        // per task: its outgoing routes resolved to (edge id, target inbox sender)
431        let mut outs: Vec<Vec<OutRoute>> = vec![Vec::new(); n];
432        let mut out_bridges: Vec<OutboundBridge> = Vec::new();
433        for t in &self.tasks {
434            // Edge ids are assigned for every route on every worker (identical order), so they agree
435            // across workers; only a local upstream's outlets are built here (a remote task's routes
436            // are wired on its own worker).
437            let upstream_local = is_local(t.id);
438            for r in &t.routes {
439                let mut senders = Vec::new();
440                for target in r.targets() {
441                    let e = next_edge;
442                    next_edge += 1;
443                    edge_target.push(target); // index == e
444                    in_edges[target as usize].push((e, t.id));
445                    if !upstream_local {
446                        continue;
447                    }
448                    // Local target: its inbox. Remote target: a bounded channel (its own
449                    // backpressure) drained by an outbound bridge onto the link to that worker.
450                    let tx = if is_local(target) {
451                        inbox_tx[target as usize].clone().expect("local inbox")
452                    } else {
453                        let h = handle.as_ref().expect("a remote edge needs an exchange handle");
454                        let link = peer_senders
455                            .get(&of_task[target as usize])
456                            .expect("a link to the target's worker")
457                            .clone();
458                        let (btx, brx) = sync_channel(self.capacity);
459                        out_bridges.push(OutboundBridge::spawn(h.clone(), link, brx));
460                        btx
461                    };
462                    senders.push((e, target, tx));
463                }
464                if upstream_local {
465                    outs[t.id as usize].push((r.clone(), senders));
466                }
467            }
468        }
469        // Inbound: one bridge per peer link, delivering `(edge, msg)` into the local dest inbox.
470        // The closures hold clones of the local inboxes, so a task with only remote inputs keeps its
471        // inbox open until its peer's Eos arrives (then the link closes and the bridge ends).
472        let mut in_bridges: Vec<InboundBridge> = Vec::new();
473        if let Some(h) = handle.as_ref() {
474            for link_rx in receivers {
475                let inboxes = inbox_tx.clone();
476                let targets = edge_target.clone();
477                in_bridges.push(InboundBridge::spawn(h.clone(), link_rx, move |edge, msg| {
478                    if let Some(inbox) = targets.get(edge as usize).and_then(|&t| inboxes[t as usize].as_ref()) {
479                        let _ = inbox.send((edge, msg));
480                    }
481                }));
482            }
483        }
484        drop(inbox_tx);
485        let local_tasks: Vec<TaskId> = (0..n as TaskId).filter(|&id| is_local(id)).collect();
486        let inner = Arc::new(Inner {
487            ctl: Mutex::new(Vec::new()),
488            epoch: AtomicU64::new(self.first_epoch),
489            stopping: AtomicBool::new(false),
490            n_tasks: n,
491        });
492        let probes: Arc<Vec<Probe>> = Arc::new((0..n).map(|_| Probe::new()).collect());
493        // Build every local task's node FIRST. A deferred operator's constructor can refuse (a
494        // checkpoint it cannot restore), and that must fail the start before any task thread
495        // exists — never orphan tasks already spawned with no coordinator to stop them. Same order
496        // as the spawn loop below, so task and edge ids agree across workers exactly as before.
497        let mut built = Vec::with_capacity(local_tasks.len());
498        for t in self.tasks.into_iter() {
499            let i = t.id as usize;
500            if !is_local(t.id) {
501                continue; // a task on another worker: not run in this process.
502            }
503            let rx = inbox_rx[i].take().expect("inbox rx");
504            let out_routes = std::mem::take(&mut outs[i]);
505            let ins_with_upstream = std::mem::take(&mut in_edges[i]);
506            let ins: Vec<EdgeId> = ins_with_upstream.iter().map(|(e, _)| *e).collect();
507            let node = match t.node {
508                Node::Deferred(build) => Node::Operator(build(t.id, &ins_with_upstream)?),
509                Node::DeferredSink(build) => Node::Sink(build(t.id)?),
510                other => other,
511            };
512            built.push((
513                TaskDef {
514                    id: t.id,
515                    node,
516                    routes: t.routes,
517                },
518                rx,
519                out_routes,
520                ins,
521            ));
522        }
523        let mut handles = Vec::with_capacity(built.len());
524        for (t, rx, out_routes, ins) in built {
525            let events = events.clone();
526            let task_probes = Arc::clone(&probes);
527            let id = t.id;
528            // sources and sinks hear the coordinator over their own control channel (sources
529            // take barriers and commit positions; sinks commit transactions).
530            let ctl = if matches!(t.node, Node::Source(_) | Node::Sink(_)) {
531                let (tx, rx) = channel();
532                inner.ctl.lock().expect("control senders").push(tx);
533                Some(rx)
534            } else {
535                None
536            };
537            handles.push(
538                std::thread::Builder::new()
539                    .name(format!("fv-task-{id}"))
540                    .spawn(move || {
541                        let mut task = Task {
542                            id,
543                            node: t.node,
544                            rx,
545                            ctl,
546                            out_routes,
547                            in_edges: ins,
548                            last_tick: None,
549                            probes: task_probes,
550                            events,
551                            wm: HashMap::new(),
552                            last_wm: None,
553                            align: Align::default(),
554                            eos_seen: HashSet::new(),
555                        };
556                        task.run();
557                    })
558                    .expect("spawn task"),
559            );
560        }
561        // the periodic epochs: the runtime's own ticker, until `stop`.
562        let ticker = self.barrier_every.map(|every| {
563            let inner = Arc::clone(&inner);
564            std::thread::Builder::new()
565                .name("fv-epochs".into())
566                .spawn(move || {
567                    let nap = Duration::from_millis(10).min(every);
568                    let mut next = Instant::now() + every;
569                    while !inner.stopping.load(Ordering::Relaxed) {
570                        if Instant::now() >= next {
571                            inner.barrier();
572                            next = Instant::now() + every;
573                        }
574                        std::thread::sleep(nap);
575                    }
576                })
577                .expect("spawn epochs")
578        });
579        Ok(Running {
580            inner,
581            handles: Mutex::new(handles),
582            probes,
583            ticker: Mutex::new(ticker),
584            out_bridges: Mutex::new(out_bridges),
585            in_bridges: Mutex::new(in_bridges),
586            local_tasks,
587        })
588    }
589}
590
591/// The coordinator's side of the sources: their control channels and the epoch counter.
592struct Inner {
593    ctl: Mutex<Vec<Sender<Control>>>,
594    epoch: AtomicU64,
595    stopping: AtomicBool,
596    n_tasks: usize,
597}
598
599impl Inner {
600    fn send_all(&self, c: Control) {
601        for tx in self.ctl.lock().expect("control senders").iter() {
602            let _ = tx.send(c); // a finished source dropped its receiver: nothing to tell it
603        }
604    }
605
606    fn barrier(&self) -> u64 {
607        let epoch = self.epoch.fetch_add(1, Ordering::SeqCst) + 1;
608        self.send_all(Control::Barrier(epoch));
609        epoch
610    }
611}
612
613/// A started dataflow: issue epochs, commit them, stop it, join it.
614pub struct Running {
615    inner: Arc<Inner>,
616    handles: Mutex<Vec<std::thread::JoinHandle<()>>>,
617    ticker: Mutex<Option<std::thread::JoinHandle<()>>>,
618    /// Every task's state probe, by task id.
619    probes: Arc<Vec<Probe>>,
620    /// The exchange bridges (empty for a single-worker run). Joined after the tasks: an outbound
621    /// bridge ends when its task drops the edge; an inbound bridge when its peer link closes.
622    out_bridges: Mutex<Vec<OutboundBridge>>,
623    in_bridges: Mutex<Vec<InboundBridge>>,
624    /// The tasks this worker actually runs — `0..n` for a solo run, this worker's own tasks for one
625    /// worker of a multi-process run. The tracker waits on exactly these.
626    local_tasks: Vec<TaskId>,
627}
628
629impl Running {
630    /// Every task's current state: `(task, state, send target, milliseconds in that state)`.
631    pub fn task_states(&self) -> Vec<(TaskId, TaskState, Option<TaskId>, i64)> {
632        self.probes
633            .iter()
634            .enumerate()
635            .map(|(i, p)| {
636                let (s, t, ms) = p.read();
637                (i as TaskId, s, t, ms)
638            })
639            .collect()
640    }
641
642    pub fn n_tasks(&self) -> usize {
643        self.inner.n_tasks
644    }
645
646    /// A tracker for the tasks this worker runs (every task for a solo run).
647    pub fn tracker(&self) -> EpochTracker {
648        EpochTracker::new_for(self.local_tasks.clone())
649    }
650
651    /// Issue the next epoch now (a barrier to every source); returns it.
652    pub fn barrier(&self) -> u64 {
653        self.inner.barrier()
654    }
655
656    /// Issue a barrier for a *specific* epoch to this worker's sources, setting the local epoch to
657    /// it. A joined worker calls this when the leader injects epoch `e`, so every worker's sources
658    /// stamp the same epoch — the barrier then aligns across a cross-worker shuffle. (The leader
659    /// uses [`Running::barrier`] to pick `e`, then tells the workers to inject it.)
660    pub fn inject_epoch(&self, epoch: u64) {
661        self.inner.epoch.store(epoch, Ordering::SeqCst);
662        self.inner.send_all(Control::Barrier(epoch));
663    }
664
665    /// The last epoch issued.
666    pub fn epoch(&self) -> u64 {
667        self.inner.epoch.load(Ordering::SeqCst)
668    }
669
670    /// Every task reported `epoch`: the sources may commit the positions they recorded at it.
671    pub fn commit(&self, epoch: u64) {
672        self.inner.send_all(Control::Commit(epoch));
673    }
674
675    /// Tell the sources to stop (Eos flows behind the last data) — WITHOUT issuing a barrier, unlike
676    /// [`Running::stop`]. A clustered run uses this at shutdown: the final epoch was already issued
677    /// and committed by the leader-coordinated round, so a fresh uncommitted barrier here would leave
678    /// the sources waiting for a commit that never comes.
679    pub fn signal_stop(&self) {
680        self.inner.stopping.store(true, Ordering::SeqCst);
681        self.inner.send_all(Control::Stop);
682    }
683
684    /// Stop: no more periodic epochs, one final barrier, then `Stop` to every source (Eos flows
685    /// downstream in order behind the barrier). Returns the final epoch. Idempotent.
686    pub fn stop(&self) -> u64 {
687        if self.inner.stopping.swap(true, Ordering::SeqCst) {
688            return self.epoch();
689        }
690        if let Some(t) = self.ticker.lock().expect("ticker").take() {
691            let _ = t.join();
692        }
693        let epoch = self.inner.barrier();
694        self.inner.send_all(Control::Stop);
695        epoch
696    }
697
698    /// Drive `events` until every task finished: a completed epoch is committed to the sources
699    /// and handed to `on_epoch` with its snapshots; a task's failure stops the graph and is the
700    /// result. Call after `stop` (or for a graph whose sources are bounded).
701    pub fn finish(
702        &self,
703        events: &Receiver<Event>,
704        tracker: &mut EpochTracker,
705        mut on_epoch: impl FnMut(Completed) -> Result<(), String>,
706    ) -> Result<(), String> {
707        let mut failure: Option<String> = None;
708        while !tracker.all_finished() {
709            match events.recv_timeout(Duration::from_millis(50)) {
710                Ok(e) => {
711                    if let Event::Failed { error, .. } = &e {
712                        failure.get_or_insert_with(|| error.clone());
713                        self.stop();
714                        // a sink waiting for its last epoch's commit must not wait forever: the
715                        // commit is never coming (uncommitted = replay), so close the control
716                        // channels and let it finish.
717                        self.inner.ctl.lock().expect("control senders").clear();
718                    }
719                    // the coordinator makes the epoch durable in `on_epoch` (its checkpoint)
720                    // BEFORE anyone commits on it; an error there ends the run.
721                    for c in tracker.on_event(&e) {
722                        let epoch = c.epoch;
723                        match on_epoch(c) {
724                            Ok(()) => self.commit(epoch),
725                            Err(err) => {
726                                // a checkpoint that failed is the run's failure, said now: the
727                                // tasks waiting for its commit are released (uncommitted = replay).
728                                eprintln!("dataflow: epoch {} not committed — {err}", epoch);
729                                failure.get_or_insert(err);
730                                self.stop();
731                                self.inner.ctl.lock().expect("control senders").clear();
732                            }
733                        }
734                    }
735                }
736                Err(RecvTimeoutError::Timeout) => {}
737                Err(RecvTimeoutError::Disconnected) => break,
738            }
739        }
740        match failure {
741            Some(f) => Err(f),
742            None => Ok(()),
743        }
744    }
745
746    /// Join every task. Closes the control channels first, so a source waiting for a final commit
747    /// that will never come (the coordinator gave up) still exits — uncommitted means replay.
748    pub fn join(&self) {
749        self.inner.stopping.store(true, Ordering::SeqCst);
750        if let Some(t) = self.ticker.lock().expect("ticker").take() {
751            let _ = t.join();
752        }
753        self.inner.ctl.lock().expect("control senders").clear();
754        let handles: Vec<_> = self.handles.lock().expect("handles").drain(..).collect();
755        for h in handles {
756            let _ = h.join();
757        }
758        // Tasks are done: an outbound bridge's edge is now dropped (so it drains and ends, closing
759        // its link). Join those first.
760        for b in self.out_bridges.lock().expect("out bridges").drain(..) {
761            b.join();
762        }
763        // Every local task has finished, so all data and Eos are delivered: signal the inbound
764        // bridges to stop rather than wait for the peer to close its link. Waiting would deadlock —
765        // two workers would each hold their sockets open until the other closed first.
766        let in_bridges: Vec<InboundBridge> = self.in_bridges.lock().expect("in bridges").drain(..).collect();
767        for b in &in_bridges {
768            b.signal_stop();
769        }
770        for b in in_bridges {
771            b.join();
772        }
773    }
774}
775
776/// An epoch every task reported (or had finished before): its snapshots, by task.
777#[derive(Debug, Clone)]
778pub struct Completed {
779    pub epoch: u64,
780    /// `(task, snapshot)` from every source and operator that snapshotted, in task order.
781    pub snapshots: Vec<(TaskId, OpSnapshot)>,
782}
783
784/// The coordinator's book: which tasks reported which epoch, which tasks finished. Epochs
785/// complete in order. An epoch completes when every *expected* task has reported it (or finished
786/// before it) — for a single-process run that is every task `0..n`; for one worker of a multi-
787/// process run it is only the tasks placed on this worker (the leader aggregates across workers).
788#[derive(Debug)]
789pub struct EpochTracker {
790    /// The tasks this coordinator waits on — `0..n` for a solo run, a worker's own tasks otherwise.
791    expected: Vec<TaskId>,
792    finished: HashSet<TaskId>,
793    /// A drained source's final position, added to every epoch it no longer reports.
794    drained: HashMap<TaskId, OpSnapshot>,
795    pending: BTreeMap<u64, HashMap<TaskId, Option<OpSnapshot>>>,
796    cpu_ms: HashMap<TaskId, u64>,
797}
798
799impl EpochTracker {
800    /// A tracker for a contiguous task set `0..n_tasks` (the single-process run).
801    pub fn new(n_tasks: usize) -> Self {
802        EpochTracker::new_for((0..n_tasks as TaskId).collect())
803    }
804
805    /// A tracker for an explicit set of task ids (one worker's own tasks).
806    pub fn new_for(expected: Vec<TaskId>) -> Self {
807        EpochTracker {
808            expected,
809            finished: HashSet::new(),
810            drained: HashMap::new(),
811            pending: BTreeMap::new(),
812            cpu_ms: HashMap::new(),
813        }
814    }
815
816    /// Book one event; returns the epochs it completed, in order.
817    pub fn on_event(&mut self, e: &Event) -> Vec<Completed> {
818        match e {
819            Event::Snapshot { task, epoch, snap } => {
820                self.pending
821                    .entry(*epoch)
822                    .or_default()
823                    .insert(*task, Some(snap.clone()));
824            }
825            Event::Ack { task, epoch } => {
826                self.pending.entry(*epoch).or_default().insert(*task, None);
827            }
828            Event::Finished { task, cpu_ms } => {
829                self.finished.insert(*task);
830                self.cpu_ms.insert(*task, *cpu_ms);
831            }
832            Event::Drained { task, snap } => {
833                self.drained.insert(*task, snap.clone());
834            }
835            Event::Failed { .. } => {}
836        }
837        let mut done = Vec::new();
838        while let Some((&epoch, reported)) = self.pending.iter().next() {
839            let complete = self
840                .expected
841                .iter()
842                .all(|t| reported.contains_key(t) || self.finished.contains(t));
843            if !complete {
844                break;
845            }
846            let reported = self.pending.remove(&epoch).expect("present");
847            // a drained source reported nothing for this epoch: its final position stands in
848            let drained: Vec<(TaskId, OpSnapshot)> = self
849                .drained
850                .iter()
851                .filter(|(t, _)| !reported.contains_key(t))
852                .map(|(t, snap)| (*t, snap.clone()))
853                .collect();
854            let mut snapshots: Vec<(TaskId, OpSnapshot)> = reported
855                .into_iter()
856                .filter_map(|(t, b)| b.map(|b| (t, b)))
857                .chain(drained)
858                .collect();
859            snapshots.sort_by_key(|(t, _)| *t);
860            done.push(Completed { epoch, snapshots });
861        }
862        done
863    }
864
865    pub fn all_finished(&self) -> bool {
866        self.expected.iter().all(|t| self.finished.contains(t))
867    }
868
869    pub fn finished(&self) -> usize {
870        self.finished.len()
871    }
872
873    /// The CPU each finished task used on its thread, in milliseconds.
874    pub fn cpu_ms(&self) -> &HashMap<TaskId, u64> {
875        &self.cpu_ms
876    }
877}
878
879/// The CPU time of every thread in this process so far, summed by thread name (milliseconds):
880/// the engine's tasks by their `fv-task-N` names, librdkafka's `rdk:*` threads, the rest. Empty
881/// where `/proc` is absent. This is where a build's cores go, at a glance.
882pub fn process_cpu_by_thread() -> BTreeMap<String, u64> {
883    let mut out = BTreeMap::new();
884    let Ok(tasks) = std::fs::read_dir("/proc/self/task") else {
885        return out;
886    };
887    for t in tasks.flatten() {
888        let dir = t.path();
889        let Ok(comm) = std::fs::read_to_string(dir.join("comm")) else {
890            continue;
891        };
892        let Ok(stat) = std::fs::read_to_string(dir.join("stat")) else {
893            continue;
894        };
895        if let Some(ms) = cpu_ms_of_stat(&stat) {
896            *out.entry(comm.trim().to_string()).or_default() += ms;
897        }
898    }
899    out
900}
901
902/// utime + stime of a `/proc/<pid>/task/<tid>/stat` line, in milliseconds.
903fn cpu_ms_of_stat(stat: &str) -> Option<u64> {
904    // the comm field may contain spaces and parentheses: split after the last ')'.
905    let rest = &stat[stat.rfind(')')? + 1..];
906    let fields: Vec<&str> = rest.split_whitespace().collect();
907    // fields after the comm: state(0) ppid(1) … utime(11) stime(12)
908    let utime: u64 = fields.get(11)?.parse().ok()?;
909    let stime: u64 = fields.get(12)?.parse().ok()?;
910    Some((utime + stime) * 10)
911}
912
913/// The CPU time this thread has used so far (user + system), from `/proc/thread-self/stat`;
914/// `None` where that file does not exist. Linux reports it in USER_HZ ticks, which is 100 on
915/// every architecture regardless of the kernel's HZ.
916pub fn thread_cpu_ms() -> Option<u64> {
917    let stat = std::fs::read_to_string("/proc/thread-self/stat").ok()?;
918    // the comm field may contain spaces and parentheses: split after the last ')'.
919    let rest = &stat[stat.rfind(')')? + 1..];
920    let fields: Vec<&str> = rest.split_whitespace().collect();
921    // fields after the comm: state(0) ppid(1) … utime(11) stime(12)
922    let utime: u64 = fields.get(11)?.parse().ok()?;
923    let stime: u64 = fields.get(12)?.parse().ok()?;
924    Some((utime + stime) * 10)
925}
926
927/// Alignment state: which input edges delivered the current barrier, and the data held back on
928/// them until the rest catch up.
929#[derive(Default)]
930struct Align {
931    epoch: Option<u64>,
932    delivered: HashSet<EdgeId>,
933    held: VecDeque<(EdgeId, Msg)>,
934}
935
936struct Task {
937    id: TaskId,
938    node: Node,
939    rx: Receiver<(EdgeId, Msg)>,
940    /// The coordinator's channel (sources only).
941    ctl: Option<Receiver<Control>>,
942    out_routes: Vec<OutRoute>,
943    in_edges: Vec<EdgeId>,
944    events: SyncSender<Event>,
945    /// When this task last forwarded a tick (at most ten a second, whatever arrives).
946    last_tick: Option<Instant>,
947    /// Every task's probe; this task writes its own (`probes[id]`).
948    probes: Arc<Vec<Probe>>,
949    /// latest watermark per (input edge, source): `None` = that source is idle.
950    wm: HashMap<(EdgeId, TaskId), Option<i64>>,
951    last_wm: Option<i64>,
952    align: Align,
953    eos_seen: HashSet<EdgeId>,
954}
955
956impl Task {
957    fn run(&mut self) {
958        // a panic inside an operator, source or sink ends the BUILD, loudly: the coordinator
959        // hears `Failed` and stops the graph. (Left uncaught, the thread dies, its downstream sees
960        // a closed edge and finishes, its upstream keeps polling into nothing, and the pipeline
961        // sits there consuming — the runtime's stall report found exactly that.)
962        let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match &self.node {
963            Node::Source(_) => self.run_source(),
964            _ => self.run_consumer(),
965        }));
966        if let Err(payload) = outcome {
967            let what = payload
968                .downcast_ref::<String>()
969                .cloned()
970                .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
971                .unwrap_or_else(|| "panic".into());
972            let _ = self.events.send(Event::Failed {
973                task: self.id,
974                error: format!("task {} panicked: {what}", self.id),
975            });
976            // No Eos: Eos means a clean end. The edges out of this task close with it, and a
977            // downstream task whose input closes with this edge still open reports the loss
978            // (`input_lost`) — the run fails rather than finishing on partial input.
979        }
980        self.state(TaskState::Done, u32::MAX);
981        let _ = self.events.send(Event::Finished {
982            task: self.id,
983            cpu_ms: thread_cpu_ms().unwrap_or(0),
984        });
985    }
986
987    fn state(&self, s: TaskState, target: TaskId) {
988        self.probes[self.id as usize].set(s, target);
989    }
990
991    fn run_source(&mut self) {
992        let ctl = self.ctl.take().expect("a source has a control channel");
993        // the last barrier taken, and the last epoch committed (a commit can arrive while the
994        // source is still polling — it must not be waited for twice).
995        let mut last_epoch: u64 = 0;
996        let mut committed: u64 = 0;
997        let mut last_tick = Instant::now();
998        let mut stopping = false;
999        let mut failed = false;
1000        let mut drained = false;
1001        loop {
1002            // the coordinator first: barriers and commits, in the order issued.
1003            while let Ok(c) = ctl.try_recv() {
1004                match c {
1005                    Control::Barrier(epoch) => {
1006                        last_epoch = epoch;
1007                        self.source_barrier(epoch);
1008                    }
1009                    Control::Commit(epoch) => {
1010                        committed = committed.max(epoch);
1011                        self.source_commit(epoch);
1012                    }
1013                    Control::Stop => stopping = true,
1014                }
1015            }
1016            if stopping {
1017                break;
1018            }
1019            let Node::Source(s) = &mut self.node else {
1020                unreachable!()
1021            };
1022            let mut out = Out::default();
1023            self.probes[self.id as usize].set(TaskState::Polling, u32::MAX);
1024            let poll = s.poll(&mut out);
1025            let quiet = out.batches.is_empty();
1026            self.emit(out);
1027            if quiet && last_tick.elapsed() >= Duration::from_millis(100) {
1028                last_tick = Instant::now();
1029                self.broadcast(Msg::Tick { now_ms: now_ms() });
1030            }
1031            match poll {
1032                Poll::Again => {}
1033                Poll::Idle => std::thread::sleep(Duration::from_millis(5)),
1034                Poll::Done => {
1035                    drained = true;
1036                    break;
1037                }
1038                Poll::Failed(error) => {
1039                    let _ = self.events.send(Event::Failed { task: self.id, error });
1040                    failed = true;
1041                    break;
1042                }
1043            }
1044        }
1045        if let Node::Source(s) = &mut self.node {
1046            if drained {
1047                // the final position, for every checkpoint completed after this source ended
1048                let pos = s.position();
1049                if !pos.is_empty() {
1050                    let _ = self.events.send(Event::Drained {
1051                        task: self.id,
1052                        snap: OpSnapshot::whole(pos),
1053                    });
1054                }
1055            }
1056            s.on_stop();
1057        }
1058        if !failed {
1059            self.broadcast(Msg::Eos); // a clean end only: a failed source's edges just close
1060        }
1061        // the last barrier taken commits once every sink acked it: wait for that commit, unless it
1062        // already came. A closed control channel means the coordinator is done with us
1063        // (uncommitted = replay). Barriers issued after Eos are not taken: downstream treats this
1064        // edge as aligned, and the epoch completes once this task finishes.
1065        if !failed && last_epoch > committed {
1066            self.state(TaskState::CommitWait, u32::MAX);
1067            while let Ok(c) = ctl.recv() {
1068                if let Control::Commit(epoch) = c {
1069                    self.source_commit(epoch);
1070                    if epoch >= last_epoch {
1071                        break;
1072                    }
1073                }
1074            }
1075        }
1076    }
1077
1078    fn source_barrier(&mut self, epoch: u64) {
1079        let Node::Source(s) = &mut self.node else {
1080            unreachable!()
1081        };
1082        let snap = OpSnapshot::whole(s.on_barrier(epoch));
1083        let _ = self.events.send(Event::Snapshot {
1084            task: self.id,
1085            epoch,
1086            snap,
1087        });
1088        self.broadcast(Msg::Barrier { epoch });
1089    }
1090
1091    fn source_commit(&mut self, epoch: u64) {
1092        let Node::Source(s) = &mut self.node else {
1093            unreachable!()
1094        };
1095        s.on_commit(epoch);
1096    }
1097
1098    /// The inbox closed before every in-edge delivered its Eos: an upstream task — or the worker it
1099    /// runs on — is gone without ending its stream. A clean end always arrives as Eos on every edge
1100    /// (a task broadcasts Eos before its senders drop; a peer's outbound bridge sends it before the
1101    /// link closes), so a closed inbox with an edge still open is a loss, never a drain: the run
1102    /// fails instead of finishing on partial input. Without this a worker that died during a
1103    /// bounded drain looked like "all upstreams done" and the leader stopped clean, its share of
1104    /// the output silently missing. Reached at most once per task — never on the data path.
1105    #[cold]
1106    fn input_lost(&mut self) {
1107        let open: Vec<EdgeId> = self
1108            .in_edges
1109            .iter()
1110            .copied()
1111            .filter(|e| !self.eos_seen.contains(e))
1112            .collect();
1113        if open.is_empty() {
1114            return; // every edge ended cleanly; the senders simply dropped after their Eos.
1115        }
1116        let _ = self.events.send(Event::Failed {
1117            task: self.id,
1118            error: format!(
1119                "task {}: input closed before end-of-stream on edge(s) {open:?} — an upstream task, or the \
1120                 worker running it, is gone",
1121                self.id
1122            ),
1123        });
1124    }
1125
1126    fn run_consumer(&mut self) {
1127        let ctl = self.ctl.take();
1128        let mut last_epoch: u64 = 0;
1129        let mut committed: u64 = 0;
1130        loop {
1131            // the coordinator's commits reach a sink between messages (ticks keep the inbox
1132            // moving at least ten times a second, so a commit waits at most that long).
1133            if let Some(ctl) = &ctl {
1134                while let Ok(c) = ctl.try_recv() {
1135                    if let Control::Commit(epoch) = c {
1136                        committed = committed.max(epoch);
1137                        if let Node::Sink(s) = &mut self.node {
1138                            s.on_commit(epoch);
1139                        }
1140                    }
1141                }
1142            }
1143            self.state(TaskState::Receiving, u32::MAX);
1144            let Ok((edge, msg)) = self.rx.recv() else {
1145                self.input_lost();
1146                break;
1147            };
1148            if let Msg::Barrier { epoch } = &msg {
1149                last_epoch = last_epoch.max(*epoch);
1150            }
1151            if self.receive(edge, msg) {
1152                break;
1153            }
1154        }
1155        // a sink's last epoch commits once every task reported it: wait for that commit, so an
1156        // exactly-once transaction is never left open at a clean stop (a closed control channel
1157        // means the coordinator is done with us: uncommitted = replay).
1158        if let (Some(ctl), Node::Sink(_)) = (&ctl, &self.node) {
1159            if last_epoch > committed {
1160                self.state(TaskState::CommitWait, u32::MAX);
1161                while let Ok(c) = ctl.recv() {
1162                    if let Control::Commit(epoch) = c {
1163                        if let Node::Sink(s) = &mut self.node {
1164                            s.on_commit(epoch);
1165                        }
1166                        if epoch >= last_epoch {
1167                            break;
1168                        }
1169                    }
1170                }
1171            }
1172        }
1173    }
1174
1175    /// Hold or handle: an edge that already delivered the pending barrier holds EVERYTHING it
1176    /// sends after it — data, watermarks, its next barrier, Eos — in order, until the other edges
1177    /// catch up. (A later barrier let through early would restart alignment and lose the pending
1178    /// epoch: the source behind it then waits forever for a commit that cannot come.)
1179    fn receive(&mut self, edge: EdgeId, msg: Msg) -> bool {
1180        if self.align.epoch.is_some() && self.align.delivered.contains(&edge) {
1181            self.align.held.push_back((edge, msg));
1182            return false;
1183        }
1184        self.handle(edge, msg)
1185    }
1186
1187    /// Handle one message; `true` when the task is done (every input reached Eos).
1188    fn handle(&mut self, edge: EdgeId, msg: Msg) -> bool {
1189        match msg {
1190            Msg::Data(b) => {
1191                self.state(TaskState::Working, u32::MAX);
1192                let mut out = Out::default();
1193                match &mut self.node {
1194                    Node::Operator(o) => o.on_data(edge, b, &mut out),
1195                    Node::Sink(s) => s.on_data(b),
1196                    Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
1197                }
1198                self.emit(out);
1199            }
1200            Msg::Watermark { source, ts } => {
1201                self.wm.insert((edge, source), ts);
1202                // the combined time is the minimum over the active inputs — and it holds until
1203                // EVERY input edge has spoken once (a watermark or the idle marker): the minimum
1204                // over the inputs heard so far would let a fast upstream fire windows the slow
1205                // one still has rows for.
1206                let every_edge_spoke = self.in_edges.iter().all(|e| self.wm.keys().any(|(edge, _)| edge == e));
1207                let active: Vec<i64> = self.wm.values().filter_map(|t| *t).collect();
1208                let combined = if active.is_empty() || !every_edge_spoke {
1209                    None
1210                } else {
1211                    active.iter().copied().min()
1212                };
1213                if let Some(wm) = combined {
1214                    if self.last_wm.is_none_or(|l| wm > l) {
1215                        self.last_wm = Some(wm);
1216                        let mut out = Out::default();
1217                        if let Node::Operator(o) = &mut self.node {
1218                            o.on_watermark(wm, now_ms(), &mut out);
1219                            // an operator's own time is the combined input time it has processed.
1220                            if out.watermark.is_none() {
1221                                out.watermark = Some(Some(wm));
1222                            }
1223                        }
1224                        self.emit(out);
1225                    }
1226                } else if active.is_empty() && (self.last_wm.is_some() || every_edge_spoke) {
1227                    // every upstream source is idle: say so downstream.
1228                    self.broadcast(Msg::Watermark {
1229                        source: self.id,
1230                        ts: None,
1231                    });
1232                }
1233            }
1234            Msg::Barrier { epoch } => {
1235                let a = &mut self.align;
1236                if a.epoch != Some(epoch) {
1237                    a.epoch = Some(epoch);
1238                    a.delivered.clear();
1239                }
1240                a.delivered.insert(edge);
1241                if self.aligned() {
1242                    return self.complete_barrier(epoch);
1243                }
1244            }
1245            Msg::Tick { now_ms } => {
1246                let mut out = Out::default();
1247                if let Node::Operator(o) = &mut self.node {
1248                    o.on_tick(now_ms, &mut out);
1249                }
1250                self.emit(out);
1251                if self.last_tick.is_none_or(|t| t.elapsed() >= Duration::from_millis(100)) {
1252                    self.last_tick = Some(Instant::now());
1253                    self.broadcast(Msg::Tick { now_ms });
1254                }
1255            }
1256            Msg::Eos => {
1257                self.eos_seen.insert(edge);
1258                // an edge at Eos never delivers another barrier: it counts as aligned — and a
1259                // barrier the other edges already delivered completes NOW (a source that finished
1260                // between two epochs never took the later one; the pending epoch must not wait
1261                // for it forever).
1262                if let Some(epoch) = self.align.epoch {
1263                    if self.aligned() && self.complete_barrier(epoch) {
1264                        return true;
1265                    }
1266                }
1267                if self.in_edges.iter().all(|e| self.eos_seen.contains(e)) {
1268                    let mut out = Out::default();
1269                    match &mut self.node {
1270                        Node::Operator(o) => o.on_eos(&mut out),
1271                        Node::Sink(s) => s.on_eos(),
1272                        Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
1273                    }
1274                    self.emit(out);
1275                    self.broadcast(Msg::Eos);
1276                    return true;
1277                }
1278            }
1279        }
1280        false
1281    }
1282
1283    /// Every input edge delivered the pending barrier, or is at Eos.
1284    fn aligned(&self) -> bool {
1285        self.in_edges
1286            .iter()
1287            .all(|e| self.align.delivered.contains(e) || self.eos_seen.contains(e))
1288    }
1289
1290    /// The pending barrier aligned: snapshot (or flush and ack), forward it, then replay what
1291    /// alignment held — in arrival order and through the hold check, so a held barrier for the
1292    /// next epoch starts the next alignment and what followed it on that edge is held again.
1293    /// `true` when the replay reached the end of every input.
1294    fn complete_barrier(&mut self, epoch: u64) -> bool {
1295        self.state(TaskState::Barrier, u32::MAX);
1296        let mut out = Out::default();
1297        match &mut self.node {
1298            Node::Operator(o) => match o.on_barrier(epoch, &mut out) {
1299                Ok(snap) => {
1300                    let _ = self.events.send(Event::Snapshot {
1301                        task: self.id,
1302                        epoch,
1303                        snap,
1304                    });
1305                }
1306                Err(error) => {
1307                    // The checkpoint could not be taken: the build fails with the operator's own
1308                    // diagnostic — the end state a panic in the operator reaches (see `run`),
1309                    // reached without one. The barrier is not forwarded; this task is done.
1310                    let _ = self.events.send(Event::Failed {
1311                        task: self.id,
1312                        error: format!("task {}: checkpoint at epoch {epoch}: {error}", self.id),
1313                    });
1314                    return true; // no Eos (a clean end); downstream sees the edge close
1315                }
1316            },
1317            Node::Sink(s) => {
1318                if s.on_barrier(epoch) {
1319                    let _ = self.events.send(Event::Ack { task: self.id, epoch });
1320                }
1321            }
1322            Node::Source(_) | Node::Deferred(_) | Node::DeferredSink(_) => unreachable!(),
1323        }
1324        self.emit(out);
1325        self.broadcast(Msg::Barrier { epoch });
1326        self.align.epoch = None;
1327        self.align.delivered.clear();
1328        let held: Vec<(EdgeId, Msg)> = self.align.held.drain(..).collect();
1329        for (e, m) in held {
1330            if self.receive(e, m) {
1331                return true;
1332            }
1333        }
1334        false
1335    }
1336
1337    /// Route an operator's outputs along every out-route.
1338    fn emit(&mut self, out: Out) {
1339        for b in out.batches {
1340            for (route, senders) in &self.out_routes {
1341                match route {
1342                    Route::Forward(_) | Route::Broadcast(_) => {
1343                        for (e, target, tx) in senders {
1344                            self.probes[self.id as usize].set(TaskState::Sending, *target);
1345                            let _ = tx.send((*e, Msg::Data(b.clone())));
1346                        }
1347                    }
1348                    Route::Shuffle {
1349                        columns,
1350                        vnodes,
1351                        targets,
1352                        hasher,
1353                    } => {
1354                        let parts = shuffle(&b, columns, *vnodes, targets, hasher.as_ref());
1355                        for ((e, target, tx), part) in senders.iter().zip(parts) {
1356                            if let Some(p) = part {
1357                                self.probes[self.id as usize].set(TaskState::Sending, *target);
1358                                let _ = tx.send((*e, Msg::Data(p)));
1359                            }
1360                        }
1361                    }
1362                }
1363            }
1364        }
1365        if let Some(ts) = out.watermark {
1366            self.broadcast(Msg::Watermark { source: self.id, ts });
1367        }
1368        self.state(TaskState::Working, u32::MAX);
1369    }
1370
1371    fn broadcast(&self, msg: Msg) {
1372        for (_, senders) in &self.out_routes {
1373            for (e, target, tx) in senders {
1374                self.probes[self.id as usize].set(TaskState::Sending, *target);
1375                let _ = tx.send((*e, msg.clone()));
1376            }
1377        }
1378        self.state(TaskState::Working, u32::MAX);
1379    }
1380}
1381
1382fn now_ms() -> i64 {
1383    chrono::Utc::now().timestamp_millis()
1384}
1385
1386/// Split a batch by vnode: `hash(columns) % vnodes` → the target whose range holds it. Returns one
1387/// optional sub-batch per target, in target order. Missing routing columns hash as null.
1388pub fn shuffle(
1389    b: &RecordBatch,
1390    columns: &[String],
1391    vnodes: u32,
1392    targets: &[(TaskId, std::ops::Range<u32>)],
1393    hasher: &(dyn Fn(&[arrow::array::ArrayRef], &mut [u64]) + Send + Sync),
1394) -> Vec<Option<RecordBatch>> {
1395    let n = b.num_rows();
1396    let vnode_of: Vec<u32> = if columns.is_empty() {
1397        vec![0; n]
1398    } else {
1399        let cols: Vec<arrow::array::ArrayRef> = columns
1400            .iter()
1401            .map(|c| match b.column_by_name(c) {
1402                Some(a) => Arc::clone(a),
1403                None => arrow::array::new_null_array(&arrow::datatypes::DataType::Null, n),
1404            })
1405            .collect();
1406        let mut hashes = vec![0u64; n];
1407        hasher(&cols, &mut hashes);
1408        hashes.iter().map(|h| (h % vnodes as u64) as u32).collect()
1409    };
1410    targets
1411        .iter()
1412        .map(|(_, range)| {
1413            let idx: Vec<u32> = (0..n as u32)
1414                .filter(|i| range.contains(&vnode_of[*i as usize]))
1415                .collect();
1416            if idx.is_empty() {
1417                None
1418            } else if idx.len() == n {
1419                Some(b.clone())
1420            } else {
1421                let idx = UInt32Array::from(idx);
1422                let cols: Vec<arrow::array::ArrayRef> = b
1423                    .columns()
1424                    .iter()
1425                    .map(|c| take(c.as_ref(), &idx, None).expect("take"))
1426                    .collect();
1427                Some(RecordBatch::try_new(b.schema(), cols).expect("shuffled batch"))
1428            }
1429        })
1430        .collect()
1431}
1432
1433/// Contiguous vnode ranges for `tasks` tasks over `vnodes` vnodes.
1434pub fn vnode_ranges(vnodes: u32, tasks: u32) -> Vec<std::ops::Range<u32>> {
1435    let tasks = tasks.max(1).min(vnodes.max(1));
1436    (0..tasks)
1437        .map(|t| {
1438            let start = vnodes * t / tasks;
1439            let end = vnodes * (t + 1) / tasks;
1440            start..end
1441        })
1442        .collect()
1443}
1444
1445#[cfg(test)]
1446mod tests {
1447    use super::*;
1448
1449    /// A deterministic hash of each row's routing cells (their display strings), for the tests.
1450    pub(super) fn test_hasher() -> Hasher {
1451        Arc::new(|cols: &[arrow::array::ArrayRef], out: &mut [u64]| {
1452            use std::hash::{Hash, Hasher as _};
1453            for (i, slot) in out.iter_mut().enumerate() {
1454                let mut h = std::collections::hash_map::DefaultHasher::new();
1455                for c in cols {
1456                    arrow::util::display::array_value_to_string(c, i)
1457                        .unwrap_or_default()
1458                        .hash(&mut h);
1459                }
1460                *slot = h.finish();
1461            }
1462        })
1463    }
1464    use arrow::array::Array;
1465    use arrow::array::{Float64Array, Int64Array, StringArray};
1466    use arrow::datatypes::{DataType, Field, Schema};
1467    use std::sync::atomic::AtomicUsize;
1468
1469    fn batch(keys: &[&str], ts: &[i64]) -> RecordBatch {
1470        let schema = Arc::new(Schema::new(vec![
1471            Field::new("k", DataType::Utf8, true),
1472            Field::new("ts", DataType::Int64, true),
1473            Field::new("v", DataType::Float64, true),
1474        ]));
1475        let k: StringArray = keys.to_vec().into();
1476        let t: Int64Array = ts.to_vec().into();
1477        let v: Float64Array = ts.iter().map(|x| *x as f64).collect::<Vec<_>>().into();
1478        RecordBatch::try_new(schema, vec![Arc::new(k), Arc::new(t), Arc::new(v)]).unwrap()
1479    }
1480
1481    /// A bounded source: emits its batches one per poll (with a watermark), then Done.
1482    struct VecSource {
1483        batches: VecDeque<RecordBatch>,
1484    }
1485    impl Source for VecSource {
1486        fn poll(&mut self, out: &mut Out) -> Poll {
1487            match self.batches.pop_front() {
1488                Some(b) => {
1489                    let ts = b
1490                        .column_by_name("ts")
1491                        .and_then(|c| arrow::compute::max(c.as_any().downcast_ref::<Int64Array>().unwrap()));
1492                    out.push(b);
1493                    out.watermark = Some(ts);
1494                    Poll::Again
1495                }
1496                None => Poll::Done,
1497            }
1498        }
1499        fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
1500            format!("src@{epoch}").into_bytes()
1501        }
1502        fn position(&mut self) -> Vec<u8> {
1503            b"src@done".to_vec()
1504        }
1505    }
1506
1507    /// Counts rows per key and emits the count on every watermark; snapshots its counts.
1508    #[derive(Default)]
1509    struct CountOp {
1510        counts: HashMap<String, i64>,
1511        fired: Vec<i64>,
1512    }
1513    impl Operator for CountOp {
1514        fn on_data(&mut self, _edge: EdgeId, batch: RecordBatch, _out: &mut Out) {
1515            let k = batch.column_by_name("k").unwrap();
1516            let k = k.as_any().downcast_ref::<StringArray>().unwrap();
1517            for i in 0..k.len() {
1518                *self.counts.entry(k.value(i).to_string()).or_default() += 1;
1519            }
1520        }
1521        fn on_watermark(&mut self, wm: i64, _now: i64, _out: &mut Out) {
1522            self.fired.push(wm);
1523        }
1524        fn on_barrier(&mut self, epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
1525            Ok(OpSnapshot::whole(
1526                format!("op@{epoch}:{}", self.counts.values().sum::<i64>()).into_bytes(),
1527            ))
1528        }
1529        fn on_eos(&mut self, out: &mut Out) {
1530            // emit the counts as a batch: one row per key.
1531            let mut keys: Vec<&String> = self.counts.keys().collect();
1532            keys.sort();
1533            let ks: StringArray = keys.iter().map(|k| k.as_str()).collect::<Vec<_>>().into();
1534            let ns: Int64Array = keys.iter().map(|k| self.counts[*k]).collect::<Vec<_>>().into();
1535            let schema = Arc::new(Schema::new(vec![
1536                Field::new("k", DataType::Utf8, true),
1537                Field::new("n", DataType::Int64, true),
1538            ]));
1539            out.push(RecordBatch::try_new(schema, vec![Arc::new(ks), Arc::new(ns)]).unwrap());
1540        }
1541    }
1542
1543    /// Collects everything it receives.
1544    struct CollectSink {
1545        rows: Arc<Mutex<Vec<(String, i64)>>>,
1546        barriers: Arc<AtomicUsize>,
1547    }
1548    impl Sink for CollectSink {
1549        fn on_data(&mut self, batch: RecordBatch) {
1550            let k = batch.column_by_name("k").unwrap();
1551            let k = k.as_any().downcast_ref::<StringArray>().unwrap();
1552            let n = batch.column_by_name("n").unwrap();
1553            let n = n.as_any().downcast_ref::<Int64Array>().unwrap();
1554            let mut rows = self.rows.lock().unwrap();
1555            for i in 0..k.len() {
1556                rows.push((k.value(i).to_string(), n.value(i)));
1557            }
1558        }
1559        fn on_barrier(&mut self, _epoch: u64) -> bool {
1560            self.barriers.fetch_add(1, Ordering::Relaxed);
1561            true
1562        }
1563        fn on_eos(&mut self) {}
1564    }
1565
1566    /// An unbounded source (idle until stopped) that remembers what it was told to commit.
1567    struct IdleSource {
1568        polls: usize,
1569        fail_at: Option<usize>,
1570        commits: Arc<Mutex<Vec<u64>>>,
1571        barriers: Arc<Mutex<Vec<u64>>>,
1572    }
1573    impl Source for IdleSource {
1574        fn poll(&mut self, out: &mut Out) -> Poll {
1575            self.polls += 1;
1576            if self.fail_at == Some(self.polls) {
1577                return Poll::Failed("boom".into());
1578            }
1579            out.push(batch(&["a"], &[self.polls as i64]));
1580            Poll::Idle
1581        }
1582        fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
1583            self.barriers.lock().unwrap().push(epoch);
1584            epoch.to_le_bytes().to_vec()
1585        }
1586        fn on_commit(&mut self, epoch: u64) {
1587            self.commits.lock().unwrap().push(epoch);
1588        }
1589    }
1590
1591    /// A sink that records the order of what it sees: `b<epoch>` for a barrier, `e` for Eos, `d`
1592    /// for data.
1593    struct TraceSink {
1594        trace: Arc<Mutex<Vec<String>>>,
1595    }
1596    impl Sink for TraceSink {
1597        fn on_data(&mut self, _b: RecordBatch) {
1598            self.trace.lock().unwrap().push("d".into());
1599        }
1600        fn on_barrier(&mut self, epoch: u64) -> bool {
1601            self.trace.lock().unwrap().push(format!("b{epoch}"));
1602            true
1603        }
1604        fn on_eos(&mut self) {
1605            self.trace.lock().unwrap().push("e".into());
1606        }
1607    }
1608
1609    #[test]
1610    fn a_shuffled_pipeline_counts_every_key_exactly_once_across_two_operator_tasks() {
1611        let mut g = Graph::new(16);
1612        let data: Vec<RecordBatch> = (0..20)
1613            .map(|i| batch(&["a", "b", "c", "d"], &[i * 10, i * 10 + 1, i * 10 + 2, i * 10 + 3]))
1614            .collect();
1615        let src = g.source(Box::new(VecSource { batches: data.into() }));
1616        let op0 = g.operator(Box::new(CountOp::default()));
1617        let op1 = g.operator(Box::new(CountOp::default()));
1618        let rows = Arc::new(Mutex::new(Vec::new()));
1619        let barriers = Arc::new(AtomicUsize::new(0));
1620        let sink = g.sink(Box::new(CollectSink {
1621            rows: Arc::clone(&rows),
1622            barriers: Arc::clone(&barriers),
1623        }));
1624        let ranges = vnode_ranges(64, 2);
1625        g.route(
1626            src,
1627            Route::Shuffle {
1628                columns: vec!["k".into()],
1629                vnodes: 64,
1630                targets: vec![(op0, ranges[0].clone()), (op1, ranges[1].clone())],
1631                hasher: test_hasher(),
1632            },
1633        );
1634        g.route(op0, Route::Forward(sink));
1635        g.route(op1, Route::Forward(sink));
1636        let (etx, erx) = sync_channel(1024);
1637        let running = g.start(etx).expect("start");
1638        let mut tracker = running.tracker();
1639        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
1640        running.join();
1641        let mut got: Vec<(String, i64)> = rows.lock().unwrap().clone();
1642        got.sort();
1643        assert_eq!(
1644            got,
1645            vec![("a".into(), 20), ("b".into(), 20), ("c".into(), 20), ("d".into(), 20)],
1646            "every key counted exactly once, on exactly one task"
1647        );
1648        assert_eq!(tracker.finished(), 4);
1649    }
1650
1651    #[test]
1652    fn a_peer_that_dies_before_eos_fails_the_worker_instead_of_finishing_it() {
1653        // Worker 1 runs only the sink, fed by a remote edge from a peer. The peer dies (its runtime
1654        // — and so its sockets — go away) without ever sending Eos. The link closes, the inbound
1655        // bridge ends, the sink's inbox closes with the edge still open: that must surface as the
1656        // run's FAILURE naming the edge, not as a clean finish on partial input (the leader used to
1657        // stop clean and the dead worker's share of the output was silently missing).
1658        let rt_peer = tokio::runtime::Builder::new_multi_thread()
1659            .worker_threads(2)
1660            .enable_all()
1661            .build()
1662            .unwrap();
1663        let rt = tokio::runtime::Builder::new_multi_thread()
1664            .worker_threads(2)
1665            .enable_all()
1666            .build()
1667            .unwrap();
1668        let listener = rt.block_on(tokio::net::TcpListener::bind("127.0.0.1:0")).unwrap();
1669        let addr = listener.local_addr().unwrap();
1670        let srv = rt.spawn(async move { fv_streams_exchange::Link::accept(&listener, 1, 1 << 20).await.unwrap() });
1671        let peer = rt_peer
1672            .block_on(fv_streams_exchange::Link::connect(addr, 1, 1 << 20))
1673            .unwrap();
1674        let (_w1_send, w1_from_peer) = rt.block_on(srv).unwrap().into_split();
1675
1676        let rows = Arc::new(Mutex::new(Vec::new()));
1677        let mut g = Graph::new(16);
1678        let src = g.source(Box::new(VecSource {
1679            batches: Vec::new().into(),
1680        }));
1681        let op = g.operator(Box::new(CountOp::default()));
1682        let sink = g.sink(Box::new(CollectSink {
1683            rows: Arc::clone(&rows),
1684            barriers: Arc::new(AtomicUsize::new(0)),
1685        }));
1686        g.route(src, Route::Forward(op));
1687        g.route(op, Route::Forward(sink));
1688        let of_task = vec![0u32, 0, 1]; // src + op on the (dead) peer, the sink here
1689        let (etx, erx) = sync_channel(1024);
1690        let r1 = g
1691            .start_worker(
1692                etx,
1693                &of_task,
1694                1,
1695                Some(WorkerExchange {
1696                    handle: rt.handle().clone(),
1697                    senders: HashMap::new(),
1698                    receivers: vec![w1_from_peer],
1699                }),
1700            )
1701            .expect("start worker 1");
1702        // the peer dies mid-stream: nothing sent, no Eos.
1703        drop(peer);
1704        drop(rt_peer);
1705        let mut t1 = r1.tracker();
1706        let err = r1
1707            .finish(&erx, &mut t1, |_| Ok(()))
1708            .expect_err("a lost upstream fails the run");
1709        assert!(
1710            err.contains("input closed before end-of-stream") && err.contains("is gone"),
1711            "names the loss: {err}"
1712        );
1713        assert!(t1.all_finished(), "the sink task still ended (no hang)");
1714        r1.join();
1715        assert!(
1716            rows.lock().unwrap().is_empty(),
1717            "no partial output was presented as final"
1718        );
1719        drop(rt);
1720    }
1721
1722    #[test]
1723    fn two_workers_split_by_placement_produce_the_same_result_as_one_process() {
1724        // The per-worker gate: the SAME `src -> op -> sink` graph, built identically on two workers
1725        // and split by placement — src+op on worker 0, sink on worker 1 — running as two independent
1726        // coordinators in one process, connected by a loopback Link. The op->sink edge is the only
1727        // cross-worker edge; the aggregate, watermarks and the final Eos travel as Arrow-IPC/postcard
1728        // frames over TCP and land in the sink exactly as a local edge would. Each worker instantiates
1729        // ONLY its own tasks (proven by each tracker's finished count) and the result is identical to
1730        // the single-process run — placement is transparent to correctness.
1731        let rt = tokio::runtime::Builder::new_multi_thread()
1732            .worker_threads(2)
1733            .enable_all()
1734            .build()
1735            .unwrap();
1736        let handle = rt.handle().clone();
1737
1738        // One link between worker 0 and worker 1: w0 sends the op->sink edge, w1 receives it.
1739        let (w0_to_w1, w1_from_w0) = handle.block_on(async {
1740            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1741            let addr = listener.local_addr().unwrap();
1742            let srv =
1743                tokio::spawn(async move { fv_streams_exchange::Link::accept(&listener, 1, 1 << 20).await.unwrap() });
1744            let a = fv_streams_exchange::Link::connect(addr, 1, 1 << 20).await.unwrap();
1745            let b = srv.await.unwrap();
1746            let (a_send, _a_recv) = a.into_split();
1747            let (_b_send, b_recv) = b.into_split();
1748            (a_send, b_recv)
1749        });
1750
1751        // Every worker builds the identical graph (deterministic task/edge ids). The sink collects
1752        // into a shared buffer so the test can read worker 1's output.
1753        let rows = Arc::new(Mutex::new(Vec::new()));
1754        let build = || {
1755            let mut g = Graph::new(16);
1756            let data: Vec<RecordBatch> = (0..20)
1757                .map(|i| batch(&["a", "b", "c", "d"], &[i * 10, i * 10 + 1, i * 10 + 2, i * 10 + 3]))
1758                .collect();
1759            let src = g.source(Box::new(VecSource { batches: data.into() }));
1760            let op = g.operator(Box::new(CountOp::default()));
1761            let sink = g.sink(Box::new(CollectSink {
1762                rows: Arc::clone(&rows),
1763                barriers: Arc::new(AtomicUsize::new(0)),
1764            }));
1765            g.route(src, Route::Forward(op));
1766            g.route(op, Route::Forward(sink));
1767            (g, src, op, sink)
1768        };
1769        let (g0, _s, _o, _k) = build();
1770        let (g1, _s1, _o1, _k1) = build();
1771        // Placement: src (0) and op (1) on worker 0; sink (2) on worker 1.
1772        let of_task = vec![0u32, 0, 1];
1773
1774        // Worker 1 (the sink): receives the op->sink edge from worker 0; sends nothing.
1775        let (etx1, erx1) = sync_channel(1024);
1776        let r1 = g1
1777            .start_worker(
1778                etx1,
1779                &of_task,
1780                1,
1781                Some(WorkerExchange {
1782                    handle: handle.clone(),
1783                    senders: HashMap::new(),
1784                    receivers: vec![w1_from_w0],
1785                }),
1786            )
1787            .expect("start worker 1");
1788        let w1 = std::thread::spawn(move || {
1789            let mut t1 = r1.tracker();
1790            r1.finish(&erx1, &mut t1, |_| Ok(())).unwrap();
1791            r1.join();
1792            t1.finished()
1793        });
1794
1795        // Worker 0 (source + operator): sends the op->sink edge to worker 1.
1796        let (etx0, erx0) = sync_channel(1024);
1797        let mut senders = HashMap::new();
1798        senders.insert(1u32, w0_to_w1);
1799        let r0 = g0
1800            .start_worker(
1801                etx0,
1802                &of_task,
1803                0,
1804                Some(WorkerExchange {
1805                    handle: handle.clone(),
1806                    senders,
1807                    receivers: vec![],
1808                }),
1809            )
1810            .expect("start worker 0");
1811        let mut t0 = r0.tracker();
1812        r0.finish(&erx0, &mut t0, |_| Ok(())).unwrap();
1813        r0.join();
1814        let w1_finished = w1.join().unwrap();
1815
1816        // Each worker ran only its own tasks.
1817        assert_eq!(t0.finished(), 2, "worker 0 ran exactly src + op");
1818        assert_eq!(w1_finished, 1, "worker 1 ran exactly the sink");
1819        // The output is identical to the single-process run.
1820        let mut got: Vec<(String, i64)> = rows.lock().unwrap().clone();
1821        got.sort();
1822        assert_eq!(
1823            got,
1824            vec![("a".into(), 20), ("b".into(), 20), ("c".into(), 20), ("d".into(), 20)],
1825            "the aggregate crossed the worker boundary intact"
1826        );
1827        drop(rt);
1828    }
1829
1830    #[test]
1831    fn the_combined_watermark_holds_until_every_input_has_spoken() {
1832        // REGRESSION (the crash e2e): a fast source's watermark reached the operator before the
1833        // slow source's first one; the minimum over the inputs heard so far fired windows the
1834        // slow source still had rows for. Every input edge must speak before time moves.
1835        struct Src {
1836            wm: i64,
1837            delay_ms: u64,
1838            emitted: bool,
1839        }
1840        impl Source for Src {
1841            fn poll(&mut self, out: &mut Out) -> Poll {
1842                if self.emitted {
1843                    return Poll::Done;
1844                }
1845                std::thread::sleep(Duration::from_millis(self.delay_ms));
1846                self.emitted = true;
1847                out.push(batch(&["x"], &[self.wm]));
1848                out.watermark = Some(Some(self.wm));
1849                Poll::Again
1850            }
1851            fn on_barrier(&mut self, _epoch: u64) -> Vec<u8> {
1852                Vec::new()
1853            }
1854        }
1855        struct WmOp(Arc<Mutex<Vec<i64>>>);
1856        impl Operator for WmOp {
1857            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
1858            fn on_watermark(&mut self, wm: i64, _n: i64, _o: &mut Out) {
1859                self.0.lock().unwrap().push(wm);
1860            }
1861            fn on_barrier(&mut self, _epoch: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
1862                Ok(OpSnapshot::whole(Vec::new()))
1863            }
1864            fn on_eos(&mut self, _o: &mut Out) {}
1865        }
1866        let mut g = Graph::new(4);
1867        let fast = g.source(Box::new(Src {
1868            wm: 1_000,
1869            delay_ms: 0,
1870            emitted: false,
1871        }));
1872        let slow = g.source(Box::new(Src {
1873            wm: 5,
1874            delay_ms: 150,
1875            emitted: false,
1876        }));
1877        let seen = Arc::new(Mutex::new(Vec::new()));
1878        let op = g.operator(Box::new(WmOp(Arc::clone(&seen))));
1879        let sink = g.sink(Box::new(CollectSink {
1880            rows: Arc::new(Mutex::new(Vec::new())),
1881            barriers: Arc::new(AtomicUsize::new(0)),
1882        }));
1883        g.route(fast, Route::Forward(op));
1884        g.route(slow, Route::Forward(op));
1885        g.route(op, Route::Forward(sink));
1886        let (etx, erx) = sync_channel(64);
1887        let running = g.start(etx).expect("start");
1888        let mut tracker = running.tracker();
1889        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
1890        running.join();
1891        let seen = seen.lock().unwrap().clone();
1892        assert_eq!(
1893            seen,
1894            vec![5],
1895            "time moved only once both inputs had spoken, to their minimum"
1896        );
1897    }
1898
1899    #[test]
1900    fn watermarks_take_the_minimum_over_two_sources_and_barriers_align() {
1901        // two sources → one operator: the operator's watermark firings must never exceed the
1902        // slower source's time, and a barrier from both must align before the snapshot.
1903        struct SlowSource {
1904            n: usize,
1905            step: i64,
1906            emitted: usize,
1907        }
1908        impl Source for SlowSource {
1909            fn poll(&mut self, out: &mut Out) -> Poll {
1910                if self.emitted >= self.n {
1911                    return Poll::Done;
1912                }
1913                self.emitted += 1;
1914                let t = self.emitted as i64 * self.step;
1915                out.push(batch(&["x"], &[t]));
1916                out.watermark = Some(Some(t));
1917                std::thread::sleep(Duration::from_millis(2));
1918                Poll::Again
1919            }
1920            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
1921                format!("{epoch}").into_bytes()
1922            }
1923        }
1924        struct WmOp {
1925            seen: Arc<Mutex<Vec<i64>>>,
1926            aligned: Arc<AtomicUsize>,
1927        }
1928        impl Operator for WmOp {
1929            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
1930            fn on_watermark(&mut self, wm: i64, _n: i64, _o: &mut Out) {
1931                self.seen.lock().unwrap().push(wm);
1932            }
1933            fn on_barrier(&mut self, _epoch: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
1934                self.aligned.fetch_add(1, Ordering::Relaxed);
1935                Ok(OpSnapshot::whole(Vec::new()))
1936            }
1937            fn on_eos(&mut self, _o: &mut Out) {}
1938        }
1939        let mut g = Graph::new(4).barrier_every(Duration::from_millis(20));
1940        let fast = g.source(Box::new(SlowSource {
1941            n: 200,
1942            step: 10,
1943            emitted: 0,
1944        }));
1945        let slow = g.source(Box::new(SlowSource {
1946            n: 200,
1947            step: 1,
1948            emitted: 0,
1949        }));
1950        let seen = Arc::new(Mutex::new(Vec::new()));
1951        let aligned = Arc::new(AtomicUsize::new(0));
1952        let op = g.operator(Box::new(WmOp {
1953            seen: Arc::clone(&seen),
1954            aligned: Arc::clone(&aligned),
1955        }));
1956        let sink = g.sink(Box::new(CollectSink {
1957            rows: Arc::new(Mutex::new(Vec::new())),
1958            barriers: Arc::new(AtomicUsize::new(0)),
1959        }));
1960        g.route(fast, Route::Forward(op));
1961        g.route(slow, Route::Forward(op));
1962        g.route(op, Route::Broadcast(vec![sink]));
1963        let (etx, erx) = sync_channel(4096);
1964        let running = g.start(etx).expect("start");
1965        let mut tracker = running.tracker();
1966        let mut completed: Vec<Completed> = Vec::new();
1967        running
1968            .finish(&erx, &mut tracker, |c| {
1969                completed.push(c);
1970                Ok(())
1971            })
1972            .unwrap();
1973        running.join();
1974        let seen = seen.lock().unwrap().clone();
1975        assert!(!seen.is_empty());
1976        assert!(seen.windows(2).all(|w| w[1] > w[0]), "monotonic");
1977        assert!(
1978            *seen.last().unwrap() <= 200,
1979            "never past the slow source's time: {:?}",
1980            seen.last()
1981        );
1982        // the operator snapshotted once per aligned barrier, and every completed epoch carries
1983        // the operator's snapshot plus one per source still running at the time.
1984        assert!(!completed.is_empty(), "the ticker issued epochs");
1985        assert_eq!(
1986            completed
1987                .iter()
1988                .filter(|c| c.snapshots.iter().any(|(t, _)| *t == op))
1989                .count(),
1990            aligned.load(Ordering::Relaxed)
1991        );
1992        let epochs: Vec<u64> = completed.iter().map(|c| c.epoch).collect();
1993        assert!(
1994            epochs.windows(2).all(|w| w[1] == w[0] + 1),
1995            "epochs complete in order: {epochs:?}"
1996        );
1997        assert!(completed.iter().all(|c| c.snapshots.len() <= 3));
1998    }
1999
2000    #[test]
2001    fn a_fast_edge_delivering_several_barriers_ahead_of_a_slow_one_still_aligns_every_epoch() {
2002        // REGRESSION: the fast source takes many barriers while the slow source sits in one poll;
2003        // the operator's fast edge therefore delivers barriers e, e+1, e+2 before the slow edge
2004        // delivers e. Each must wait its turn: every epoch aligns, snapshots, and commits, and the
2005        // fast source (waiting for its last commit) exits — the run must not hang.
2006        struct PacedSource {
2007            polls: usize,
2008            left: usize,
2009            nap: Duration,
2010        }
2011        impl Source for PacedSource {
2012            fn poll(&mut self, out: &mut Out) -> Poll {
2013                if self.left == 0 {
2014                    return Poll::Done;
2015                }
2016                self.left -= 1;
2017                self.polls += 1;
2018                std::thread::sleep(self.nap);
2019                out.push(batch(&["p"], &[self.polls as i64]));
2020                Poll::Again
2021            }
2022            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
2023                epoch.to_le_bytes().to_vec()
2024            }
2025        }
2026        let mut g = Graph::new(2).barrier_every(Duration::from_millis(5));
2027        let fast = g.source(Box::new(PacedSource {
2028            polls: 0,
2029            left: 400,
2030            nap: Duration::from_millis(1),
2031        }));
2032        let slow = g.source(Box::new(PacedSource {
2033            polls: 0,
2034            left: 6,
2035            nap: Duration::from_millis(60),
2036        }));
2037        let op = g.operator(Box::new(CountOp::default()));
2038        let rows = Arc::new(Mutex::new(Vec::new()));
2039        let sink = g.sink(Box::new(CollectSink {
2040            rows: Arc::clone(&rows),
2041            barriers: Arc::new(AtomicUsize::new(0)),
2042        }));
2043        g.route(fast, Route::Forward(op));
2044        g.route(slow, Route::Forward(op));
2045        g.route(op, Route::Forward(sink));
2046        let (etx, erx) = sync_channel(4096);
2047        let running = g.start(etx).expect("start");
2048        let mut tracker = running.tracker();
2049        let mut completed: Vec<Completed> = Vec::new();
2050        // a watchdog: past the deadline it clears the sources' control channels, which releases a
2051        // source stuck waiting for a commit — `finish` then returns and the test reports the hang
2052        // instead of blocking the whole suite forever.
2053        let inner = Arc::clone(&running.inner);
2054        let done = Arc::new(AtomicBool::new(false));
2055        let hung = Arc::new(AtomicBool::new(false));
2056        let watchdog = std::thread::spawn({
2057            let done = Arc::clone(&done);
2058            let hung = Arc::clone(&hung);
2059            move || {
2060                let deadline = Instant::now() + Duration::from_secs(20);
2061                while !done.load(Ordering::Relaxed) {
2062                    if Instant::now() > deadline {
2063                        hung.store(true, Ordering::Relaxed);
2064                        inner.ctl.lock().unwrap().clear();
2065                        return;
2066                    }
2067                    std::thread::sleep(Duration::from_millis(20));
2068                }
2069            }
2070        });
2071        running
2072            .finish(&erx, &mut tracker, |c| {
2073                completed.push(c);
2074                Ok(())
2075            })
2076            .unwrap();
2077        running.join();
2078        done.store(true, Ordering::Relaxed);
2079        watchdog.join().unwrap();
2080        assert!(
2081            !hung.load(Ordering::Relaxed),
2082            "the dataflow hung: alignment lost an epoch"
2083        );
2084        assert_eq!(rows.lock().unwrap().as_slice(), &[("p".to_string(), 406)]);
2085        let epochs: Vec<u64> = completed.iter().map(|c| c.epoch).collect();
2086        assert!(
2087            epochs.len() >= 5,
2088            "several epochs completed while the slow source lagged: {epochs:?}"
2089        );
2090        assert!(epochs.windows(2).all(|w| w[1] == w[0] + 1), "in order: {epochs:?}");
2091        assert!(
2092            completed.iter().all(|c| c.snapshots.iter().any(|(t, _)| *t == op)),
2093            "every completed epoch carries the operator's aligned snapshot"
2094        );
2095    }
2096
2097    #[test]
2098    fn a_source_finishing_between_epochs_does_not_stall_the_pending_alignment() {
2099        // REGRESSION: the bounded source takes no barrier for the epoch issued during its last,
2100        // slow poll and finishes; the live source did take it, so the operator holds that epoch
2101        // pending on the finished source's edge. Its Eos must complete the alignment, or every
2102        // later epoch waits behind it and the live source can never commit.
2103        struct SlowThenDone {
2104            polls: usize,
2105        }
2106        impl Source for SlowThenDone {
2107            fn poll(&mut self, out: &mut Out) -> Poll {
2108                self.polls += 1;
2109                if self.polls > 2 {
2110                    return Poll::Done;
2111                }
2112                std::thread::sleep(Duration::from_millis(40));
2113                out.push(batch(&["s"], &[self.polls as i64]));
2114                Poll::Again
2115            }
2116            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
2117                epoch.to_le_bytes().to_vec()
2118            }
2119        }
2120        let commits = Arc::new(Mutex::new(Vec::new()));
2121        let barriers = Arc::new(Mutex::new(Vec::new()));
2122        let mut g = Graph::new(4).barrier_every(Duration::from_millis(5));
2123        let live = g.source(Box::new(IdleSource {
2124            polls: 0,
2125            fail_at: None,
2126            commits: Arc::clone(&commits),
2127            barriers: Arc::clone(&barriers),
2128        }));
2129        let bounded = g.source(Box::new(SlowThenDone { polls: 0 }));
2130        let op = g.operator(Box::new(CountOp::default()));
2131        let rows = Arc::new(Mutex::new(Vec::new()));
2132        let sink = g.sink(Box::new(CollectSink {
2133            rows: Arc::clone(&rows),
2134            barriers: Arc::new(AtomicUsize::new(0)),
2135        }));
2136        g.route(live, Route::Forward(op));
2137        g.route(bounded, Route::Forward(op));
2138        g.route(op, Route::Forward(sink));
2139        let (etx, erx) = sync_channel(4096);
2140        let running = g.start(etx).expect("start");
2141        let mut tracker = running.tracker();
2142        let mut completed: Vec<u64> = Vec::new();
2143        let inner = Arc::clone(&running.inner);
2144        let done = Arc::new(AtomicBool::new(false));
2145        let hung = Arc::new(AtomicBool::new(false));
2146        let watchdog = std::thread::spawn({
2147            let done = Arc::clone(&done);
2148            let hung = Arc::clone(&hung);
2149            move || {
2150                let deadline = Instant::now() + Duration::from_secs(20);
2151                while !done.load(Ordering::Relaxed) {
2152                    if Instant::now() > deadline {
2153                        hung.store(true, Ordering::Relaxed);
2154                        inner.ctl.lock().unwrap().clear();
2155                        return;
2156                    }
2157                    std::thread::sleep(Duration::from_millis(20));
2158                }
2159            }
2160        });
2161        // let the bounded source finish and a few more epochs pass, then stop the live one.
2162        let until = Instant::now() + Duration::from_millis(250);
2163        while Instant::now() < until {
2164            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
2165                for c in tracker.on_event(&e) {
2166                    running.commit(c.epoch);
2167                    completed.push(c.epoch);
2168                }
2169            }
2170        }
2171        let last = running.stop();
2172        running
2173            .finish(&erx, &mut tracker, |c| {
2174                completed.push(c.epoch);
2175                Ok(())
2176            })
2177            .unwrap();
2178        running.join();
2179        done.store(true, Ordering::Relaxed);
2180        watchdog.join().unwrap();
2181        assert!(
2182            !hung.load(Ordering::Relaxed),
2183            "the pending epoch stalled on the finished source"
2184        );
2185        assert!(
2186            completed.windows(2).all(|w| w[1] == w[0] + 1),
2187            "in order: {completed:?}"
2188        );
2189        assert_eq!(completed.last(), Some(&last));
2190        assert_eq!(
2191            commits.lock().unwrap().last(),
2192            Some(&last),
2193            "the live source committed the final epoch"
2194        );
2195        let taken = barriers.lock().unwrap().len();
2196        assert!(
2197            taken > 10,
2198            "the live source kept taking epochs after the other finished: {taken}"
2199        );
2200        let got = rows.lock().unwrap().clone();
2201        assert_eq!(got.iter().find(|(k, _)| k == "s").map(|(_, n)| *n), Some(2));
2202    }
2203
2204    #[test]
2205    fn a_graph_started_at_an_epoch_issues_the_next_ones() {
2206        // a restored run continues its checkpoint's numbering: restored from epoch 7, its first
2207        // barrier is 8, so its checkpoints always read as newer than the one it restored from.
2208        let barriers = Arc::new(Mutex::new(Vec::new()));
2209        let mut g = Graph::new(4).start_at_epoch(7);
2210        let src = g.source(Box::new(IdleSource {
2211            polls: 0,
2212            fail_at: None,
2213            commits: Arc::new(Mutex::new(Vec::new())),
2214            barriers: Arc::clone(&barriers),
2215        }));
2216        let sink = g.sink(Box::new(TraceSink {
2217            trace: Arc::new(Mutex::new(Vec::new())),
2218        }));
2219        g.route(src, Route::Forward(sink));
2220        let (etx, erx) = sync_channel(64);
2221        let running = g.start(etx).expect("start");
2222        assert_eq!(running.epoch(), 7);
2223        assert_eq!(running.barrier(), 8);
2224        let last = running.stop();
2225        assert_eq!(last, 9);
2226        let mut tracker = running.tracker();
2227        running.finish(&erx, &mut tracker, |_| Ok(())).unwrap();
2228        running.join();
2229        assert_eq!(barriers.lock().unwrap().as_slice(), &[8, 9]);
2230    }
2231
2232    #[test]
2233    fn epochs_complete_when_every_sink_acked_and_the_sources_hear_the_commit() {
2234        let commits = Arc::new(Mutex::new(Vec::new()));
2235        let barriers = Arc::new(Mutex::new(Vec::new()));
2236        let trace = Arc::new(Mutex::new(Vec::new()));
2237        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2238        let src = g.source(Box::new(IdleSource {
2239            polls: 0,
2240            fail_at: None,
2241            commits: Arc::clone(&commits),
2242            barriers: Arc::clone(&barriers),
2243        }));
2244        let sink = g.sink(Box::new(TraceSink {
2245            trace: Arc::clone(&trace),
2246        }));
2247        g.route(src, Route::Forward(sink));
2248        let (etx, erx) = sync_channel(4096);
2249        let running = g.start(etx).expect("start");
2250        let mut tracker = running.tracker();
2251        let mut completed = Vec::new();
2252        // commit epochs while it runs, as the build loop does.
2253        let deadline = Instant::now() + Duration::from_millis(120);
2254        while Instant::now() < deadline {
2255            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
2256                for c in tracker.on_event(&e) {
2257                    running.commit(c.epoch);
2258                    completed.push(c.epoch);
2259                }
2260            }
2261        }
2262        let last = running.stop();
2263        running
2264            .finish(&erx, &mut tracker, |c| {
2265                completed.push(c.epoch);
2266                Ok(())
2267            })
2268            .unwrap();
2269        running.join();
2270        let commits = commits.lock().unwrap().clone();
2271        let barriers = barriers.lock().unwrap().clone();
2272        assert!(barriers.len() >= 3, "the ticker issued epochs: {barriers:?}");
2273        assert_eq!(barriers.last(), Some(&last), "stop issued the final barrier");
2274        assert_eq!(
2275            commits, completed,
2276            "the source hears exactly the completed epochs, in order"
2277        );
2278        assert_eq!(
2279            commits.last(),
2280            Some(&last),
2281            "the final epoch committed before the source exited"
2282        );
2283        let trace = trace.lock().unwrap().clone();
2284        let eos = trace.iter().position(|t| t == "e").expect("eos reached the sink");
2285        assert_eq!(
2286            trace[eos - 1],
2287            format!("b{last}"),
2288            "the final barrier arrives right before Eos: {trace:?}"
2289        );
2290        assert!(trace[..eos].iter().any(|t| t == "d"));
2291    }
2292
2293    #[test]
2294    fn an_operator_panic_ends_the_run_failed_and_every_task_finishes() {
2295        // the third batch panics inside the operator: the coordinator hears the panic as the
2296        // build's failure, stops the sources, and every task still finishes.
2297        struct PanicOp {
2298            seen: usize,
2299        }
2300        impl Operator for PanicOp {
2301            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {
2302                self.seen += 1;
2303                if self.seen == 3 {
2304                    panic!("interleave: types differ");
2305                }
2306            }
2307            fn on_watermark(&mut self, _wm: i64, _n: i64, _o: &mut Out) {}
2308            fn on_barrier(&mut self, _e: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
2309                Ok(OpSnapshot::whole(Vec::new()))
2310            }
2311            fn on_eos(&mut self, _o: &mut Out) {}
2312        }
2313        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2314        let src = g.source(Box::new(IdleSource {
2315            polls: 0,
2316            fail_at: None,
2317            commits: Arc::new(Mutex::new(Vec::new())),
2318            barriers: Arc::new(Mutex::new(Vec::new())),
2319        }));
2320        let op = g.operator(Box::new(PanicOp { seen: 0 }));
2321        let trace = Arc::new(Mutex::new(Vec::new()));
2322        let sink = g.sink(Box::new(TraceSink {
2323            trace: Arc::clone(&trace),
2324        }));
2325        g.route(src, Route::Forward(op));
2326        g.route(op, Route::Forward(sink));
2327        let (etx, erx) = sync_channel(4096);
2328        let running = g.start(etx).expect("start");
2329        let mut tracker = running.tracker();
2330        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
2331        running.join();
2332        assert!(err.contains("task 1 panicked") && err.contains("interleave"), "{err}");
2333        assert!(
2334            tracker.all_finished(),
2335            "the panic stopped the source and the sink finished"
2336        );
2337        let (_, state, _, _) = running.task_states()[op as usize];
2338        assert_eq!(state, TaskState::Done);
2339    }
2340
2341    #[test]
2342    fn an_operator_that_cannot_checkpoint_fails_the_run_with_its_diagnostic_and_no_panic() {
2343        // the operator's barrier returns an error (a full disk under its spill dir): the run fails
2344        // with that message — the end state a panic reaches, without one, and the diagnostic intact.
2345        struct NoDiskOp;
2346        impl Operator for NoDiskOp {
2347            fn on_data(&mut self, _e: EdgeId, _b: RecordBatch, _o: &mut Out) {}
2348            fn on_watermark(&mut self, _wm: i64, _n: i64, _o: &mut Out) {}
2349            fn on_barrier(&mut self, _e: u64, _o: &mut Out) -> Result<OpSnapshot, String> {
2350                Err("state io: write checkpoint file /spill/.ckpt/s1-hot.arrow: No space left on device".into())
2351            }
2352            fn on_eos(&mut self, _o: &mut Out) {}
2353        }
2354        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2355        let src = g.source(Box::new(IdleSource {
2356            polls: 0,
2357            fail_at: None,
2358            commits: Arc::new(Mutex::new(Vec::new())),
2359            barriers: Arc::new(Mutex::new(Vec::new())),
2360        }));
2361        let op = g.operator(Box::new(NoDiskOp));
2362        let trace = Arc::new(Mutex::new(Vec::new()));
2363        let sink = g.sink(Box::new(TraceSink {
2364            trace: Arc::clone(&trace),
2365        }));
2366        g.route(src, Route::Forward(op));
2367        g.route(op, Route::Forward(sink));
2368        let (etx, erx) = sync_channel(4096);
2369        let running = g.start(etx).expect("start");
2370        let mut tracker = running.tracker();
2371        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
2372        running.join();
2373        assert!(
2374            err.contains("checkpoint at epoch") && err.contains("No space left"),
2375            "{err}"
2376        );
2377        assert!(!err.contains("panicked"), "a returned error, not a caught panic: {err}");
2378        assert!(
2379            tracker.all_finished(),
2380            "the failure stopped the source and the sink finished"
2381        );
2382        let (_, state, _, _) = running.task_states()[op as usize];
2383        assert_eq!(state, TaskState::Done);
2384    }
2385
2386    #[test]
2387    fn a_deferred_operator_that_cannot_be_built_fails_start_before_any_task_runs() {
2388        // a restore the operator refuses is reported by `start` itself: no thread is spawned, so
2389        // nothing is left running with no coordinator to stop it.
2390        use std::sync::atomic::AtomicUsize;
2391        struct CountingSource(Arc<AtomicUsize>);
2392        impl Source for CountingSource {
2393            fn poll(&mut self, _out: &mut Out) -> Poll {
2394                self.0.fetch_add(1, Ordering::Relaxed);
2395                Poll::Idle
2396            }
2397            fn on_barrier(&mut self, _epoch: u64) -> Vec<u8> {
2398                Vec::new()
2399            }
2400        }
2401        let polls = Arc::new(AtomicUsize::new(0));
2402        let mut g = Graph::new(4);
2403        let src = g.source(Box::new(CountingSource(Arc::clone(&polls))));
2404        let op = g.operator_with(Box::new(|task, _in_edges| {
2405            Err::<Box<dyn Operator>, String>(format!(
2406                "stage 1 streamJoin task {task}: cannot restore its checkpoint — foreign manifest"
2407            ))
2408        }));
2409        let trace = Arc::new(Mutex::new(Vec::new()));
2410        let sink = g.sink(Box::new(TraceSink {
2411            trace: Arc::clone(&trace),
2412        }));
2413        g.route(src, Route::Forward(op));
2414        g.route(op, Route::Forward(sink));
2415        let (etx, _erx) = sync_channel(64);
2416        let err = g.start(etx).err().expect("start should have refused the operator");
2417        assert!(
2418            err.contains("cannot restore its checkpoint") && err.contains("foreign manifest"),
2419            "{err}"
2420        );
2421        std::thread::sleep(Duration::from_millis(50));
2422        assert_eq!(
2423            polls.load(Ordering::Relaxed),
2424            0,
2425            "no task ran: the source was never polled"
2426        );
2427        assert!(trace.lock().unwrap().is_empty(), "the sink saw nothing");
2428    }
2429
2430    #[test]
2431    fn a_failed_source_ends_the_run_with_its_error() {
2432        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2433        let src = g.source(Box::new(IdleSource {
2434            polls: 0,
2435            fail_at: Some(3),
2436            commits: Arc::new(Mutex::new(Vec::new())),
2437            barriers: Arc::new(Mutex::new(Vec::new())),
2438        }));
2439        let other = g.source(Box::new(IdleSource {
2440            polls: 0,
2441            fail_at: None,
2442            commits: Arc::new(Mutex::new(Vec::new())),
2443            barriers: Arc::new(Mutex::new(Vec::new())),
2444        }));
2445        let trace = Arc::new(Mutex::new(Vec::new()));
2446        let sink = g.sink(Box::new(TraceSink {
2447            trace: Arc::clone(&trace),
2448        }));
2449        g.route(src, Route::Forward(sink));
2450        g.route(other, Route::Forward(sink));
2451        let (etx, erx) = sync_channel(4096);
2452        let running = g.start(etx).expect("start");
2453        let mut tracker = running.tracker();
2454        let err = running.finish(&erx, &mut tracker, |_| Ok(())).unwrap_err();
2455        running.join();
2456        assert_eq!(err, "boom");
2457        assert!(tracker.all_finished(), "the failure stopped the other source too");
2458        // Eos means a clean end: the failed source's edge just closed, so the sink ended on a lost
2459        // input — it never saw `on_eos`, and so never treated uncommitted output as final.
2460        assert_ne!(trace.lock().unwrap().last().map(String::as_str), Some("e"));
2461    }
2462
2463    #[test]
2464    fn a_drained_sources_final_position_rides_every_later_epoch() {
2465        // A source that ran dry takes no more barriers, so the epochs completed after it would
2466        // name no position for it — and a restore would replay it from the start. Its final
2467        // position (`Drained`) stands in for every such epoch.
2468        let mut t = EpochTracker::new(3); // 0 source, 1 operator, 2 sink
2469        let pos = OpSnapshot::whole(b"src@done".to_vec());
2470        assert!(t
2471            .on_event(&Event::Drained {
2472                task: 0,
2473                snap: pos.clone()
2474            })
2475            .is_empty());
2476        assert!(t.on_event(&Event::Finished { task: 0, cpu_ms: 0 }).is_empty());
2477        for epoch in 1..=2 {
2478            assert!(t
2479                .on_event(&Event::Snapshot {
2480                    task: 1,
2481                    epoch,
2482                    snap: OpSnapshot::whole(vec![1]),
2483                })
2484                .is_empty());
2485            let done = t.on_event(&Event::Ack { task: 2, epoch });
2486            assert_eq!(done.len(), 1);
2487            let tasks: Vec<TaskId> = done[0].snapshots.iter().map(|(t, _)| *t).collect();
2488            assert_eq!(tasks, vec![0, 1], "epoch {epoch}: the drained source's position is in");
2489            assert_eq!(done[0].snapshots[0].1.head, pos.head);
2490        }
2491        // a source that did report an epoch keeps its own snapshot (no double entry)
2492        let mut t = EpochTracker::new(2);
2493        t.on_event(&Event::Snapshot {
2494            task: 0,
2495            epoch: 1,
2496            snap: OpSnapshot::whole(b"at@1".to_vec()),
2497        });
2498        t.on_event(&Event::Drained {
2499            task: 0,
2500            snap: pos.clone(),
2501        });
2502        t.on_event(&Event::Finished { task: 0, cpu_ms: 0 });
2503        let done = t.on_event(&Event::Ack { task: 1, epoch: 1 });
2504        assert_eq!(done[0].snapshots.len(), 1);
2505        assert_eq!(done[0].snapshots[0].1.head, b"at@1".to_vec());
2506    }
2507
2508    #[test]
2509    fn a_source_that_drains_before_the_first_barrier_still_has_its_position_checkpointed() {
2510        // Two sources into one sink: one drains at once (a bounded input that ran out before any
2511        // checkpoint — the cluster crash e2e's leader), the other keeps the run alive so barriers
2512        // keep coming. Every completed epoch must carry the drained source's final position.
2513        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2514        let src = g.source(Box::new(VecSource {
2515            batches: vec![batch(&["a"], &[1])].into(),
2516        }));
2517        let other = g.source(Box::new(IdleSource {
2518            polls: 0,
2519            fail_at: None,
2520            commits: Arc::new(Mutex::new(Vec::new())),
2521            barriers: Arc::new(Mutex::new(Vec::new())),
2522        }));
2523        let sink = g.sink(Box::new(TraceSink {
2524            trace: Arc::new(Mutex::new(Vec::new())),
2525        }));
2526        g.route(src, Route::Forward(sink));
2527        g.route(other, Route::Forward(sink));
2528        let (etx, erx) = sync_channel(4096);
2529        let running = g.start(etx).expect("start");
2530        let mut tracker = running.tracker();
2531        std::thread::sleep(Duration::from_millis(80));
2532        running.stop();
2533        let completed = Arc::new(Mutex::new(Vec::new()));
2534        let seen = Arc::clone(&completed);
2535        running
2536            .finish(&erx, &mut tracker, move |c| {
2537                seen.lock().unwrap().push(c);
2538                Ok(())
2539            })
2540            .unwrap();
2541        running.join();
2542        let completed = completed.lock().unwrap();
2543        assert!(!completed.is_empty(), "epochs completed after the source drained");
2544        for c in completed.iter() {
2545            let pos = c.snapshots.iter().find(|(t, _)| *t == src).map(|(_, s)| s.head.clone());
2546            assert_eq!(
2547                pos,
2548                Some(b"src@done".to_vec()),
2549                "epoch {}: the drained source's final position",
2550                c.epoch
2551            );
2552        }
2553    }
2554
2555    #[test]
2556    fn epoch_tracker_completes_in_order_and_counts_finished_tasks_as_reported() {
2557        let mut t = EpochTracker::new(3); // 0 source, 1 operator, 2 sink
2558        let snap = |task, epoch| Event::Snapshot {
2559            task,
2560            epoch,
2561            snap: OpSnapshot::whole(vec![task as u8]),
2562        };
2563        assert!(t.on_event(&snap(0, 1)).is_empty());
2564        assert!(t.on_event(&snap(1, 1)).is_empty());
2565        // epoch 2 reported by the source before epoch 1's ack: nothing completes yet.
2566        assert!(t.on_event(&snap(0, 2)).is_empty());
2567        let done = t.on_event(&Event::Ack { task: 2, epoch: 1 });
2568        assert_eq!(done.len(), 1);
2569        assert_eq!(done[0].epoch, 1);
2570        assert_eq!(
2571            done[0].snapshots,
2572            vec![(0, OpSnapshot::whole(vec![0])), (1, OpSnapshot::whole(vec![1]))]
2573        );
2574        // the operator finished (Eos) before epoch 2's barrier reached it: it counts as reported.
2575        assert!(t.on_event(&Event::Finished { task: 1, cpu_ms: 0 }).is_empty());
2576        let done = t.on_event(&Event::Ack { task: 2, epoch: 2 });
2577        assert_eq!(done.iter().map(|c| c.epoch).collect::<Vec<_>>(), vec![2]);
2578        assert!(!t.all_finished());
2579        t.on_event(&Event::Finished { task: 0, cpu_ms: 0 });
2580        t.on_event(&Event::Finished { task: 2, cpu_ms: 0 });
2581        assert!(t.all_finished());
2582    }
2583
2584    #[test]
2585    fn a_chained_sink_runs_in_the_source_task_and_flushes_before_the_position_is_taken() {
2586        // one task, no edge: data reaches the sink in order; at a barrier the sink's flush comes
2587        // before the source records its position; a refused flush leaves no position; Eos reaches
2588        // the sink when the task stops.
2589        struct TracingSource {
2590            trace: Arc<Mutex<Vec<String>>>,
2591            polls: usize,
2592        }
2593        impl Source for TracingSource {
2594            fn poll(&mut self, out: &mut Out) -> Poll {
2595                self.polls += 1;
2596                out.push(batch(&["c"], &[self.polls as i64]));
2597                Poll::Idle
2598            }
2599            fn on_barrier(&mut self, epoch: u64) -> Vec<u8> {
2600                self.trace.lock().unwrap().push(format!("position {epoch}"));
2601                epoch.to_le_bytes().to_vec()
2602            }
2603            fn on_commit(&mut self, epoch: u64) {
2604                self.trace.lock().unwrap().push(format!("commit {epoch}"));
2605            }
2606        }
2607        struct TracingSink {
2608            trace: Arc<Mutex<Vec<String>>>,
2609            refuse: u64,
2610        }
2611        impl Sink for TracingSink {
2612            fn on_data(&mut self, b: RecordBatch) {
2613                self.trace.lock().unwrap().push(format!("data {}", b.num_rows()));
2614            }
2615            fn on_barrier(&mut self, epoch: u64) -> bool {
2616                self.trace.lock().unwrap().push(format!("flush {epoch}"));
2617                epoch != self.refuse
2618            }
2619            fn on_eos(&mut self) {
2620                self.trace.lock().unwrap().push("eos".into());
2621            }
2622        }
2623        let trace = Arc::new(Mutex::new(Vec::new()));
2624        let mut g = Graph::new(4).barrier_every(Duration::from_millis(10));
2625        let task = g.source(Box::new(Chained {
2626            source: TracingSource {
2627                trace: Arc::clone(&trace),
2628                polls: 0,
2629            },
2630            sink: TracingSink {
2631                trace: Arc::clone(&trace),
2632                refuse: 2,
2633            },
2634        }));
2635        let (etx, erx) = sync_channel(4096);
2636        let running = g.start(etx).expect("start");
2637        let mut tracker = running.tracker();
2638        let mut completed = Vec::new();
2639        let until = Instant::now() + Duration::from_millis(80);
2640        while Instant::now() < until {
2641            if let Ok(e) = erx.recv_timeout(Duration::from_millis(10)) {
2642                for c in tracker.on_event(&e) {
2643                    running.commit(c.epoch);
2644                    completed.push(c);
2645                }
2646            }
2647        }
2648        let last = running.stop();
2649        running
2650            .finish(&erx, &mut tracker, |c| {
2651                completed.push(c);
2652                Ok(())
2653            })
2654            .unwrap();
2655        running.join();
2656        assert_eq!(
2657            running.n_tasks(),
2658            1,
2659            "no sink task: the sink lives in the source's task"
2660        );
2661        let trace = trace.lock().unwrap().clone();
2662        let at = |s: &str| trace.iter().position(|t| t == s);
2663        assert!(
2664            at("data 1").unwrap() < at("flush 1").unwrap(),
2665            "data before the first barrier"
2666        );
2667        assert!(
2668            at("flush 1").unwrap() < at("position 1").unwrap(),
2669            "durable before the position: {trace:?}"
2670        );
2671        assert!(
2672            at("position 2").is_none(),
2673            "a refused flush records no position for that epoch"
2674        );
2675        assert!(at("flush 3").unwrap() < at("position 3").unwrap());
2676        // the last epoch commits after its flush and position, and the sink sees Eos when the
2677        // task stops. Which of the two lands last depends on whether the ticker or the stop
2678        // issued that epoch (its commit can reach the task before the stop does); both orders
2679        // keep offsets behind durable output.
2680        let final_commit = at(&format!("commit {last}")).expect("the final epoch commits");
2681        assert!(
2682            at(&format!("flush {last}")).unwrap() < final_commit,
2683            "the final commit follows its flush: {trace:?}"
2684        );
2685        let eos = at("eos").expect("the sink saw Eos when the task stopped");
2686        assert!(
2687            trace.len() - 1 == final_commit || trace.len() - 1 == eos,
2688            "the final commit or Eos lands last: {trace:?}"
2689        );
2690        assert!(
2691            completed.iter().any(|c| c.epoch == 2),
2692            "epoch 2 still completes (one task, one snapshot)"
2693        );
2694        let _ = task;
2695    }
2696
2697    #[test]
2698    fn thread_cpu_ms_reads_this_threads_own_time() {
2699        // Linux only: the ledger every task reports on Finished. Burn a little CPU and see it.
2700        let Some(before) = thread_cpu_ms() else {
2701            return; // not Linux: the ledger is absent, never wrong
2702        };
2703        // burn until the ledger moved (a loaded CI runner may give this thread little CPU per
2704        // wall second; the ledger's resolution is 10 ms).
2705        let mut x: u64 = 0;
2706        let start = Instant::now();
2707        let mut after = before;
2708        while after < before + 20 && start.elapsed() < Duration::from_secs(5) {
2709            for _ in 0..100_000 {
2710                x = x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
2711            }
2712            after = thread_cpu_ms().unwrap();
2713        }
2714        assert!(x != 1, "keep the loop");
2715        assert!(after >= before + 20, "cpu time advanced: {before} → {after}");
2716        // another thread's time is its own: a fresh thread starts near zero.
2717        let fresh = std::thread::spawn(|| thread_cpu_ms().unwrap()).join().unwrap();
2718        assert!(
2719            fresh < after,
2720            "a new thread's ledger is not this one's: {fresh} vs {after}"
2721        );
2722    }
2723
2724    #[test]
2725    fn shuffle_splits_rows_by_vnode_range_and_keeps_all_rows() {
2726        let b = batch(&["a", "b", "c", "d", "e", "f"], &[1, 2, 3, 4, 5, 6]);
2727        let ranges = vnode_ranges(8, 3);
2728        let targets: Vec<(TaskId, std::ops::Range<u32>)> =
2729            ranges.into_iter().enumerate().map(|(i, r)| (i as TaskId, r)).collect();
2730        let parts = shuffle(&b, &["k".to_string()], 8, &targets, test_hasher().as_ref());
2731        let total: usize = parts.iter().flatten().map(|p| p.num_rows()).sum();
2732        assert_eq!(total, 6);
2733        // the same key always lands on the same target.
2734        let again = shuffle(&b, &["k".to_string()], 8, &targets, test_hasher().as_ref());
2735        for (p, q) in parts.iter().zip(again.iter()) {
2736            assert_eq!(p.as_ref().map(|x| x.num_rows()), q.as_ref().map(|x| x.num_rows()));
2737        }
2738        // no routing columns: everything to the first target.
2739        let global = shuffle(&b, &[], 8, &targets, test_hasher().as_ref());
2740        assert_eq!(global[0].as_ref().map(|x| x.num_rows()), Some(6));
2741        assert!(global[1].is_none() && global[2].is_none());
2742    }
2743
2744    #[test]
2745    fn vnode_ranges_cover_every_vnode_exactly_once() {
2746        let r = vnode_ranges(64, 5);
2747        assert_eq!(r.len(), 5);
2748        let covered: Vec<u32> = r.iter().flat_map(|x| x.clone()).collect();
2749        assert_eq!(covered, (0..64).collect::<Vec<_>>());
2750        assert_eq!(vnode_ranges(4, 9).len(), 4, "never more tasks than vnodes");
2751    }
2752}