Skip to main content

spate_core/coordination/
driver.rs

1//! Reusable source-side choreography for coordinated sources.
2//!
3//! A coordinated source owns two jobs: reading its data (lanes, fetchers,
4//! offsets, all connector-specific) and translating [`CoordinationEvent`]s
5//! into the controller's assignment protocol while keeping fenced-tenancy
6//! bookkeeping straight (source-generic). [`CoordinationDriver`] owns the
7//! second job. A source embeds one next to a [`SplitSource`]
8//! implementation and delegates `poll_events`/`commit` to it.
9//!
10//! # Tenancies, lanes, partitions
11//!
12//! Every continuous ownership span of a split (from `Gained` to whatever
13//! ends it) is one **tenancy**, and each tenancy gets a fresh, never
14//! reused [`PartitionId`] and a fresh, never reused [`LaneId`]. A lane
15//! materializes exactly once, when its tenancy's split is staged for
16//! opening, and lives untouched until the tenancy ends. Gains are
17//! additive ([`SourceEvent::LanesAdded`]); a peer's split arriving never
18//! drains flowing lanes. Watermarks come back keyed by partition, so a
19//! late drain-commit from a lane that lost its split resolves to a retired
20//! tenancy and is skipped. A stale write is never folded, committed, or
21//! resurrected.
22//!
23//! # Event choreography
24//!
25//! One controller event per [`poll_events`](CoordinationDriver::poll_events)
26//! call, in priority order:
27//!
28//! 1. Pending losses → partial [`SourceEvent::LanesRevoked`] (barrier
29//!    sized one party per lane, matching the runtime's drain contract).
30//!    Once delivered, the retired tenancies they belonged to have
31//!    absorbed every late watermark they can see and are pruned.
32//! 2. Staged gains → [`SourceEvent::LanesAdded`] with lanes for the
33//!    newly-gained splits only; existing lanes and their in-flight acks
34//!    are untouched.
35//! 3. Otherwise poll the coordinator (the idle wait delegates there),
36//!    fold its events into the tenancy table, sweep for completions, and
37//!    advance any in-flight cooperative revocations (a revoked split, once
38//!    its intake is stopped and its tail acked, takes a final fenced
39//!    commit and is handed back barrier-less).
40//! 4. [`CoordinationEvent::AllComplete`] → [`SourceEvent::Drained`];
41//!    [`CoordinationEvent::Stalled`] → a fatal error by default
42//!    (see [`stall_drains`](CoordinationDriver::stall_drains)).
43//! 5. Nothing staged and some lane newly at end-of-input
44//!    ([`SplitSource::take_finishing`]) → [`SourceEvent::CommitReady`],
45//!    so the runtime chases the final acks instead of waiting out its
46//!    commit tick.
47
48use super::{
49    ControlWaker, CoordinationError, CoordinationErrorKind, CoordinationEvent, LeaseEpoch,
50    SplitCoordinator, SplitId, SplitPlanner, SplitProgress, SplitSpec,
51};
52use crate::error::{ErrorClass, SourceError};
53use crate::record::PartitionId;
54use crate::source::{DrainBarrier, LaneId, SourceEvent, SourceLane};
55use std::collections::BTreeMap;
56use std::fmt;
57use std::time::Duration;
58
59/// Everything the driver hands a source when a split's lane is
60/// materialized. This happens exactly once per tenancy, when the gain is
61/// staged into a [`SourceEvent::LanesAdded`].
62#[derive(Debug)]
63#[non_exhaustive]
64pub struct SplitOpening<'a> {
65    /// The split to read.
66    pub split: &'a SplitSpec,
67    /// Authoritative progress to resume from (already validated via
68    /// [`SplitSource::validate_resume`]); `None` for a fresh split.
69    pub resume: Option<&'a SplitProgress>,
70    /// Lane id minted for this tenancy's lifetime; never reused by this
71    /// source.
72    pub lane: LaneId,
73    /// Stable partition id for this tenancy, the key under which this
74    /// split's watermarks come back to [`CoordinationDriver::commit`].
75    pub partition: PartitionId,
76    /// Fencing token of the current tenancy.
77    pub epoch: LeaseEpoch,
78    /// Wakes the control-plane wait. Clone it into the lane and signal it
79    /// the moment the lane decides end-of-input or reports poison.
80    /// Otherwise the driver notices only between waits, and the split's
81    /// completion waits out an idle timeout.
82    pub waker: &'a ControlWaker,
83}
84
85/// What the driver needs from the embedding source.
86///
87/// Implement it on the source's lane-assembly context (the sub-struct that
88/// holds what lane construction needs), not on the source itself. The
89/// driver lives beside that context as a sibling field, so both can be
90/// borrowed disjointly.
91pub trait SplitSource {
92    /// The data-plane lane type produced for gained splits.
93    type Lane: SourceLane;
94
95    /// Materialize the lane for a gained (or re-assigned) split. Spawn
96    /// fetchers here; never block on data.
97    fn open_split(&mut self, opening: SplitOpening<'_>) -> Result<Self::Lane, SourceError>;
98
99    /// Drift-check carried progress against this instance's view of the
100    /// split (etag pins, schema versions) before it is trusted; the default
101    /// accepts everything.
102    ///
103    /// A rejection is raised before the tenancy is recorded, so the split is
104    /// never opened, and the driver reports it as poison: one delivery
105    /// attempt is consumed, the split is handed back for another instance,
106    /// and at the attempt cap it is quarantined. The error's class then
107    /// decides this run. [`ErrorClass::Fatal`] stops the pipeline; any other
108    /// class is logged and the run continues with the split left to the
109    /// coordinator.
110    fn validate_resume(
111        &self,
112        split: &SplitSpec,
113        progress: &SplitProgress,
114    ) -> Result<(), SourceError> {
115        let _ = (split, progress);
116        Ok(())
117    }
118
119    /// Snapshot the split's committable progress at an acked watermark. The
120    /// snapshot carries the opaque resume state plus whether that watermark
121    /// completes the split (fully delivered **and** fully acknowledged; the
122    /// source owns its eof/emitted accounting).
123    fn encode_commit(
124        &mut self,
125        split: &SplitId,
126        watermark: i64,
127    ) -> Result<SplitProgress, SourceError>;
128
129    /// Completion sweep for an owned split with no new watermark this
130    /// tick (empty splits; tails acked exactly at the previous commit).
131    /// Returns `Some(terminal progress)` when complete, `None` while data
132    /// is in flight.
133    fn sweep(&mut self, split: &SplitId) -> Result<Option<SplitProgress>, SourceError>;
134
135    /// The split's lane is being retired (lost, fenced, completed, or
136    /// shutdown). Detach its fetcher; never abort it, because the pipeline
137    /// thread may still be draining the lane. Must not block.
138    ///
139    /// This is the end of the tenancy. The driver never calls
140    /// [`SplitSource::encode_commit`] or [`SplitSource::sweep`] for the
141    /// split afterwards (its tenancy is retired first, and retired
142    /// tenancies absorb late watermarks), so the source may drop the
143    /// split's state here.
144    fn close_split(&mut self, split: &SplitId);
145
146    /// Splits whose lanes decided end-of-input since the last call (the
147    /// edge, not the level). The driver surfaces them as
148    /// [`SourceEvent::CommitReady`] so the runtime chases their final acks
149    /// instead of waiting out the commit tick. The split then completes
150    /// (and frees its working-set slot) within milliseconds of its last
151    /// record becoming sink-durable. A latency hint only; the default
152    /// reports none.
153    fn take_finishing(&mut self) -> Vec<SplitId> {
154        Vec::new()
155    }
156
157    /// Begin a cooperative revocation of an owned split. Stop its intake at a
158    /// safe boundary while **keeping** its commit state, so the tail can
159    /// still be chased to a final fenced commit. Unlike
160    /// [`close_split`](SplitSource::close_split) (which ends the tenancy and
161    /// lets the source drop the split's state), the split stays commit- and
162    /// sweep-adjacent here. The driver keeps committing its acked
163    /// watermarks and then calls [`drain_ready`](SplitSource::drain_ready)
164    /// until the drain finishes.
165    ///
166    /// Return `true` to accept the revocation, `false` to decline it (the
167    /// default). **Contract: return `false` for any split this source has
168    /// not opened or has already closed or completed.** The driver also
169    /// guards this (it declines tenancies without an open lane and feeds
170    /// the decline back to the backend), but the source must not rely on
171    /// that alone.
172    ///
173    /// Declining is safe but not free. The split still leaves, because a
174    /// revocation is the leader's decision; the backend forces the release
175    /// instead and this split's uncommitted tail replays under its next
176    /// owner. (The backend cancels a revocation the leader takes back
177    /// before the decline lands, and then the split stays and keeps being
178    /// read.) A source that *can* stop intake at a safe boundary should,
179    /// because that is the difference between a replay-free move and a
180    /// bounded-duplicate one.
181    ///
182    /// While a split is draining, [`encode_commit`](SplitSource::encode_commit)
183    /// must never report it `completed`. A drain cut can look terminal to
184    /// the source (everything emitted is acked) while the split is
185    /// half-read; the driver strips a `completed` flag it sees here and
186    /// logs an error.
187    fn begin_revoke(&mut self, split: &SplitId) -> bool {
188        let _ = split;
189        false
190    }
191
192    /// Poll a draining split for its final progress. Returns
193    /// `Some(progress)` with `completed: false` once every record it emitted
194    /// is acked **and** that watermark is committed, so the resume point
195    /// handed to the next owner covers everything this instance produced (a
196    /// replay-free transfer); `None` while any of that tail is still in
197    /// flight, to be retried on the next poll. Never reports `completed`; a
198    /// revocation gives the split away rather than finishing it. The default
199    /// reports `None`.
200    ///
201    /// **Level-triggered, unlike [`take_finishing`](SplitSource::take_finishing):**
202    /// keep returning `Some` on every poll until the driver retires the
203    /// split. The final store commit can defer on a store hiccup and is
204    /// re-attempted from a fresh `drain_ready` answer, so an
205    /// edge-triggered implementation would stall the drain until the
206    /// backend forced it.
207    ///
208    /// A source that accepts a revocation in
209    /// [`begin_revoke`](SplitSource::begin_revoke) **must** eventually
210    /// answer `Some` here. The default `None` never finishes a drain; it
211    /// pairs with the default `begin_revoke`, which declines.
212    fn drain_ready(&mut self, split: &SplitId) -> Result<Option<SplitProgress>, SourceError> {
213        let _ = split;
214        Ok(None)
215    }
216}
217
218#[derive(Debug, PartialEq, Eq, Clone, Copy)]
219enum TenancyState {
220    /// Owned; lane live or staged to open.
221    Live,
222    /// Owned, but intake has stopped at a safe boundary for a cooperative
223    /// revocation, and the lane is draining toward one final fenced commit.
224    /// Still commit-eligible (tick commits keep folding its acked
225    /// watermarks) and never swept, because a revocation gives the split
226    /// away rather than completing it. Becomes `Retired` (drained, so
227    /// barrier-less) once that final commit lands, or fenced-`Retired` (the
228    /// loss path) if a peer fences it mid-drain.
229    Draining,
230    /// Ownership over (lost, fenced, failed, completed, released); entry
231    /// retained only to absorb late watermarks until its revocation is
232    /// delivered.
233    Retired,
234}
235
236#[derive(Debug)]
237struct Tenancy {
238    split: SplitSpec,
239    epoch: LeaseEpoch,
240    lane: Option<LaneId>,
241    state: TenancyState,
242    /// A commit for this tenancy was fenced: fold nothing, commit nothing.
243    fenced: bool,
244    /// Resume cache: the `Gained` carry, then every acked commit fold.
245    /// Acked means sink-durable, so respawning from it can only skip data
246    /// that is already safe. At-least-once holds even when the durable
247    /// store lags a Retryable commit behind.
248    progress: Option<SplitProgress>,
249    /// Terminal progress reached the store; nothing further to commit.
250    completed: bool,
251    /// This tenancy released its split through a cooperative revocation. Its
252    /// final commit is durable, so the peer resumes replay-free. Routes the
253    /// lane out through the barrier-less retired path, exactly like
254    /// `completed` (nothing is in flight behind a drained revocation).
255    handed_off: bool,
256}
257
258/// Source-side coordination choreography, embedded by a coordinated
259/// source; see the [module docs](self) for the protocol it implements.
260pub struct CoordinationDriver {
261    coordinator: Box<dyn SplitCoordinator>,
262    /// Parking half of the control-plane wakeup; the waker half is held by
263    /// the backend and by every lane this driver opened.
264    wait: crossbeam_channel::Receiver<()>,
265    waker: ControlWaker,
266    tenancies: BTreeMap<PartitionId, Tenancy>,
267    by_split: BTreeMap<SplitId, PartitionId>,
268    /// Lanes whose loss must still surface as a partial revoke.
269    pending_lost: Vec<LaneId>,
270    /// Lanes of completed tenancies: terminal progress reached the store,
271    /// so they leave without a drain barrier.
272    pending_retired: Vec<LaneId>,
273    /// Tenancies gained but not yet materialized: their lanes go out in
274    /// the next [`SourceEvent::LanesAdded`].
275    pending_open: Vec<PartitionId>,
276    /// Poison reports the backend refused, with the reason to re-offer.
277    /// A gain refused on resume has no tenancy, so until its report lands
278    /// nothing else here hands the split back and the backend keeps
279    /// renewing its lease.
280    pending_poison: Vec<(SplitId, String)>,
281    all_complete: bool,
282    stalled: Option<(u64, u64)>,
283    stall_drains: bool,
284    started: bool,
285    next_partition: u32,
286    /// Lane ids are minted once per tenancy and never reused.
287    next_lane: u32,
288}
289
290impl fmt::Debug for CoordinationDriver {
291    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
292        f.debug_struct("CoordinationDriver")
293            .field("tenancies", &self.tenancies.len())
294            .field("live", &self.by_split.len())
295            .field("pending_lost", &self.pending_lost.len())
296            .field("pending_retired", &self.pending_retired.len())
297            .field("pending_open", &self.pending_open.len())
298            .field("all_complete", &self.all_complete)
299            .field("stalled", &self.stalled)
300            .field("started", &self.started)
301            .finish_non_exhaustive()
302    }
303}
304
305impl CoordinationDriver {
306    /// Wrap a coordinator handle.
307    #[must_use]
308    pub fn new(mut coordinator: Box<dyn SplitCoordinator>) -> CoordinationDriver {
309        let (waker, wait) = super::control_channel();
310        coordinator.set_waker(waker.clone());
311        CoordinationDriver {
312            coordinator,
313            wait,
314            waker,
315            tenancies: BTreeMap::new(),
316            by_split: BTreeMap::new(),
317            pending_lost: Vec::new(),
318            pending_retired: Vec::new(),
319            pending_open: Vec::new(),
320            pending_poison: Vec::new(),
321            all_complete: false,
322            stalled: None,
323            stall_drains: false,
324            started: false,
325            next_partition: 0,
326            next_lane: 0,
327        }
328    }
329
330    /// Treat [`CoordinationEvent::Stalled`] as a drain-with-warning
331    /// instead of a fatal error. Default `false`: a bounded job that
332    /// cannot finish because splits are quarantined fails loudly rather
333    /// than exiting as if it were complete.
334    #[must_use]
335    pub fn stall_drains(mut self, drains: bool) -> CoordinationDriver {
336        self.stall_drains = drains;
337        self
338    }
339
340    /// Join the job. Returns the empty `LanesAssigned` ready signal, which
341    /// the source must return from the *same* `poll_events` call. It bumps
342    /// the controller's assignment epoch and marks the pipeline running
343    /// while splits are still being claimed.
344    pub fn start<L>(
345        &mut self,
346        planner: Box<dyn SplitPlanner>,
347    ) -> Result<SourceEvent<L>, SourceError> {
348        assert!(!self.started, "CoordinationDriver::start called twice");
349        self.coordinator.start(planner).map_err(as_source_error)?;
350        self.started = true;
351        Ok(SourceEvent::LanesAssigned(Vec::new()))
352    }
353
354    /// Coordinated `poll_events` body: surfaces at most one controller
355    /// event per call, per the [module docs](self).
356    pub fn poll_events<S: SplitSource>(
357        &mut self,
358        source: &mut S,
359        timeout: Duration,
360    ) -> Result<SourceEvent<S::Lane>, SourceError> {
361        assert!(self.started, "poll_events before start");
362
363        // Staged work can be consumed without producing an event (a batch of
364        // gains, every one of which was retired before it could open), and
365        // then the drain has to look again. Loop rather than recurse: nothing
366        // bounds how many such batches a backend produces back to back.
367        let mut park = timeout;
368        loop {
369            // 1. Losses first: stop lost lanes before anything else runs.
370            if !self.pending_lost.is_empty() {
371                let lanes = std::mem::take(&mut self.pending_lost);
372                let barrier = DrainBarrier::new(lanes.len());
373                return Ok(SourceEvent::LanesRevoked { lanes, barrier });
374            }
375
376            // 1b. Completed tenancies leave barrier-less.
377            if !self.pending_retired.is_empty() {
378                let lanes = std::mem::take(&mut self.pending_retired);
379                return Ok(SourceEvent::LanesRetired { lanes });
380            }
381
382            // Every queued revocation has been delivered and the controller
383            // has drained and committed those lanes, so retired tenancies
384            // have absorbed every late watermark they can see. Prune only
385            // those; `Draining` tenancies must survive to be advanced.
386            self.tenancies
387                .retain(|_, t| t.state != TenancyState::Retired);
388
389            // 2. Staged gains: additive lanes for the newly-gained splits
390            // only; existing lanes are untouched.
391            if !self.pending_open.is_empty() {
392                let lanes = self.open_pending(source)?;
393                if !lanes.is_empty() {
394                    return Ok(SourceEvent::LanesAdded(lanes));
395                }
396            }
397
398            // 3. Terminal states, once the choreography above has quiesced.
399            if let Some((completed, quarantined)) = self.stalled {
400                if self.stall_drains {
401                    tracing::warn!(
402                        completed,
403                        quarantined,
404                        "job stalled; draining as configured"
405                    );
406                    return Ok(SourceEvent::Drained);
407                }
408                return Err(SourceError::Client {
409                    class: ErrorClass::Fatal,
410                    reason: format!(
411                        "coordinated job stalled: {completed} splits completed but {quarantined} \
412                     are quarantined and out of delivery attempts; inspect \
413                     spate_coordination_splits_quarantined and requeue or exclude them"
414                    ),
415                });
416            }
417            if self.all_complete {
418                return Ok(SourceEvent::Drained);
419            }
420
421            // 3b. Re-offer poison reports the backend refused; a refused
422            // report leaves a split held here with no tenancy behind it.
423            for (split, reason) in std::mem::take(&mut self.pending_poison) {
424                if !self.report_poison(&split, &reason) {
425                    self.pending_poison.push((split, reason));
426                }
427            }
428
429            // 4. Drain the coordinator (never blocks; the wait is ours, at
430            // the end of this function).
431            let events = self.coordinator.poll().map_err(as_source_error)?;
432            // Apply every event even after one fails, and surface one
433            // failure afterwards. `poll` drained the batch, so an event
434            // skipped here is never re-offered: a skipped gain leaves a split
435            // this instance holds but never reads, and a skipped loss leaves
436            // a lane reading a split it no longer owns.
437            let mut surfaced: Option<SourceError> = None;
438            for event in events {
439                if let Err(e) = self.apply(source, event) {
440                    // First failure, except that a fatal one anywhere in the
441                    // batch takes its place: the controller reads only the
442                    // class of the error it is handed, so a fatal error
443                    // behind an earlier retryable one would let the run
444                    // continue past a stop the source asked for.
445                    match &surfaced {
446                        Some(kept) if is_fatal(kept) || !is_fatal(&e) => {}
447                        _ => surfaced = Some(e),
448                    }
449                }
450            }
451            if let Some(e) = surfaced {
452                return Err(e);
453            }
454
455            // 5. Completion sweep over live, uncommitted-terminal tenancies.
456            self.sweep(source)?;
457
458            // 5b. Advance in-flight cooperative revocations.
459            self.advance_drains(source)?;
460
461            if !self.pending_lost.is_empty()
462                || !self.pending_retired.is_empty()
463                || !self.pending_open.is_empty()
464                || self.all_complete
465                || self.stalled.is_some()
466            {
467                // Something is staged: go round and surface it on this same
468                // call, without parking on the way.
469                park = Duration::ZERO;
470                continue;
471            }
472            break;
473        }
474
475        // 5. Nothing staged: surface newly-finishing splits so the runtime
476        // chases their final acks instead of waiting out its commit tick.
477        let finishing = source.take_finishing();
478        if !finishing.is_empty() {
479            let partitions: Vec<PartitionId> = finishing
480                .iter()
481                .filter_map(|split| self.by_split.get(split).copied())
482                .collect();
483            if !partitions.is_empty() {
484                return Ok(SourceEvent::CommitReady { partitions });
485            }
486        }
487
488        // 6. Nothing to report: park here, not inside the backend. Both the
489        // backend and any lane deciding end-of-input or reporting poison
490        // signal this waker, ending the park before the timeout runs out.
491        if !park.is_zero() {
492            let _ = self.wait.recv_timeout(park);
493        }
494        Ok(SourceEvent::Idle)
495    }
496
497    /// Coordinated `commit` body: per-split fenced commits keyed by the
498    /// tenancy partition ids the driver minted.
499    pub fn commit<S: SplitSource>(
500        &mut self,
501        source: &mut S,
502        watermarks: &[(PartitionId, i64)],
503    ) -> Result<(), SourceError> {
504        for &(partition, watermark) in watermarks {
505            let Some(tenancy) = self.tenancies.get(&partition) else {
506                // Pruned tenancy: a drain commit that arrived after its
507                // retirement was fully delivered. Its data replays under
508                // the new owner.
509                continue;
510            };
511            if tenancy.state == TenancyState::Retired || tenancy.fenced || tenancy.completed {
512                continue;
513            }
514            let split = tenancy.split.id.clone();
515            let progress = source.encode_commit(&split, watermark)?;
516            self.commit_progress(source, partition, &split, progress)?;
517        }
518        Ok(())
519    }
520
521    /// Report an owned split as poison: consumes a delivery attempt and
522    /// hands it back for another worker (or quarantine, at the cap). The
523    /// split's lane is retired through the normal loss path.
524    pub fn fail<S: SplitSource>(
525        &mut self,
526        source: &mut S,
527        split: &SplitId,
528        reason: &str,
529    ) -> Result<(), SourceError> {
530        let Some(&partition) = self.by_split.get(split) else {
531            return Ok(()); // already lost, nothing to report
532        };
533        match self.coordinator.fail(split, reason) {
534            Ok(()) => {}
535            // Fenced: someone already took it; the retire below still applies.
536            Err(e) if e.kind == CoordinationErrorKind::Fenced => {}
537            Err(e) => return Err(as_source_error(e)),
538        }
539        self.retire(source, partition, false);
540        Ok(())
541    }
542
543    /// Best-effort graceful release of every held split, so peers claim
544    /// them without waiting out the lease. Call from the source's `Drop`.
545    pub fn release(&mut self) {
546        if !self.started {
547            return;
548        }
549        let held: Vec<SplitId> = self.by_split.keys().cloned().collect();
550        if held.is_empty() {
551            return;
552        }
553        if let Err(e) = self.coordinator.release(&held) {
554            tracing::warn!(error = %e, "graceful split release failed; leases will expire");
555        }
556    }
557
558    /// Current live split → lane view (pause/resume bookkeeping, tests).
559    #[must_use]
560    pub fn assignments(&self) -> Vec<(SplitId, LaneId)> {
561        self.by_split
562            .iter()
563            .filter_map(|(split, partition)| {
564                let lane = self.tenancies.get(partition)?.lane?;
565                Some((split.clone(), lane))
566            })
567            .collect()
568    }
569
570    /// Report a gain the source refused as poison, for a split with no
571    /// tenancy to retire. The rejection is what the caller returns, so a
572    /// report the backend refuses is queued for retry rather than put in
573    /// its place.
574    fn report_rejected_gain(&mut self, split: &SplitId, rejection: &SourceError) {
575        let reason = format!("carried progress rejected on resume: {rejection}");
576        if !self.report_poison(split, &reason) {
577            self.pending_poison.push((split.clone(), reason));
578        }
579    }
580
581    /// Offer one poison report to the backend. `false` means the backend
582    /// refused it and this instance is still holding the split.
583    fn report_poison(&mut self, split: &SplitId, reason: &str) -> bool {
584        match self.coordinator.fail(split, reason) {
585            Ok(()) => true,
586            // Fenced: someone already took it, so it is already back.
587            Err(e) if e.kind == CoordinationErrorKind::Fenced => true,
588            Err(e) => {
589                tracing::warn!(
590                    split = %split,
591                    error = %e,
592                    "poison report refused; retrying while this instance holds the split"
593                );
594                false
595            }
596        }
597    }
598
599    fn apply<S: SplitSource>(
600        &mut self,
601        source: &mut S,
602        event: CoordinationEvent,
603    ) -> Result<(), SourceError> {
604        match event {
605            CoordinationEvent::Gained {
606                split,
607                epoch,
608                progress,
609            } => {
610                if let Some(&stale) = self.by_split.get(&split.id) {
611                    // Backend contract violation (re-gain without a loss);
612                    // retire the stale tenancy defensively and continue.
613                    tracing::warn!(split = %split.id, "gained a split already held; retiring stale tenancy");
614                    self.retire(source, stale, false);
615                }
616                if let Some(progress) = progress.as_ref()
617                    && let Err(e) = source.validate_resume(&split, progress)
618                {
619                    // Report before the rejection leaves: no tenancy was
620                    // recorded, so nothing else here releases the split and
621                    // the backend keeps renewing its lease.
622                    self.report_rejected_gain(&split.id, &e);
623                    return Err(e);
624                }
625                let partition = PartitionId(self.next_partition);
626                self.next_partition += 1;
627                self.by_split.insert(split.id.clone(), partition);
628                self.tenancies.insert(
629                    partition,
630                    Tenancy {
631                        split,
632                        epoch,
633                        lane: None,
634                        state: TenancyState::Live,
635                        fenced: false,
636                        progress,
637                        completed: false,
638                        handed_off: false,
639                    },
640                );
641                self.pending_open.push(partition);
642            }
643            CoordinationEvent::RevokeRequested { split } => {
644                // The leader wants this split back. Accept only a split held
645                // live with an OPEN lane, un-fenced and not yet completed,
646                // and only if the source can stop its intake at a safe
647                // boundary: a tenancy gained but not yet opened has no drain
648                // to finish and would sit in `Draining` forever. A refusal is
649                // declined back to the backend, which forces the release; the
650                // split leaves either way.
651                //
652                // A repeat request for a tenancy already draining is accepted
653                // silently, with no second `begin_revoke` and no decline.
654                if let Some(&partition) = self.by_split.get(&split)
655                    && self
656                        .tenancies
657                        .get(&partition)
658                        .is_some_and(|t| t.state == TenancyState::Draining && !t.fenced)
659                {
660                    return Ok(());
661                }
662                let accepted = match self.by_split.get(&split) {
663                    Some(&partition) => {
664                        let eligible = self.tenancies.get(&partition).is_some_and(|t| {
665                            t.state == TenancyState::Live
666                                && t.lane.is_some()
667                                && !t.fenced
668                                && !t.completed
669                        });
670                        if eligible && source.begin_revoke(&split) {
671                            // Re-borrow after `begin_revoke`; the guard
672                            // reads above kept no mutable borrow across it.
673                            if let Some(t) = self.tenancies.get_mut(&partition) {
674                                t.state = TenancyState::Draining;
675                            }
676                            true
677                        } else {
678                            false
679                        }
680                    }
681                    None => false,
682                };
683                if !accepted && let Err(e) = self.coordinator.decline_revoke(&split) {
684                    // Liveness cost only: the backend forces the release at
685                    // its own deadline regardless.
686                    tracing::warn!(split = %split, error = %e, "revocation decline failed");
687                }
688            }
689            CoordinationEvent::Lost { split } => {
690                if let Some(&partition) = self.by_split.get(&split) {
691                    self.retire(source, partition, false);
692                }
693                // Else: already retired (e.g. we fenced on commit first).
694            }
695            CoordinationEvent::Quarantined { split, attempts } => {
696                tracing::warn!(split = %split, attempts, "split quarantined");
697                if let Some(&partition) = self.by_split.get(&split) {
698                    self.retire(source, partition, false);
699                }
700            }
701            CoordinationEvent::AllComplete => {
702                if self.next_partition == 0 {
703                    tracing::info!(
704                        "coordinated job completed without this instance holding any split — \
705                         the job finished before this instance's first rebalance window, or \
706                         the fleet has more replicas than splits (see the scaling-out guide)"
707                    );
708                }
709                self.all_complete = true;
710            }
711            CoordinationEvent::Stalled {
712                completed,
713                quarantined,
714            } => {
715                self.stalled = Some((completed, quarantined));
716            }
717        }
718        Ok(())
719    }
720
721    /// End a tenancy: detach its fetchers, queue its lane for revocation,
722    /// keep the entry to absorb late watermarks until the next assignment.
723    fn retire<S: SplitSource>(&mut self, source: &mut S, partition: PartitionId, fenced: bool) {
724        let Some(tenancy) = self.tenancies.get_mut(&partition) else {
725            return;
726        };
727        if tenancy.state == TenancyState::Retired {
728            if fenced {
729                tenancy.fenced = true;
730            }
731            return;
732        }
733        tenancy.state = TenancyState::Retired;
734        tenancy.fenced |= fenced;
735        self.by_split.remove(&tenancy.split.id);
736        let split = tenancy.split.id.clone();
737        if let Some(lane) = tenancy.lane.take() {
738            if tenancy.completed || tenancy.handed_off {
739                self.pending_retired.push(lane);
740            } else {
741                self.pending_lost.push(lane);
742            }
743        }
744        source.close_split(&split);
745    }
746
747    /// Materialize lanes for the staged gains only. Each tenancy opens
748    /// exactly once, with a lane id minted for its lifetime; a staged
749    /// tenancy that ended before it could open (gained then immediately
750    /// lost or fenced) is skipped, its retirement having already handled it.
751    ///
752    /// All-or-nothing: a failure part way through undoes the whole batch
753    /// and re-stages it. Anything else strands the lanes already built.
754    /// They never reach the runtime, yet their tenancies stay `Live`
755    /// holding a lane id, which the `lane.is_some()` guard then skips
756    /// forever. The splits would keep their leases, heartbeated and
757    /// unreadable, and the job would stall instead of failing.
758    fn open_pending<S: SplitSource>(
759        &mut self,
760        source: &mut S,
761    ) -> Result<Vec<S::Lane>, SourceError> {
762        let staged = std::mem::take(&mut self.pending_open);
763        let mut lanes = Vec::with_capacity(staged.len());
764        // Tenancies this call minted a lane for, so a failure can undo them.
765        let mut opened: Vec<PartitionId> = Vec::new();
766        for idx in 0..staged.len() {
767            let partition = staged[idx];
768            let Some(tenancy) = self.tenancies.get_mut(&partition) else {
769                continue; // retired and pruned before it could open
770            };
771            if tenancy.state != TenancyState::Live || tenancy.lane.is_some() {
772                continue;
773            }
774            let lane_id = LaneId(self.next_lane);
775            self.next_lane = self
776                .next_lane
777                .checked_add(1)
778                .expect("lane ids exhausted (u32)");
779            tenancy.lane = Some(lane_id);
780            let opening = SplitOpening {
781                split: &tenancy.split,
782                resume: tenancy.progress.as_ref(),
783                lane: lane_id,
784                partition,
785                epoch: tenancy.epoch,
786                waker: &self.waker,
787            };
788            match source.open_split(opening) {
789                Ok(lane) => {
790                    lanes.push(lane);
791                    opened.push(partition);
792                }
793                Err(e) => {
794                    // Dropping the lanes detaches whatever `open_split`
795                    // spawned for them; clearing `lane` lets the retry mint
796                    // a fresh id (ids are burned, never reused).
797                    drop(lanes);
798                    for p in opened.iter().chain(std::iter::once(&partition)) {
799                        if let Some(t) = self.tenancies.get_mut(p) {
800                            t.lane = None;
801                        }
802                    }
803                    self.pending_open = staged;
804                    return Err(e);
805                }
806            }
807        }
808        Ok(lanes)
809    }
810
811    fn sweep<S: SplitSource>(&mut self, source: &mut S) -> Result<(), SourceError> {
812        let candidates: Vec<PartitionId> = self
813            .tenancies
814            .iter()
815            .filter(|(_, t)| t.state == TenancyState::Live && !t.fenced && !t.completed)
816            .map(|(&p, _)| p)
817            .collect();
818        for partition in candidates {
819            let split = self.tenancies[&partition].split.id.clone();
820            if let Some(progress) = source.sweep(&split)? {
821                self.commit_progress(source, partition, &split, progress)?;
822            }
823        }
824        Ok(())
825    }
826
827    /// Advance every in-flight cooperative revocation. For each `Draining`
828    /// tenancy whose drain has finished ([`SplitSource::drain_ready`]
829    /// returns the final progress once its tail is acked and committed),
830    /// take one last fenced commit (never `completed`; a revocation gives the
831    /// split away rather than finishing it) and dispose of it:
832    ///
833    /// - durable → mark drained, hand the split back
834    ///   ([`SplitCoordinator::release_drained`]), and retire it barrier-less;
835    /// - fenced → a peer fenced this tenancy mid-drain; retire it through
836    ///   the loss path (its bounded tail replays under the new owner);
837    /// - deferred → the store lagged; stay `Draining` and re-attempt on
838    ///   the next poll.
839    fn advance_drains<S: SplitSource>(&mut self, source: &mut S) -> Result<(), SourceError> {
840        let draining: Vec<PartitionId> = self
841            .tenancies
842            .iter()
843            .filter(|(_, t)| t.state == TenancyState::Draining && !t.fenced)
844            .map(|(&p, _)| p)
845            .collect();
846        for partition in draining {
847            let split = self.tenancies[&partition].split.id.clone();
848            let Some(progress) = source.drain_ready(&split)? else {
849                continue; // tail still in flight, retry next poll
850            };
851            debug_assert!(
852                !progress.completed,
853                "a revocation commit gives the split away, it must not complete it"
854            );
855            self.commit_drained(source, partition, &split, progress)?;
856        }
857        Ok(())
858    }
859
860    /// The disposition of one fenced commit attempt. Tick, sweep, and
861    /// drain commits triage the backend's three answers identically; each
862    /// caller owns only the durable arm and shares the fence/retry
863    /// handling here.
864    fn try_commit<S: SplitSource>(
865        &mut self,
866        source: &mut S,
867        partition: PartitionId,
868        split: &SplitId,
869        progress: &SplitProgress,
870    ) -> Result<CommitDisposition, SourceError> {
871        match self.coordinator.commit(split, progress) {
872            Ok(()) => Ok(CommitDisposition::Durable),
873            Err(e) if e.kind == CoordinationErrorKind::Fenced => {
874                // Nothing was written; the split belongs to a peer. Retire
875                // with the fence flag so nothing of this tenancy is ever
876                // folded or respawned (the matching Lost event may still
877                // arrive and finds the tenancy already retired).
878                tracing::warn!(split = %split, "commit fenced; split lost to a peer");
879                self.retire(source, partition, true);
880                Ok(CommitDisposition::Fenced)
881            }
882            Err(e) if e.kind == CoordinationErrorKind::Retryable => {
883                tracing::warn!(split = %split, error = %e, "commit deferred; will retry");
884                Ok(CommitDisposition::Deferred)
885            }
886            Err(e) => Err(as_source_error(e)),
887        }
888    }
889
890    /// Shared fenced-commit path for tick commits and sweep commits.
891    fn commit_progress<S: SplitSource>(
892        &mut self,
893        source: &mut S,
894        partition: PartitionId,
895        split: &SplitId,
896        progress: SplitProgress,
897    ) -> Result<(), SourceError> {
898        // A drain cut can look terminal to the source (every record it
899        // emitted is acked). Committing it `completed: true` marks a
900        // half-read split permanently done, and its next owner never
901        // resumes it: silent data loss.
902        let progress = if progress.completed
903            && self
904                .tenancies
905                .get(&partition)
906                .is_some_and(|t| t.state == TenancyState::Draining)
907        {
908            tracing::error!(
909                split = %split,
910                "source reported a draining split completed; forcing \
911                 completed=false — a drain cut is never terminal"
912            );
913            SplitProgress::new(progress.watermark, progress.state)
914        } else {
915            progress
916        };
917        match self.try_commit(source, partition, split, &progress)? {
918            CommitDisposition::Durable => {
919                let tenancy = self.tenancies.get_mut(&partition).expect("live tenancy");
920                let completed = progress.completed;
921                tenancy.progress = Some(progress);
922                if completed {
923                    tenancy.completed = true;
924                    // A completed split frees its lane for the working set.
925                    self.retire(source, partition, false);
926                }
927            }
928            // Fenced: already retired inside `try_commit`.
929            CommitDisposition::Fenced => {}
930            CommitDisposition::Deferred => {
931                // Nothing goes out again on a bare tick, since `commit`
932                // issues only the watermarks it is handed. This progress
933                // rides the split's next commit, which merges past it, or
934                // the terminal progress `sweep` returns if no watermark
935                // follows. The resume cache still advances: the watermark
936                // is acked (sink-durable), so respawning past it cannot
937                // lose data.
938                let tenancy = self.tenancies.get_mut(&partition).expect("live tenancy");
939                tenancy.progress = Some(progress);
940            }
941        }
942        Ok(())
943    }
944
945    /// Final fenced commit that ends a cooperative revocation. Same triage as
946    /// [`commit_progress`](CoordinationDriver::commit_progress), but a
947    /// durable commit hands the split back and retires it barrier-less
948    /// instead of folding progress into a still-live tenancy.
949    fn commit_drained<S: SplitSource>(
950        &mut self,
951        source: &mut S,
952        partition: PartitionId,
953        split: &SplitId,
954        progress: SplitProgress,
955    ) -> Result<(), SourceError> {
956        // Same guard as `commit_progress`: drain progress is never terminal.
957        let progress = if progress.completed {
958            tracing::error!(
959                split = %split,
960                "drain_ready returned completed=true; forcing completed=false — \
961                 a drain cut is never terminal"
962            );
963            SplitProgress::new(progress.watermark, progress.state)
964        } else {
965            progress
966        };
967        match self.try_commit(source, partition, split, &progress)? {
968            CommitDisposition::Durable => {
969                let tenancy = self
970                    .tenancies
971                    .get_mut(&partition)
972                    .expect("draining tenancy");
973                tenancy.progress = Some(progress);
974                tenancy.handed_off = true;
975                if let Err(e) = self
976                    .coordinator
977                    .release_drained(std::slice::from_ref(split))
978                {
979                    // Liveness cost only: the lease expires on its own and a
980                    // peer takes over, with no data at risk. Retire anyway so
981                    // the lane leaves this instance.
982                    tracing::warn!(
983                        split = %split,
984                        error = %e,
985                        "drain release failed; lease will expire and a peer will take over"
986                    );
987                }
988                self.retire(source, partition, false);
989            }
990            // Fenced mid-drain: already retired through the loss path.
991            CommitDisposition::Fenced => {}
992            CommitDisposition::Deferred => {
993                // Store lagged: stay `Draining` and re-attempt next poll.
994                let tenancy = self
995                    .tenancies
996                    .get_mut(&partition)
997                    .expect("draining tenancy");
998                tenancy.progress = Some(progress);
999            }
1000        }
1001        Ok(())
1002    }
1003}
1004
1005/// How the backend answered one fenced commit; see
1006/// [`CoordinationDriver::try_commit`].
1007enum CommitDisposition {
1008    /// Durable write. The caller advances its own state.
1009    Durable,
1010    /// Fenced: nothing written, the split belongs to a peer; the tenancy
1011    /// has already been retired with the fence flag.
1012    Fenced,
1013    /// Retryable: nothing was written, the previous durable state stays
1014    /// authoritative, and each caller decides when the progress goes out
1015    /// again.
1016    Deferred,
1017}
1018
1019fn is_fatal(e: &SourceError) -> bool {
1020    let SourceError::Client { class, .. } = e;
1021    *class == ErrorClass::Fatal
1022}
1023
1024fn as_source_error(e: CoordinationError) -> SourceError {
1025    SourceError::Client {
1026        class: e.class(),
1027        reason: e.to_string(),
1028    }
1029}
1030
1031#[cfg(test)]
1032mod tests {
1033    use super::*;
1034    use crate::checkpoint::AckRef;
1035    use crate::coordination::{PlanContext, PlanFinality, SplitPlan};
1036    use crate::record::RawPayload;
1037    use crate::source::PayloadBatch;
1038    use std::cell::RefCell;
1039    use std::collections::{HashMap, HashSet, VecDeque};
1040    use std::rc::Rc;
1041    use std::sync::{Arc, Mutex};
1042    use std::time::Instant;
1043
1044    // ------------------------------------------------------------------
1045    // Scripted coordinator double (the shape spate-test publishes).
1046
1047    #[derive(Default)]
1048    struct ScriptState {
1049        batches: VecDeque<Vec<CoordinationEvent>>,
1050        commit_outcomes: HashMap<String, VecDeque<CoordinationErrorKind>>,
1051        commits: Vec<(SplitId, SplitProgress)>,
1052        fail_outcomes: HashMap<String, VecDeque<CoordinationErrorKind>>,
1053        /// Every `fail` call, including the ones `fail_outcomes` refuses;
1054        /// the attempt is what a test asserts the driver made.
1055        fails: Vec<(SplitId, String)>,
1056        released: Vec<SplitId>,
1057        /// Captured separately from `released` so a test can prove the
1058        /// driver takes the revocation-release path, not a plain hand-back.
1059        released_drained: Vec<SplitId>,
1060        /// Every `decline_revoke` call, so a test can prove the driver feeds
1061        /// a refusal back to the backend, which then forces the release.
1062        declined: Vec<SplitId>,
1063        started: bool,
1064        waker: Option<ControlWaker>,
1065    }
1066
1067    #[derive(Clone, Default)]
1068    struct Script(Arc<Mutex<ScriptState>>);
1069
1070    impl Script {
1071        fn push(&self, events: Vec<CoordinationEvent>) {
1072            let mut st = self.0.lock().unwrap();
1073            st.batches.push_back(events);
1074            if let Some(w) = &st.waker {
1075                w.wake();
1076            }
1077        }
1078
1079        fn fail_next_commit(&self, split: &str, kind: CoordinationErrorKind) {
1080            self.0
1081                .lock()
1082                .unwrap()
1083                .commit_outcomes
1084                .entry(split.to_string())
1085                .or_default()
1086                .push_back(kind);
1087        }
1088
1089        fn fail_next_report(&self, split: &str, kind: CoordinationErrorKind) {
1090            self.0
1091                .lock()
1092                .unwrap()
1093                .fail_outcomes
1094                .entry(split.to_string())
1095                .or_default()
1096                .push_back(kind);
1097        }
1098
1099        fn commits(&self) -> Vec<(SplitId, SplitProgress)> {
1100            self.0.lock().unwrap().commits.clone()
1101        }
1102
1103        fn released(&self) -> Vec<SplitId> {
1104            self.0.lock().unwrap().released.clone()
1105        }
1106
1107        fn released_drained(&self) -> Vec<SplitId> {
1108            self.0.lock().unwrap().released_drained.clone()
1109        }
1110
1111        fn declined(&self) -> Vec<SplitId> {
1112            self.0.lock().unwrap().declined.clone()
1113        }
1114
1115        fn fails(&self) -> Vec<(SplitId, String)> {
1116            self.0.lock().unwrap().fails.clone()
1117        }
1118    }
1119
1120    struct ScriptedCoordinator(Script);
1121
1122    impl SplitCoordinator for ScriptedCoordinator {
1123        fn start(&mut self, _planner: Box<dyn SplitPlanner>) -> Result<(), CoordinationError> {
1124            self.0.0.lock().unwrap().started = true;
1125            Ok(())
1126        }
1127
1128        fn set_waker(&mut self, waker: ControlWaker) {
1129            self.0.0.lock().unwrap().waker = Some(waker);
1130        }
1131
1132        fn poll(&mut self) -> Result<Vec<CoordinationEvent>, CoordinationError> {
1133            Ok(self
1134                .0
1135                .0
1136                .lock()
1137                .unwrap()
1138                .batches
1139                .pop_front()
1140                .unwrap_or_default())
1141        }
1142
1143        fn commit(
1144            &mut self,
1145            split: &SplitId,
1146            progress: &SplitProgress,
1147        ) -> Result<(), CoordinationError> {
1148            let mut s = self.0.0.lock().unwrap();
1149            if let Some(kinds) = s.commit_outcomes.get_mut(split.as_str())
1150                && let Some(kind) = kinds.pop_front()
1151            {
1152                return Err(CoordinationError::new(kind, "scripted"));
1153            }
1154            s.commits.push((split.clone(), progress.clone()));
1155            Ok(())
1156        }
1157
1158        fn fail(&mut self, split: &SplitId, reason: &str) -> Result<(), CoordinationError> {
1159            let mut s = self.0.0.lock().unwrap();
1160            s.fails.push((split.clone(), reason.to_string()));
1161            if let Some(kinds) = s.fail_outcomes.get_mut(split.as_str())
1162                && let Some(kind) = kinds.pop_front()
1163            {
1164                return Err(CoordinationError::new(kind, "scripted"));
1165            }
1166            Ok(())
1167        }
1168
1169        fn release(&mut self, splits: &[SplitId]) -> Result<(), CoordinationError> {
1170            self.0
1171                .0
1172                .lock()
1173                .unwrap()
1174                .released
1175                .extend(splits.iter().cloned());
1176            Ok(())
1177        }
1178
1179        fn release_drained(&mut self, splits: &[SplitId]) -> Result<(), CoordinationError> {
1180            self.0
1181                .0
1182                .lock()
1183                .unwrap()
1184                .released_drained
1185                .extend(splits.iter().cloned());
1186            Ok(())
1187        }
1188
1189        fn decline_revoke(&mut self, split: &SplitId) -> Result<(), CoordinationError> {
1190            self.0.0.lock().unwrap().declined.push(split.clone());
1191            Ok(())
1192        }
1193    }
1194
1195    struct NoopPlanner;
1196
1197    impl SplitPlanner for NoopPlanner {
1198        fn fingerprint(&self) -> String {
1199            "test:v1".into()
1200        }
1201
1202        fn plan(&mut self, _ctx: PlanContext<'_>) -> Result<SplitPlan, CoordinationError> {
1203            Ok(SplitPlan::new(vec![], PlanFinality::Final))
1204        }
1205    }
1206
1207    // ------------------------------------------------------------------
1208    // Stub data plane.
1209
1210    enum NoBatch {}
1211
1212    impl<'buf> PayloadBatch<'buf> for NoBatch {
1213        fn next_payload(&mut self) -> Option<RawPayload<'buf>> {
1214            match *self {}
1215        }
1216
1217        fn ack(&self) -> &AckRef {
1218            match *self {}
1219        }
1220    }
1221
1222    #[derive(Debug)]
1223    struct StubLane {
1224        lane: LaneId,
1225        partition: PartitionId,
1226    }
1227
1228    impl SourceLane for StubLane {
1229        type Batch<'a> = NoBatch;
1230
1231        fn id(&self) -> LaneId {
1232            self.lane
1233        }
1234
1235        fn partition(&self) -> PartitionId {
1236            self.partition
1237        }
1238
1239        fn poll(
1240            &mut self,
1241            _max: usize,
1242            _timeout: Duration,
1243        ) -> Result<Option<NoBatch>, SourceError> {
1244            Ok(None)
1245        }
1246    }
1247
1248    /// Recording SplitSource: every callback is logged; sweep and
1249    /// encode_commit outcomes are scripted per split.
1250    #[derive(Default)]
1251    struct TestSource {
1252        opened: Vec<(String, Option<i64>, LaneId, PartitionId, u64)>,
1253        closed: Vec<String>,
1254        encoded: Vec<(String, i64)>,
1255        sweeps: Rc<RefCell<HashMap<String, SplitProgress>>>,
1256        complete_at: HashMap<String, i64>,
1257        /// Split ids whose carried progress `validate_resume` refuses, and
1258        /// the class each refusal carries, so a batch can mix a drifted
1259        /// split with a sound one, and one class with another.
1260        reject_resume: HashMap<String, ErrorClass>,
1261        finishing: Vec<String>,
1262        /// Split ids whose `open_split` fails. Consumed per attempt, so a
1263        /// retry of the same split succeeds.
1264        fail_open: Vec<String>,
1265        /// Split ids for which `begin_revoke` accepts (returns true). Empty
1266        /// by default, so the double declines like the trait default.
1267        accept_revoke: HashSet<String>,
1268        /// Every `begin_revoke` call, in order (accepted or declined).
1269        begin_revoke_calls: Vec<String>,
1270        /// Every `drain_ready` call, in order. Proves a split did (or did
1271        /// not) transition to `Draining`.
1272        drain_ready_calls: Vec<String>,
1273        /// Scripted `drain_ready` results, sticky per split (returned on
1274        /// every poll until the tenancy retires), so a retryable final
1275        /// commit can be re-offered the same tail next poll.
1276        ready_progress: Rc<RefCell<HashMap<String, SplitProgress>>>,
1277    }
1278
1279    impl SplitSource for TestSource {
1280        type Lane = StubLane;
1281
1282        fn open_split(&mut self, o: SplitOpening<'_>) -> Result<StubLane, SourceError> {
1283            let id = o.split.id.as_str().to_string();
1284            if let Some(i) = self.fail_open.iter().position(|s| *s == id) {
1285                self.fail_open.remove(i);
1286                return Err(SourceError::Client {
1287                    class: ErrorClass::Retryable,
1288                    reason: format!("open_split failed for {id}"),
1289                });
1290            }
1291            self.opened.push((
1292                o.split.id.as_str().to_string(),
1293                o.resume.map(|p| p.watermark),
1294                o.lane,
1295                o.partition,
1296                o.epoch.0,
1297            ));
1298            Ok(StubLane {
1299                lane: o.lane,
1300                partition: o.partition,
1301            })
1302        }
1303
1304        fn validate_resume(
1305            &self,
1306            split: &SplitSpec,
1307            _progress: &SplitProgress,
1308        ) -> Result<(), SourceError> {
1309            if let Some(&class) = self.reject_resume.get(split.id.as_str()) {
1310                return Err(SourceError::Client {
1311                    class,
1312                    reason: format!("resume drift on {}", split.id),
1313                });
1314            }
1315            Ok(())
1316        }
1317
1318        fn encode_commit(
1319            &mut self,
1320            split: &SplitId,
1321            watermark: i64,
1322        ) -> Result<SplitProgress, SourceError> {
1323            self.encoded.push((split.as_str().to_string(), watermark));
1324            let completed = self.complete_at.get(split.as_str()) == Some(&watermark);
1325            Ok(if completed {
1326                SplitProgress::completed(watermark, vec![])
1327            } else {
1328                SplitProgress::new(watermark, vec![])
1329            })
1330        }
1331
1332        fn sweep(&mut self, split: &SplitId) -> Result<Option<SplitProgress>, SourceError> {
1333            Ok(self.sweeps.borrow_mut().remove(split.as_str()))
1334        }
1335
1336        fn close_split(&mut self, split: &SplitId) {
1337            self.closed.push(split.as_str().to_string());
1338        }
1339
1340        fn take_finishing(&mut self) -> Vec<SplitId> {
1341            std::mem::take(&mut self.finishing)
1342                .into_iter()
1343                .map(|s| SplitId::new(&s).unwrap())
1344                .collect()
1345        }
1346
1347        fn begin_revoke(&mut self, split: &SplitId) -> bool {
1348            self.begin_revoke_calls.push(split.as_str().to_string());
1349            self.accept_revoke.contains(split.as_str())
1350        }
1351
1352        fn drain_ready(&mut self, split: &SplitId) -> Result<Option<SplitProgress>, SourceError> {
1353            self.drain_ready_calls.push(split.as_str().to_string());
1354            Ok(self.ready_progress.borrow().get(split.as_str()).cloned())
1355        }
1356    }
1357
1358    // ------------------------------------------------------------------
1359    // Helpers.
1360
1361    fn split(id: &str) -> SplitSpec {
1362        SplitSpec::new(SplitId::new(id).unwrap(), format!("desc:{id}").into_bytes())
1363    }
1364
1365    fn gained(id: &str, epoch: u64, watermark: Option<i64>) -> CoordinationEvent {
1366        CoordinationEvent::Gained {
1367            split: split(id),
1368            epoch: LeaseEpoch(epoch),
1369            progress: watermark.map(|w| SplitProgress::new(w, vec![])),
1370        }
1371    }
1372
1373    fn driver(script: &Script) -> CoordinationDriver {
1374        let mut d = CoordinationDriver::new(Box::new(ScriptedCoordinator(script.clone())));
1375        let ready: SourceEvent<StubLane> = d.start(Box::new(NoopPlanner)).unwrap();
1376        assert!(
1377            matches!(ready, SourceEvent::LanesAssigned(ref lanes) if lanes.is_empty()),
1378            "start must return the empty ready signal"
1379        );
1380        d
1381    }
1382
1383    fn poll(d: &mut CoordinationDriver, s: &mut TestSource) -> SourceEvent<StubLane> {
1384        d.poll_events(s, Duration::ZERO).unwrap()
1385    }
1386
1387    /// A source whose `validate_resume` refuses exactly these splits,
1388    /// classed `Fatal` the way a connector is told to.
1389    fn rejecting(splits: &[&str]) -> TestSource {
1390        TestSource {
1391            reject_resume: splits
1392                .iter()
1393                .map(|s| ((*s).to_string(), ErrorClass::Fatal))
1394                .collect(),
1395            ..TestSource::default()
1396        }
1397    }
1398
1399    // ------------------------------------------------------------------
1400    // Scenarios (each replays a defect class from the PR #34 review).
1401
1402    #[test]
1403    fn a_signal_cuts_the_control_plane_park_short() {
1404        // The driver owns the control-plane wait so that both producers can
1405        // end it: the backend, and a *lane* deciding end-of-input on a
1406        // pipeline thread. If a `wake()` call site is ever dropped, the
1407        // symptom is silent (completions wait out an idle timeout again), so
1408        // assert the park is interruptible rather than trusting the wiring.
1409        let script = Script::default();
1410        let mut d = driver(&script);
1411        let mut s = TestSource::default();
1412        let park = Duration::from_millis(400);
1413
1414        // Control: nothing pending and nothing signaling, so the full
1415        // timeout elapses. Without this the test would pass even if
1416        // `poll_events` never parked at all.
1417        let t0 = Instant::now();
1418        assert!(matches!(
1419            d.poll_events(&mut s, park).unwrap(),
1420            SourceEvent::Idle
1421        ));
1422        let idle = t0.elapsed();
1423        assert!(
1424            idle >= park / 2,
1425            "expected a real park, returned after {idle:?}"
1426        );
1427
1428        // A signal landing mid-park ends it. The event itself surfaces on
1429        // the following call (the drain runs at the top of `poll_events`),
1430        // so this asserts the wakeup, not the delivery.
1431        let signaller = script.clone();
1432        let handle = std::thread::spawn(move || {
1433            std::thread::sleep(Duration::from_millis(20));
1434            signaller.push(vec![CoordinationEvent::AllComplete]);
1435        });
1436        let t1 = Instant::now();
1437        let _ = d.poll_events(&mut s, park).unwrap();
1438        let woken = t1.elapsed();
1439        handle.join().unwrap();
1440        assert!(
1441            woken < park / 2,
1442            "a signal must cut the park short, but it ran {woken:?} of {park:?}"
1443        );
1444        assert!(matches!(
1445            d.poll_events(&mut s, Duration::ZERO).unwrap(),
1446            SourceEvent::Drained
1447        ));
1448    }
1449
1450    #[test]
1451    fn a_failed_open_undoes_the_whole_batch_instead_of_stranding_lanes() {
1452        // Lanes already built when a later `open_split` fails never reach the
1453        // runtime, yet their tenancies keep a lane id and hold their leases,
1454        // heartbeated and unreadable: a stalled job rather than a failed one.
1455        let script = Script::default();
1456        let mut d = driver(&script);
1457        let mut s = TestSource {
1458            fail_open: vec!["b".into()],
1459            ..TestSource::default()
1460        };
1461
1462        script.push(vec![gained("a", 1, None), gained("b", 1, None)]);
1463        let err = d
1464            .poll_events(&mut s, Duration::ZERO)
1465            .expect_err("the failing open must surface");
1466        assert!(err.to_string().contains("open_split failed for b"), "{err}");
1467        assert_eq!(s.opened.len(), 1, "a opened before b failed");
1468
1469        // The retry re-stages the whole batch and yields both lanes.
1470        let event = poll(&mut d, &mut s);
1471        let SourceEvent::LanesAdded(lanes) = event else {
1472            panic!("expected both lanes after the retry, got {event:?}");
1473        };
1474        assert_eq!(lanes.len(), 2);
1475        let reopened: Vec<&str> = s.opened.iter().map(|o| o.0.as_str()).collect();
1476        assert_eq!(reopened, ["a", "a", "b"], "a re-opens on the retry");
1477        // The rolled-back ids are burned, never reused.
1478        assert_eq!(lanes[0].id(), LaneId(2));
1479        assert_eq!(lanes[1].id(), LaneId(3));
1480    }
1481
1482    #[test]
1483    fn gains_coalesce_into_one_added_batch() {
1484        let script = Script::default();
1485        let mut d = driver(&script);
1486        let mut s = TestSource::default();
1487
1488        script.push(vec![gained("b", 1, Some(7)), gained("a", 1, None)]);
1489        let event = poll(&mut d, &mut s);
1490        let SourceEvent::LanesAdded(lanes) = event else {
1491            panic!("expected added lanes, got {event:?}");
1492        };
1493        assert_eq!(lanes.len(), 2);
1494        // Lane ids minted in gain order; distinct tenancy partitions.
1495        assert_eq!(s.opened[0].0, "b");
1496        assert_eq!(s.opened[0].2, LaneId(0));
1497        assert_eq!(s.opened[0].1, Some(7), "carried progress reaches open");
1498        assert_eq!(s.opened[1].0, "a");
1499        assert_eq!(s.opened[1].2, LaneId(1));
1500        assert_ne!(s.opened[0].3, s.opened[1].3);
1501        assert_eq!(d.assignments().len(), 2);
1502    }
1503
1504    #[test]
1505    fn a_mid_flow_gain_never_touches_live_lanes_and_their_commits_fold() {
1506        let script = Script::default();
1507        let mut d = driver(&script);
1508        let mut s = TestSource::default();
1509        script.push(vec![gained("a", 1, None)]);
1510        poll(&mut d, &mut s);
1511        let a_partition = s.opened[0].3;
1512
1513        // Split b arrives while a is live and flowing: strictly additive.
1514        script.push(vec![gained("b", 1, None)]);
1515        let event = poll(&mut d, &mut s);
1516        let SourceEvent::LanesAdded(lanes) = event else {
1517            panic!("expected added lanes, got {event:?}");
1518        };
1519        assert_eq!(lanes.len(), 1, "only the new split's lane");
1520        assert!(
1521            s.closed.is_empty(),
1522            "a routine gain must never detach flowing fetchers"
1523        );
1524
1525        // The narrow commit window: a's acked watermark lands right after
1526        // the gain. It must fold normally.
1527        d.commit(&mut s, &[(a_partition, 42)]).unwrap();
1528        assert_eq!(s.encoded, vec![("a".to_string(), 42)]);
1529        assert_eq!(script.commits().len(), 1);
1530        assert_eq!(script.commits()[0].0.as_str(), "a");
1531        // a's lane is the original, never re-minted by the gain.
1532        assert!(
1533            d.assignments()
1534                .contains(&(SplitId::new("a").unwrap(), LaneId(0)))
1535        );
1536    }
1537
1538    #[test]
1539    fn finishing_splits_surface_as_commit_ready_once() {
1540        let script = Script::default();
1541        let mut d = driver(&script);
1542        let mut s = TestSource::default();
1543        script.push(vec![gained("a", 1, None)]);
1544        poll(&mut d, &mut s);
1545        let a_partition = s.opened[0].3;
1546
1547        s.finishing.push("a".to_string());
1548        let event = poll(&mut d, &mut s);
1549        let SourceEvent::CommitReady { partitions } = event else {
1550            panic!("expected commit-ready, got {event:?}");
1551        };
1552        assert_eq!(partitions, vec![a_partition]);
1553        // Edge, not level: the hint is consumed.
1554        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
1555    }
1556
1557    #[test]
1558    fn loss_surfaces_as_partial_revoke_and_detaches_fetchers() {
1559        let script = Script::default();
1560        let mut d = driver(&script);
1561        let mut s = TestSource::default();
1562        script.push(vec![gained("a", 1, None), gained("b", 1, None)]);
1563        poll(&mut d, &mut s);
1564
1565        script.push(vec![CoordinationEvent::Lost {
1566            split: SplitId::new("a").unwrap(),
1567        }]);
1568        let event = poll(&mut d, &mut s);
1569        let SourceEvent::LanesRevoked { lanes, barrier } = event else {
1570            panic!("expected revoke, got {event:?}");
1571        };
1572        assert_eq!(lanes, vec![LaneId(0)]);
1573        assert_eq!(barrier.remaining(), 1, "one party per revoked lane");
1574        assert_eq!(s.closed, vec!["a"], "fetcher detached on loss");
1575        assert_eq!(d.assignments().len(), 1);
1576    }
1577
1578    #[test]
1579    fn late_drain_commit_after_loss_is_skipped() {
1580        let script = Script::default();
1581        let mut d = driver(&script);
1582        let mut s = TestSource::default();
1583        script.push(vec![gained("a", 1, None)]);
1584        poll(&mut d, &mut s);
1585        let partition = s.opened[0].3;
1586
1587        script.push(vec![CoordinationEvent::Lost {
1588            split: SplitId::new("a").unwrap(),
1589        }]);
1590        poll(&mut d, &mut s);
1591
1592        // The drain hands back a final watermark for the retired tenancy.
1593        d.commit(&mut s, &[(partition, 42)]).unwrap();
1594        assert!(s.encoded.is_empty(), "retired tenancy must not encode");
1595        assert!(script.commits().is_empty(), "and must not commit");
1596    }
1597
1598    #[test]
1599    fn fenced_commit_quarantines_the_tenancy_and_never_respawns_it() {
1600        let script = Script::default();
1601        let mut d = driver(&script);
1602        let mut s = TestSource::default();
1603        script.push(vec![gained("a", 1, None), gained("b", 1, None)]);
1604        poll(&mut d, &mut s);
1605        let a_partition = s.opened[0].3;
1606        let b_partition = s.opened[1].3;
1607
1608        script.fail_next_commit("a", CoordinationErrorKind::Fenced);
1609        d.commit(&mut s, &[(a_partition, 10), (b_partition, 20)])
1610            .unwrap();
1611        // b committed; a wrote nothing and is retired with the fence flag.
1612        assert_eq!(script.commits().len(), 1);
1613        assert_eq!(script.commits()[0].0.as_str(), "b");
1614        assert_eq!(s.closed, vec!["a"]);
1615
1616        // The fenced lane is revoked...
1617        let event = poll(&mut d, &mut s);
1618        assert!(
1619            matches!(event, SourceEvent::LanesRevoked { ref lanes, .. } if lanes[..] == [LaneId(0)])
1620        );
1621
1622        // ...a late watermark for it is skipped...
1623        s.encoded.clear();
1624        d.commit(&mut s, &[(a_partition, 11)]).unwrap();
1625        assert!(s.encoded.is_empty());
1626
1627        // ...and the mid-cycle Lost that follows the fence is a no-op,
1628        // while a re-gain (higher epoch) starts a fresh tenancy, added
1629        // beside b's untouched live lane and never draining it.
1630        script.push(vec![
1631            CoordinationEvent::Lost {
1632                split: SplitId::new("a").unwrap(),
1633            },
1634            gained("a", 3, Some(10)),
1635        ]);
1636        let event = poll(&mut d, &mut s);
1637        let SourceEvent::LanesAdded(lanes) = event else {
1638            panic!("expected an added lane for the re-gain, got {event:?}");
1639        };
1640        assert_eq!(lanes.len(), 1, "only the fresh tenancy's lane");
1641        assert_eq!(s.closed, vec!["a"], "b's fetcher was never detached");
1642        let a_again = s.opened.last().unwrap();
1643        assert_eq!(a_again.0, "a");
1644        assert_eq!(a_again.4, 3, "fresh tenancy under the new epoch");
1645        assert_ne!(a_again.3, a_partition, "fresh partition — no reuse");
1646        assert_eq!(a_again.2, LaneId(2), "fresh lane id — never reused");
1647    }
1648
1649    #[test]
1650    fn lost_then_regained_in_one_batch_is_a_clean_tenancy_swap() {
1651        let script = Script::default();
1652        let mut d = driver(&script);
1653        let mut s = TestSource::default();
1654        script.push(vec![gained("a", 1, None)]);
1655        poll(&mut d, &mut s);
1656        let first_partition = s.opened[0].3;
1657
1658        script.push(vec![
1659            CoordinationEvent::Lost {
1660                split: SplitId::new("a").unwrap(),
1661            },
1662            gained("a", 2, Some(5)),
1663        ]);
1664        // Loss first (revoke), then the re-gain's lane is added fresh.
1665        let event = poll(&mut d, &mut s);
1666        assert!(matches!(event, SourceEvent::LanesRevoked { .. }));
1667        let event = poll(&mut d, &mut s);
1668        assert!(matches!(event, SourceEvent::LanesAdded(ref l) if l.len() == 1));
1669        let reopened = s.opened.last().unwrap();
1670        assert_eq!(reopened.4, 2);
1671        assert_eq!(reopened.1, Some(5), "resume from the carried progress");
1672        assert_ne!(reopened.3, first_partition);
1673        assert_eq!(reopened.2, LaneId(1), "lane ids are never recycled");
1674    }
1675
1676    #[test]
1677    fn retryable_commit_defers_and_recommits_idempotently() {
1678        let script = Script::default();
1679        let mut d = driver(&script);
1680        let mut s = TestSource::default();
1681        script.push(vec![gained("a", 1, None)]);
1682        poll(&mut d, &mut s);
1683        let partition = s.opened[0].3;
1684
1685        script.fail_next_commit("a", CoordinationErrorKind::Retryable);
1686        d.commit(&mut s, &[(partition, 10)]).unwrap();
1687        assert!(script.commits().is_empty(), "deferred, not written");
1688
1689        // A bare tick re-issues nothing: the driver commits the watermarks
1690        // it is handed and no others.
1691        poll(&mut d, &mut s);
1692        assert!(script.commits().is_empty(), "a tick recommitted nothing");
1693
1694        // The split's next commit carries the merged progress and lands.
1695        d.commit(&mut s, &[(partition, 12)]).unwrap();
1696        assert_eq!(script.commits().len(), 1);
1697        assert_eq!(script.commits()[0].1.watermark, 12);
1698    }
1699
1700    #[test]
1701    fn completion_sweep_commits_terminal_progress_and_frees_the_lane() {
1702        let script = Script::default();
1703        let mut d = driver(&script);
1704        let mut s = TestSource::default();
1705        script.push(vec![gained("a", 1, None)]);
1706        poll(&mut d, &mut s);
1707
1708        s.sweeps
1709            .borrow_mut()
1710            .insert("a".into(), SplitProgress::completed(9, vec![]));
1711        // The sweep commits terminal progress and retires the lane; being
1712        // complete (nothing in flight by construction), it leaves
1713        // barrier-less on this same poll (staged-work fastpath).
1714        let event = poll(&mut d, &mut s);
1715        assert!(
1716            matches!(event, SourceEvent::LanesRetired { ref lanes } if lanes[..] == [LaneId(0)]),
1717            "completed lanes retire without a drain barrier, got {event:?}"
1718        );
1719        assert_eq!(script.commits().len(), 1);
1720        assert!(script.commits()[0].1.completed);
1721
1722        // A watermark-carrying commit that completes a split does the same.
1723        script.push(vec![gained("b", 1, None)]);
1724        poll(&mut d, &mut s);
1725        let b_partition = s.opened.last().unwrap().3;
1726        s.complete_at.insert("b".into(), 20);
1727        d.commit(&mut s, &[(b_partition, 20)]).unwrap();
1728        assert!(script.commits().last().unwrap().1.completed);
1729        let event = poll(&mut d, &mut s);
1730        assert!(matches!(event, SourceEvent::LanesRetired { .. }));
1731    }
1732
1733    #[test]
1734    fn standby_with_zero_splits_drains_on_all_complete() {
1735        let script = Script::default();
1736        let mut d = driver(&script);
1737        let mut s = TestSource::default();
1738
1739        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
1740        script.push(vec![CoordinationEvent::AllComplete]);
1741        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Drained));
1742        // Idempotent thereafter.
1743        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Drained));
1744    }
1745
1746    #[test]
1747    fn stalled_is_fatal_by_default_and_drains_when_configured() {
1748        let script = Script::default();
1749        let mut d = driver(&script);
1750        let mut s = TestSource::default();
1751        script.push(vec![CoordinationEvent::Stalled {
1752            completed: 7,
1753            quarantined: 1,
1754        }]);
1755        // The stall surfaces on the same call that absorbed the event.
1756        let err = d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1757        assert!(err.to_string().contains("quarantined"), "{err}");
1758
1759        let script = Script::default();
1760        let mut d = CoordinationDriver::new(Box::new(ScriptedCoordinator(script.clone())))
1761            .stall_drains(true);
1762        let _: SourceEvent<StubLane> = d.start(Box::new(NoopPlanner)).unwrap();
1763        script.push(vec![CoordinationEvent::Stalled {
1764            completed: 7,
1765            quarantined: 1,
1766        }]);
1767        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Drained));
1768    }
1769
1770    #[test]
1771    fn fail_reports_poison_and_retires_the_lane() {
1772        let script = Script::default();
1773        let mut d = driver(&script);
1774        let mut s = TestSource::default();
1775        script.push(vec![gained("a", 1, None)]);
1776        poll(&mut d, &mut s);
1777
1778        d.fail(&mut s, &SplitId::new("a").unwrap(), "undecodable object")
1779            .unwrap();
1780        assert_eq!(script.fails().len(), 1);
1781        assert_eq!(s.closed, vec!["a"]);
1782        let event = poll(&mut d, &mut s);
1783        assert!(matches!(event, SourceEvent::LanesRevoked { .. }));
1784        // Failing a split we no longer hold is a quiet no-op.
1785        d.fail(&mut s, &SplitId::new("a").unwrap(), "again")
1786            .unwrap();
1787        assert_eq!(script.fails().len(), 1);
1788    }
1789
1790    #[test]
1791    fn release_hands_back_every_live_split() {
1792        let script = Script::default();
1793        let mut d = driver(&script);
1794        let mut s = TestSource::default();
1795        script.push(vec![gained("a", 1, None), gained("b", 1, None)]);
1796        poll(&mut d, &mut s);
1797
1798        d.release();
1799        let released = script.released();
1800        assert_eq!(released.len(), 2);
1801        assert!(released.iter().any(|s| s.as_str() == "a"));
1802        assert!(released.iter().any(|s| s.as_str() == "b"));
1803    }
1804
1805    #[test]
1806    fn resume_validation_rejects_drifted_progress() {
1807        let script = Script::default();
1808        let mut d = driver(&script);
1809        let mut s = rejecting(&["a"]);
1810        script.push(vec![gained("a", 1, Some(7))]);
1811        let err = d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1812        assert!(err.to_string().contains("resume drift"), "{err}");
1813
1814        // The split is never opened, so nothing else would hand it back:
1815        // the driver reports it as poison, carrying the source's reason.
1816        assert!(s.opened.is_empty());
1817        let fails = script.fails();
1818        assert_eq!(fails.len(), 1);
1819        assert_eq!(fails[0].0.as_str(), "a");
1820        assert!(fails[0].1.contains("resume drift on a"), "{}", fails[0].1);
1821    }
1822
1823    #[test]
1824    fn a_refused_resume_leaves_the_rest_of_the_batch_applied() {
1825        let script = Script::default();
1826        let mut d = driver(&script);
1827        let mut s = rejecting(&["a"]);
1828        // One batch, one drifted split: `poll` already drained the batch, so
1829        // a sound split behind the rejection is never re-offered.
1830        script.push(vec![gained("a", 1, Some(7)), gained("b", 1, Some(3))]);
1831        let err = d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1832        assert!(err.to_string().contains("resume drift on a"), "{err}");
1833        assert_eq!(script.fails().len(), 1);
1834
1835        // `b` was staged by the same batch and opens on the next poll.
1836        let event = poll(&mut d, &mut s);
1837        let SourceEvent::LanesAdded(lanes) = event else {
1838            panic!("expected the sound split to open");
1839        };
1840        assert_eq!(lanes.len(), 1);
1841        assert_eq!(s.opened.len(), 1);
1842        assert_eq!(s.opened[0].0, "b");
1843        assert_eq!(s.opened[0].1, Some(3));
1844    }
1845
1846    #[test]
1847    fn a_failed_poison_report_does_not_replace_the_rejection() {
1848        let script = Script::default();
1849        let mut d = driver(&script);
1850        let mut s = rejecting(&["a"]);
1851        script.fail_next_report("a", CoordinationErrorKind::Retryable);
1852        script.push(vec![gained("a", 1, Some(7))]);
1853
1854        // The source's rejection is the error the pipeline classes; a
1855        // backend that would not take the report cannot downgrade it.
1856        let err = d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1857        assert!(err.to_string().contains("resume drift on a"), "{err}");
1858        assert_eq!(script.fails().len(), 1);
1859    }
1860
1861    #[test]
1862    fn a_refused_poison_report_is_retried_until_it_lands() {
1863        let script = Script::default();
1864        let mut d = driver(&script);
1865        let mut s = rejecting(&["a"]);
1866        script.fail_next_report("a", CoordinationErrorKind::Retryable);
1867        script.push(vec![gained("a", 1, Some(7))]);
1868        d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1869        assert_eq!(script.fails().len(), 1);
1870
1871        // A refused report leaves the split held here, so it is re-offered
1872        // until the backend takes it, and then stops.
1873        poll(&mut d, &mut s);
1874        assert_eq!(script.fails().len(), 2);
1875        poll(&mut d, &mut s);
1876        assert_eq!(script.fails().len(), 2);
1877    }
1878
1879    #[test]
1880    fn a_fatal_rejection_survives_an_earlier_retryable_one() {
1881        let script = Script::default();
1882        let mut d = driver(&script);
1883        let mut s = TestSource {
1884            reject_resume: HashMap::from([
1885                ("a".to_string(), ErrorClass::Retryable),
1886                ("b".to_string(), ErrorClass::Fatal),
1887            ]),
1888            ..TestSource::default()
1889        };
1890        script.push(vec![gained("a", 1, Some(7)), gained("b", 1, Some(3))]);
1891
1892        // Surfacing the retryable rejection would run past a stop `b`'s
1893        // source asked for. Both splits are still handed back.
1894        let err = d.poll_events(&mut s, Duration::ZERO).unwrap_err();
1895        assert!(is_fatal(&err), "{err}");
1896        assert!(err.to_string().contains("resume drift on b"), "{err}");
1897        assert_eq!(script.fails().len(), 2);
1898    }
1899
1900    // ------------------------------------------------------------------
1901    // Cooperative revocation: request → drain → commit → release,
1902    // barrier-less; a fence mid-drain aborts through the existing loss path.
1903
1904    #[test]
1905    fn a_drain_keeps_the_tenancy_commit_eligible_until_the_final_commit() {
1906        // The inverse of `late_drain_commit_after_loss_is_skipped`: a lane
1907        // put into `Draining` is *still* commit-eligible, so its acked
1908        // watermarks keep folding to the store right up to the final commit.
1909        let script = Script::default();
1910        let mut d = driver(&script);
1911        let mut s = TestSource {
1912            accept_revoke: HashSet::from(["a".to_string()]),
1913            ..TestSource::default()
1914        };
1915        script.push(vec![gained("a", 1, None)]);
1916        poll(&mut d, &mut s);
1917        let partition = s.opened[0].3;
1918
1919        // The request lands and the source accepts, but the drain is not
1920        // finished (`drain_ready` returns None), so the tenancy stays.
1921        script.push(vec![CoordinationEvent::RevokeRequested {
1922            split: SplitId::new("a").unwrap(),
1923        }]);
1924        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
1925        assert_eq!(
1926            s.begin_revoke_calls,
1927            ["a"],
1928            "the source was asked to stop intake"
1929        );
1930
1931        // A tick commit for the draining tenancy must still fold.
1932        d.commit(&mut s, &[(partition, 42)]).unwrap();
1933        assert_eq!(s.encoded, vec![("a".to_string(), 42)]);
1934        assert_eq!(script.commits().len(), 1);
1935        assert_eq!(script.commits()[0].0.as_str(), "a");
1936        assert!(
1937            !script.commits()[0].1.completed,
1938            "a revocation never completes the split"
1939        );
1940        assert!(
1941            script.released_drained().is_empty(),
1942            "not released while the drain is still in flight"
1943        );
1944    }
1945
1946    #[test]
1947    fn a_completed_drain_releases_exactly_one_split_and_retires_barrierless() {
1948        let script = Script::default();
1949        let mut d = driver(&script);
1950        let mut s = TestSource {
1951            accept_revoke: HashSet::from(["a".to_string()]),
1952            ..TestSource::default()
1953        };
1954        script.push(vec![gained("a", 1, None), gained("b", 1, None)]);
1955        poll(&mut d, &mut s);
1956
1957        // a's drain has finished (tail acked and committed), so
1958        // `drain_ready` offers the final (non-terminal) progress.
1959        s.ready_progress
1960            .borrow_mut()
1961            .insert("a".into(), SplitProgress::new(50, vec![]));
1962        script.push(vec![CoordinationEvent::RevokeRequested {
1963            split: SplitId::new("a").unwrap(),
1964        }]);
1965
1966        // One poll carries the whole grant: accept, chase the tail, final
1967        // commit, release, and the lane leaves barrier-less.
1968        let event = poll(&mut d, &mut s);
1969        let SourceEvent::LanesRetired { lanes } = event else {
1970            panic!("a cooperative revocation must retire barrier-less, got {event:?}");
1971        };
1972        assert_eq!(lanes, vec![LaneId(0)], "only a's lane leaves");
1973
1974        assert_eq!(
1975            script.released_drained(),
1976            vec![SplitId::new("a").unwrap()],
1977            "exactly one split, released via the revocation path"
1978        );
1979        assert!(script.released().is_empty(), "not a plain hand-back");
1980        // Its final commit is not a completion.
1981        let last = script.commits().last().cloned().unwrap();
1982        assert_eq!(last.0.as_str(), "a");
1983        assert_eq!(last.1.watermark, 50);
1984        assert!(!last.1.completed, "drain commits never complete the split");
1985
1986        // b's live lane is untouched, and no revoke ever follows.
1987        assert_eq!(s.closed, vec!["a"], "b's fetcher stays attached");
1988        assert!(
1989            d.assignments().iter().any(|(id, _)| id.as_str() == "b"),
1990            "b is still owned"
1991        );
1992        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
1993    }
1994
1995    #[test]
1996    fn a_fenced_final_commit_aborts_the_drain_into_a_revoke() {
1997        let script = Script::default();
1998        let mut d = driver(&script);
1999        let mut s = TestSource {
2000            accept_revoke: HashSet::from(["a".to_string()]),
2001            ..TestSource::default()
2002        };
2003        script.push(vec![gained("a", 1, None)]);
2004        poll(&mut d, &mut s);
2005        let partition = s.opened[0].3;
2006
2007        // The drain is ready, but a peer fenced this tenancy first,
2008        // so its final commit is rejected.
2009        s.ready_progress
2010            .borrow_mut()
2011            .insert("a".into(), SplitProgress::new(50, vec![]));
2012        script.fail_next_commit("a", CoordinationErrorKind::Fenced);
2013        script.push(vec![CoordinationEvent::RevokeRequested {
2014            split: SplitId::new("a").unwrap(),
2015        }]);
2016
2017        let event = poll(&mut d, &mut s);
2018        let SourceEvent::LanesRevoked { lanes, barrier } = event else {
2019            panic!("a fenced drain must abort into a revoke, got {event:?}");
2020        };
2021        assert_eq!(lanes, vec![LaneId(0)]);
2022        assert_eq!(barrier.remaining(), 1, "one party per revoked lane");
2023        assert!(
2024            script.released_drained().is_empty(),
2025            "a fenced drain never releases"
2026        );
2027        assert!(
2028            script.commits().is_empty(),
2029            "the fenced final commit wrote nothing"
2030        );
2031        assert_eq!(s.closed, vec!["a"], "the fetcher was detached on the abort");
2032
2033        // A late watermark for the now-retired-fenced tenancy is skipped.
2034        d.commit(&mut s, &[(partition, 60)]).unwrap();
2035        assert!(
2036            s.encoded.is_empty(),
2037            "a retired-fenced tenancy must not encode"
2038        );
2039    }
2040
2041    #[test]
2042    fn a_retryable_final_commit_keeps_the_drain_pending() {
2043        let script = Script::default();
2044        let mut d = driver(&script);
2045        let mut s = TestSource {
2046            accept_revoke: HashSet::from(["a".to_string()]),
2047            ..TestSource::default()
2048        };
2049        script.push(vec![gained("a", 1, None)]);
2050        poll(&mut d, &mut s);
2051
2052        // The drain is ready, but the store defers the first final commit.
2053        s.ready_progress
2054            .borrow_mut()
2055            .insert("a".into(), SplitProgress::new(50, vec![]));
2056        script.fail_next_commit("a", CoordinationErrorKind::Retryable);
2057        script.push(vec![CoordinationEvent::RevokeRequested {
2058            split: SplitId::new("a").unwrap(),
2059        }]);
2060
2061        // Deferred: nothing written, nothing released, the split stays owned
2062        // and draining.
2063        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2064        assert!(script.commits().is_empty(), "deferred, not written");
2065        assert!(script.released_drained().is_empty());
2066        assert!(
2067            d.assignments().iter().any(|(id, _)| id.as_str() == "a"),
2068            "still owned while the final commit retries"
2069        );
2070
2071        // The next poll re-offers the same tail and the commit lands:
2072        // released and retired barrier-less.
2073        let event = poll(&mut d, &mut s);
2074        let SourceEvent::LanesRetired { lanes } = event else {
2075            panic!("the retried drain must finally retire, got {event:?}");
2076        };
2077        assert_eq!(lanes, vec![LaneId(0)]);
2078        assert_eq!(script.commits().len(), 1);
2079        assert_eq!(script.commits()[0].1.watermark, 50);
2080        assert_eq!(script.released_drained(), vec![SplitId::new("a").unwrap()]);
2081    }
2082
2083    #[test]
2084    fn a_source_that_cannot_stop_intake_declines_the_revoke() {
2085        let script = Script::default();
2086        let mut d = driver(&script);
2087        // Default `accept_revoke` is empty, so `begin_revoke` declines.
2088        let mut s = TestSource::default();
2089        script.push(vec![gained("a", 1, None)]);
2090        poll(&mut d, &mut s);
2091        let partition = s.opened[0].3;
2092
2093        script.push(vec![CoordinationEvent::RevokeRequested {
2094            split: SplitId::new("a").unwrap(),
2095        }]);
2096        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2097
2098        // Asked and declined: the tenancy never enters the drain (no
2099        // `drain_ready` poll) and stays fully live.
2100        assert_eq!(s.begin_revoke_calls, ["a"]);
2101        assert!(
2102            s.drain_ready_calls.is_empty(),
2103            "a declined split never drains"
2104        );
2105        assert!(script.released_drained().is_empty());
2106
2107        // A live tenancy keeps committing as normal.
2108        d.commit(&mut s, &[(partition, 30)]).unwrap();
2109        assert_eq!(s.encoded, vec![("a".to_string(), 30)]);
2110        assert_eq!(script.commits().len(), 1);
2111    }
2112
2113    #[test]
2114    fn a_revoke_request_for_an_unheld_split_is_ignored() {
2115        let script = Script::default();
2116        let mut d = driver(&script);
2117        let mut s = TestSource {
2118            // Even a source that *would* accept is never consulted for a
2119            // split this instance does not hold.
2120            accept_revoke: HashSet::from(["a".to_string(), "ghost".to_string()]),
2121            ..TestSource::default()
2122        };
2123        script.push(vec![gained("a", 1, None)]);
2124        poll(&mut d, &mut s);
2125
2126        script.push(vec![CoordinationEvent::RevokeRequested {
2127            split: SplitId::new("ghost").unwrap(),
2128        }]);
2129        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2130        assert!(
2131            s.begin_revoke_calls.is_empty(),
2132            "an unheld split must not consult the source"
2133        );
2134        assert!(
2135            d.assignments().iter().any(|(id, _)| id.as_str() == "a"),
2136            "the held split is untouched"
2137        );
2138    }
2139
2140    #[test]
2141    fn a_source_that_declines_feeds_the_decline_back() {
2142        // Distinct from `a_source_that_cannot_stop_intake_declines_the_revoke`:
2143        // this one proves the refusal reaches the backend, which then forces
2144        // the release.
2145        let script = Script::default();
2146        let mut d = driver(&script);
2147        // Default `accept_revoke` is empty, so `begin_revoke` declines.
2148        let mut s = TestSource::default();
2149        script.push(vec![gained("a", 1, None)]);
2150        poll(&mut d, &mut s);
2151        let partition = s.opened[0].3;
2152
2153        script.push(vec![CoordinationEvent::RevokeRequested {
2154            split: SplitId::new("a").unwrap(),
2155        }]);
2156        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2157
2158        // Asked, refused, and the refusal handed back to the backend exactly
2159        // once, naming the split, so the backend cools down that split only.
2160        assert_eq!(s.begin_revoke_calls, ["a"]);
2161        assert_eq!(
2162            script.declined(),
2163            vec![SplitId::new("a").unwrap()],
2164            "the source's refusal must reach the backend, once"
2165        );
2166
2167        // The tenancy stayed Live: commits still flow.
2168        d.commit(&mut s, &[(partition, 30)]).unwrap();
2169        assert_eq!(s.encoded, vec![("a".to_string(), 30)]);
2170        assert_eq!(script.commits().len(), 1);
2171        assert_eq!(script.commits()[0].0.as_str(), "a");
2172    }
2173
2174    #[test]
2175    fn a_repeated_revoke_request_mid_drain_is_not_declined() {
2176        // Re-emission is reachable: the backend cancels a revocation the
2177        // leader takes back, so a leader can drop a split, restore it, and
2178        // drop it again. A decline would force the release of a split
2179        // draining fine, costing the replay the cooperative path avoids.
2180        let script = Script::default();
2181        let mut d = driver(&script);
2182        let mut s = TestSource {
2183            accept_revoke: HashSet::from(["a".to_string()]),
2184            ..TestSource::default()
2185        };
2186        script.push(vec![gained("a", 1, None)]);
2187        poll(&mut d, &mut s);
2188
2189        // First request: accepted, and the drain does not finish
2190        // (`drain_ready` reports None), so the tenancy stays `Draining`.
2191        script.push(vec![CoordinationEvent::RevokeRequested {
2192            split: SplitId::new("a").unwrap(),
2193        }]);
2194        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2195        assert_eq!(s.begin_revoke_calls, ["a"]);
2196
2197        // Second request for the same, still-draining split.
2198        script.push(vec![CoordinationEvent::RevokeRequested {
2199            split: SplitId::new("a").unwrap(),
2200        }]);
2201        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2202        assert_eq!(
2203            s.begin_revoke_calls,
2204            ["a"],
2205            "the source must not be asked to stop intake it has already stopped"
2206        );
2207        assert!(
2208            script.declined().is_empty(),
2209            "a drain already in flight satisfies the request; declining it would force the release"
2210        );
2211
2212        // And the drain still completes on its own terms.
2213        s.ready_progress
2214            .borrow_mut()
2215            .insert("a".into(), SplitProgress::new(50, vec![]));
2216        let event = poll(&mut d, &mut s);
2217        let SourceEvent::LanesRetired { lanes } = event else {
2218            panic!("the drain must still retire, got {event:?}");
2219        };
2220        assert_eq!(lanes, vec![LaneId(0)]);
2221        assert_eq!(script.released_drained(), vec![SplitId::new("a").unwrap()]);
2222    }
2223
2224    #[test]
2225    fn an_unopened_tenancy_declines_without_asking_the_source() {
2226        // A gain and a revocation request for the same split arrive in one
2227        // event batch: the tenancy exists but its lane has not opened yet.
2228        // The driver declines without consulting the source, and the split
2229        // still opens normally afterwards.
2230        let script = Script::default();
2231        let mut d = driver(&script);
2232        // Even a source that WOULD accept must not be consulted before the
2233        // lane exists.
2234        let mut s = TestSource {
2235            accept_revoke: HashSet::from(["a".to_string()]),
2236            ..TestSource::default()
2237        };
2238
2239        script.push(vec![
2240            gained("a", 1, None),
2241            CoordinationEvent::RevokeRequested {
2242                split: SplitId::new("a").unwrap(),
2243            },
2244        ]);
2245        // Both events are applied (decline included) and then the lane opens,
2246        // all on this one call.
2247        let event = poll(&mut d, &mut s);
2248        let SourceEvent::LanesAdded(lanes) = event else {
2249            panic!("the split must still open after the early decline, got {event:?}");
2250        };
2251        assert_eq!(lanes.len(), 1);
2252
2253        assert!(
2254            s.begin_revoke_calls.is_empty(),
2255            "a not-yet-opened tenancy must never be asked to stop intake"
2256        );
2257        assert_eq!(
2258            script.declined(),
2259            vec![SplitId::new("a").unwrap()],
2260            "the premature request is declined back to the backend"
2261        );
2262
2263        // The split is fully live now: commits flow.
2264        let partition = s.opened[0].3;
2265        d.commit(&mut s, &[(partition, 25)]).unwrap();
2266        assert_eq!(s.encoded, vec![("a".to_string(), 25)]);
2267        assert_eq!(script.commits().len(), 1);
2268        // And it was never released as drained.
2269        assert!(script.released_drained().is_empty());
2270    }
2271
2272    #[test]
2273    fn a_completed_progress_during_a_drain_is_never_terminal() {
2274        // The central guard must strip a `completed` flag reported during a
2275        // drain while still landing the commit (the watermark is acked).
2276        let script = Script::default();
2277        let mut d = driver(&script);
2278        let mut s = TestSource {
2279            accept_revoke: HashSet::from(["a".to_string()]),
2280            ..TestSource::default()
2281        };
2282        script.push(vec![gained("a", 1, None)]);
2283        poll(&mut d, &mut s);
2284        let partition = s.opened[0].3;
2285
2286        // Drive "a" into `Draining` (the drain is not finished yet:
2287        // `drain_ready` returns None).
2288        script.push(vec![CoordinationEvent::RevokeRequested {
2289            split: SplitId::new("a").unwrap(),
2290        }]);
2291        assert!(matches!(poll(&mut d, &mut s), SourceEvent::Idle));
2292        assert_eq!(s.begin_revoke_calls, ["a"]);
2293
2294        // A tick commit whose `encode_commit` reports the split COMPLETE at
2295        // this watermark.
2296        s.complete_at.insert("a".into(), 42);
2297        d.commit(&mut s, &[(partition, 42)]).unwrap();
2298
2299        // The commit landed (watermark advanced) but stripped of completion.
2300        let committed = script.commits().last().cloned().expect("the commit landed");
2301        assert_eq!(committed.0.as_str(), "a");
2302        assert_eq!(
2303            committed.1.watermark, 42,
2304            "the guard strips the flag, not the commit"
2305        );
2306        assert!(
2307            !committed.1.completed,
2308            "a drain cut is never terminal, whatever the source claims"
2309        );
2310
2311        // The tenancy is neither completed nor retired: still owned, still
2312        // draining, no lane has left.
2313        assert!(
2314            d.assignments().iter().any(|(id, _)| id.as_str() == "a"),
2315            "still owned"
2316        );
2317        assert!(s.closed.is_empty(), "not retired");
2318
2319        // The revocation then finishes normally once the drain completes.
2320        s.ready_progress
2321            .borrow_mut()
2322            .insert("a".into(), SplitProgress::new(50, vec![]));
2323        let event = poll(&mut d, &mut s);
2324        let SourceEvent::LanesRetired { lanes } = event else {
2325            panic!("the drained revocation must finally retire, got {event:?}");
2326        };
2327        assert_eq!(lanes, vec![LaneId(0)]);
2328        assert_eq!(script.released_drained(), vec![SplitId::new("a").unwrap()]);
2329        assert!(
2330            !script.commits().last().unwrap().1.completed,
2331            "the final revocation commit is not terminal either"
2332        );
2333    }
2334}