Skip to main content

kvbm_engine/offload/
settlement.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::fmt;
5use std::sync::Arc;
6
7use futures::future::select_all;
8use parking_lot::Mutex;
9use tokio::sync::{OwnedSemaphorePermit, watch};
10use uuid::Uuid;
11
12/// An offload pipeline whose executor can participate in causal settlement.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum PipelineLane {
15    G1ToG2,
16    G2ToG3,
17    G2ToG4,
18}
19
20impl PipelineLane {
21    pub(crate) const ALL: [Self; 3] = [Self::G1ToG2, Self::G2ToG3, Self::G2ToG4];
22
23    pub(crate) const fn index(self) -> usize {
24        match self {
25            Self::G1ToG2 => 0,
26            Self::G2ToG3 => 1,
27            Self::G2ToG4 => 2,
28        }
29    }
30}
31
32impl fmt::Display for PipelineLane {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            Self::G1ToG2 => f.write_str("G1→G2"),
36            Self::G2ToG3 => f.write_str("G2→G3"),
37            Self::G2ToG4 => f.write_str("G2→G4"),
38        }
39    }
40}
41
42/// Terminal executor failure observed while waiting for causal settlement.
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct PipelineFailure {
45    kind: PipelineFailureKind,
46    message: String,
47}
48
49impl PipelineFailure {
50    pub(crate) fn new(kind: PipelineFailureKind, message: impl Into<String>) -> Self {
51        Self {
52            kind,
53            message: message.into(),
54        }
55    }
56
57    /// Failure category.
58    pub fn kind(&self) -> PipelineFailureKind {
59        self.kind
60    }
61
62    /// Diagnostic detail supplied by the failing executor.
63    pub fn message(&self) -> &str {
64        &self.message
65    }
66}
67
68impl fmt::Display for PipelineFailure {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        write!(f, "{}: {}", self.kind, self.message)
71    }
72}
73
74/// Stable failure categories for pipeline settlement.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum PipelineFailureKind {
77    TaskAborted,
78    TaskPanicked,
79    Executor,
80    Shutdown,
81    CounterOverflow,
82    Invariant,
83}
84
85impl fmt::Display for PipelineFailureKind {
86    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87        match self {
88            Self::TaskAborted => f.write_str("task aborted"),
89            Self::TaskPanicked => f.write_str("task panicked"),
90            Self::Executor => f.write_str("executor failure"),
91            Self::Shutdown => f.write_str("pipeline shutdown"),
92            Self::CounterOverflow => f.write_str("counter overflow"),
93            Self::Invariant => f.write_str("pipeline invariant violation"),
94        }
95    }
96}
97
98/// Opaque checkpoint captured before an external completion source is fired.
99#[derive(Debug, Clone)]
100pub struct SettlementToken {
101    pub(crate) engine_id: Uuid,
102    pub(crate) checkpoints: [Option<PipelineCheckpoint>; 3],
103}
104
105impl SettlementToken {
106    pub(crate) fn validate_engine(&self, engine_id: Uuid) -> Result<(), SettlementError> {
107        if self.engine_id == engine_id {
108            Ok(())
109        } else {
110            Err(SettlementError::ForeignToken)
111        }
112    }
113
114    pub(crate) fn checkpoint(
115        &self,
116        lane: PipelineLane,
117    ) -> Result<PipelineCheckpoint, SettlementError> {
118        self.checkpoints[lane.index()].ok_or(SettlementError::LaneUnavailable { lane })
119    }
120}
121
122/// Expected completed transfer batches for each pipeline lane.
123#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
124pub struct SettlementTarget {
125    completed_batches: [u64; 3],
126}
127
128impl SettlementTarget {
129    /// Create an empty target. Empty targets settle immediately.
130    pub fn new() -> Self {
131        Self::default()
132    }
133
134    /// Require `count` completed batches on `lane` after the checkpoint.
135    pub fn with_completed_batches(mut self, lane: PipelineLane, count: u64) -> Self {
136        self.completed_batches[lane.index()] = count;
137        self
138    }
139
140    /// Add completed batches to a lane's requirement.
141    pub fn add_completed_batches(
142        &mut self,
143        lane: PipelineLane,
144        count: u64,
145    ) -> Result<(), SettlementError> {
146        let value = &mut self.completed_batches[lane.index()];
147        *value = value
148            .checked_add(count)
149            .ok_or(SettlementError::TargetOverflow { lane })?;
150        Ok(())
151    }
152
153    pub(crate) fn completed_batches(self, lane: PipelineLane) -> u64 {
154        self.completed_batches[lane.index()]
155    }
156
157    pub(crate) fn is_empty(self) -> bool {
158        self.completed_batches.iter().all(|count| *count == 0)
159    }
160}
161
162/// Error returned when a causal settlement boundary cannot be established.
163#[derive(Debug, Clone, PartialEq, Eq)]
164pub enum SettlementError {
165    ForeignToken,
166    LaneUnavailable {
167        lane: PipelineLane,
168    },
169    UnsupportedAutoChain {
170        lane: PipelineLane,
171    },
172    StaleToken {
173        lane: PipelineLane,
174    },
175    TargetOverflow {
176        lane: PipelineLane,
177    },
178    ProgressClosed {
179        lane: PipelineLane,
180    },
181    PipelineFailed {
182        lane: PipelineLane,
183        failure: PipelineFailure,
184    },
185}
186
187impl fmt::Display for SettlementError {
188    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
189        match self {
190            Self::ForeignToken => f.write_str("settlement token belongs to a different engine"),
191            Self::LaneUnavailable { lane } => write!(f, "pipeline lane {lane} is not configured"),
192            Self::UnsupportedAutoChain { lane } => {
193                write!(f, "pipeline lane {lane} uses unsupported auto-chaining")
194            }
195            Self::StaleToken { lane } => {
196                write!(f, "settlement token is stale for pipeline lane {lane}")
197            }
198            Self::TargetOverflow { lane } => {
199                write!(f, "settlement target overflowed for lane {lane}")
200            }
201            Self::ProgressClosed { lane } => {
202                write!(f, "pipeline progress channel closed for lane {lane}")
203            }
204            Self::PipelineFailed { lane, failure } => {
205                write!(f, "pipeline lane {lane} failed: {failure}")
206            }
207        }
208    }
209}
210
211impl std::error::Error for SettlementError {}
212
213#[derive(Debug, Clone, Copy)]
214pub(crate) struct PipelineCheckpoint {
215    pub(crate) epoch: u64,
216    pub(crate) completed_batches: u64,
217}
218
219#[derive(Debug, Clone, Default, PartialEq, Eq)]
220pub(crate) struct PipelineSnapshot {
221    pub(crate) epoch: u64,
222    pub(crate) queued: usize,
223    pub(crate) starting: usize,
224    pub(crate) in_flight: usize,
225    pub(crate) settling: usize,
226    pub(crate) completed_batches: u64,
227    pub(crate) failure: Option<PipelineFailure>,
228}
229
230impl PipelineSnapshot {
231    pub(crate) fn checkpoint(&self) -> PipelineCheckpoint {
232        PipelineCheckpoint {
233            epoch: self.epoch,
234            completed_batches: self.completed_batches,
235        }
236    }
237
238    pub(crate) fn is_stable(&self, max_concurrent_transfers: usize) -> bool {
239        if self.starting > 0 || self.settling > 0 {
240            return false;
241        }
242        self.queued == 0 || self.in_flight >= max_concurrent_transfers
243    }
244}
245
246pub(crate) struct SettlementWaiter {
247    lane: PipelineLane,
248    checkpoint: PipelineCheckpoint,
249    expected_completed: u64,
250    max_concurrent_transfers: usize,
251    progress: watch::Receiver<PipelineSnapshot>,
252}
253
254impl SettlementWaiter {
255    pub(crate) fn new(
256        lane: PipelineLane,
257        tracker: &PipelineSettlementTracker,
258        checkpoint: PipelineCheckpoint,
259        delta: u64,
260    ) -> Result<Self, SettlementError> {
261        let expected_completed = checkpoint
262            .completed_batches
263            .checked_add(delta)
264            .ok_or(SettlementError::TargetOverflow { lane })?;
265        Ok(Self {
266            lane,
267            checkpoint,
268            expected_completed,
269            max_concurrent_transfers: tracker.max_concurrent_transfers(),
270            progress: tracker.subscribe(),
271        })
272    }
273}
274
275pub(crate) async fn wait_for_settlement(
276    mut waiters: Vec<SettlementWaiter>,
277) -> Result<(), SettlementError> {
278    loop {
279        let snapshots: Vec<PipelineSnapshot> = waiters
280            .iter_mut()
281            .map(|waiter| waiter.progress.borrow_and_update().clone())
282            .collect();
283
284        let mut all_settled = true;
285        for (waiter, snapshot) in waiters.iter().zip(&snapshots) {
286            if snapshot.epoch < waiter.checkpoint.epoch
287                || snapshot.completed_batches < waiter.checkpoint.completed_batches
288            {
289                return Err(SettlementError::StaleToken { lane: waiter.lane });
290            }
291            if let Some(failure) = &snapshot.failure {
292                return Err(SettlementError::PipelineFailed {
293                    lane: waiter.lane,
294                    failure: failure.clone(),
295                });
296            }
297            all_settled &= snapshot.completed_batches >= waiter.expected_completed
298                && snapshot.is_stable(waiter.max_concurrent_transfers);
299        }
300
301        let epoch_changed = waiters
302            .iter()
303            .zip(&snapshots)
304            .any(|(waiter, snapshot)| waiter.progress.borrow().epoch != snapshot.epoch);
305        if epoch_changed {
306            continue;
307        }
308        if all_settled {
309            return Ok(());
310        }
311
312        let changes = waiters
313            .iter_mut()
314            .map(|waiter| Box::pin(waiter.progress.changed()))
315            .collect::<Vec<_>>();
316        let (result, changed_index, _) = select_all(changes).await;
317        if result.is_err() {
318            return Err(SettlementError::ProgressClosed {
319                lane: waiters[changed_index].lane,
320            });
321        }
322    }
323}
324
325struct TrackerInner {
326    snapshot: Mutex<PipelineSnapshot>,
327    progress_tx: watch::Sender<PipelineSnapshot>,
328    max_concurrent_transfers: usize,
329}
330
331/// One linearization point for all executor state used by settlement.
332#[derive(Clone)]
333pub(crate) struct PipelineSettlementTracker {
334    inner: Arc<TrackerInner>,
335}
336
337impl PipelineSettlementTracker {
338    pub(crate) fn new(max_concurrent_transfers: usize) -> Self {
339        let snapshot = PipelineSnapshot::default();
340        let (progress_tx, _progress_rx) = watch::channel(snapshot.clone());
341        Self {
342            inner: Arc::new(TrackerInner {
343                snapshot: Mutex::new(snapshot),
344                progress_tx,
345                max_concurrent_transfers: max_concurrent_transfers.max(1),
346            }),
347        }
348    }
349
350    pub(crate) fn snapshot(&self) -> PipelineSnapshot {
351        self.inner.snapshot.lock().clone()
352    }
353
354    pub(crate) fn subscribe(&self) -> watch::Receiver<PipelineSnapshot> {
355        self.inner.progress_tx.subscribe()
356    }
357
358    pub(crate) fn max_concurrent_transfers(&self) -> usize {
359        self.inner.max_concurrent_transfers
360    }
361
362    pub(crate) fn queue_batch(&self) {
363        self.update(|snapshot| {
364            increment_usize(&mut snapshot.queued, "queued")?;
365            Ok(())
366        });
367    }
368
369    pub(crate) fn discard_queued(&self) {
370        self.update(|snapshot| decrement_usize(&mut snapshot.queued, "queued"));
371    }
372
373    pub(crate) fn fail_queued(&self, failure: PipelineFailure) {
374        self.update(|snapshot| {
375            decrement_usize(&mut snapshot.queued, "queued")?;
376            record_failure(snapshot, failure);
377            Ok(())
378        });
379    }
380
381    pub(crate) fn admit(&self) {
382        self.update(|snapshot| {
383            decrement_usize(&mut snapshot.queued, "queued")?;
384            increment_usize(&mut snapshot.starting, "starting")?;
385            Ok(())
386        });
387    }
388
389    fn transition(&self, from: BatchPhase, to: BatchPhase) {
390        self.update(|snapshot| {
391            decrement_phase(snapshot, from)?;
392            increment_phase(snapshot, to)?;
393            Ok(())
394        });
395    }
396
397    fn complete(&self, phase: BatchPhase) {
398        self.update(|snapshot| {
399            decrement_phase(snapshot, phase)?;
400            snapshot.completed_batches =
401                snapshot.completed_batches.checked_add(1).ok_or_else(|| {
402                    PipelineFailure::new(
403                        PipelineFailureKind::CounterOverflow,
404                        "completed_batches overflowed",
405                    )
406                })?;
407            Ok(())
408        });
409    }
410
411    fn fail_phase(&self, phase: BatchPhase, failure: PipelineFailure) {
412        self.update(|snapshot| {
413            decrement_phase(snapshot, phase)?;
414            record_failure(snapshot, failure);
415            Ok(())
416        });
417    }
418
419    pub(crate) fn fail(&self, failure: PipelineFailure) {
420        self.update(|snapshot| {
421            record_failure(snapshot, failure);
422            Ok(())
423        });
424    }
425
426    fn update(&self, update: impl FnOnce(&mut PipelineSnapshot) -> Result<(), PipelineFailure>) {
427        let mut snapshot = self.inner.snapshot.lock();
428        if let Err(failure) = update(&mut snapshot) {
429            record_failure(&mut snapshot, failure);
430        }
431        if let Some(next_epoch) = snapshot.epoch.checked_add(1) {
432            snapshot.epoch = next_epoch;
433        } else if snapshot.failure.is_none() {
434            snapshot.failure = Some(PipelineFailure::new(
435                PipelineFailureKind::CounterOverflow,
436                "pipeline epoch overflowed",
437            ));
438        }
439        self.inner.progress_tx.send_replace(snapshot.clone());
440    }
441}
442
443pub(crate) struct QueuedBatchGuard {
444    tracker: PipelineSettlementTracker,
445    queued: bool,
446}
447
448impl QueuedBatchGuard {
449    pub(crate) fn new(tracker: PipelineSettlementTracker) -> Self {
450        tracker.queue_batch();
451        Self {
452            tracker,
453            queued: true,
454        }
455    }
456
457    pub(crate) fn sent(mut self) {
458        self.queued = false;
459    }
460
461    pub(crate) fn finish_failure(mut self, failure: PipelineFailure) {
462        self.queued = false;
463        self.tracker.fail_queued(failure);
464    }
465}
466
467impl Drop for QueuedBatchGuard {
468    fn drop(&mut self) {
469        if !self.queued {
470            return;
471        }
472        let (kind, message) = if std::thread::panicking() {
473            (
474                PipelineFailureKind::TaskPanicked,
475                "precondition task panicked while forwarding a queued batch",
476            )
477        } else {
478            (
479                PipelineFailureKind::TaskAborted,
480                "precondition task ended while forwarding a queued batch",
481            )
482        };
483        self.tracker
484            .fail_queued(PipelineFailure::new(kind, message));
485    }
486}
487
488pub(crate) struct PipelineRunGuard {
489    tracker: PipelineSettlementTracker,
490    running: bool,
491    pipeline_name: &'static str,
492}
493
494impl PipelineRunGuard {
495    pub(crate) fn new(tracker: PipelineSettlementTracker, pipeline_name: &'static str) -> Self {
496        Self {
497            tracker,
498            running: true,
499            pipeline_name,
500        }
501    }
502
503    pub(crate) fn finish_shutdown(mut self) {
504        self.running = false;
505        self.tracker.fail(PipelineFailure::new(
506            PipelineFailureKind::Shutdown,
507            format!("{} stopped", self.pipeline_name),
508        ));
509    }
510}
511
512impl Drop for PipelineRunGuard {
513    fn drop(&mut self) {
514        if !self.running {
515            return;
516        }
517        let (kind, reason) = if std::thread::panicking() {
518            (PipelineFailureKind::TaskPanicked, "panicked")
519        } else {
520            (PipelineFailureKind::TaskAborted, "was aborted")
521        };
522        self.tracker.fail(PipelineFailure::new(
523            kind,
524            format!("{} {reason}", self.pipeline_name),
525        ));
526    }
527}
528
529#[derive(Debug, Clone, Copy)]
530enum BatchPhase {
531    Starting,
532    InFlight,
533    Settling,
534}
535
536pub(crate) struct BatchPhaseGuard {
537    tracker: PipelineSettlementTracker,
538    phase: Option<BatchPhase>,
539    permit: Option<OwnedSemaphorePermit>,
540}
541
542impl BatchPhaseGuard {
543    pub(crate) fn starting(
544        tracker: PipelineSettlementTracker,
545        permit: OwnedSemaphorePermit,
546    ) -> Self {
547        tracker.admit();
548        Self {
549            tracker,
550            phase: Some(BatchPhase::Starting),
551            permit: Some(permit),
552        }
553    }
554
555    pub(crate) fn mark_in_flight(&mut self) {
556        self.transition(BatchPhase::InFlight);
557    }
558
559    pub(crate) fn mark_settling(&mut self) {
560        self.transition(BatchPhase::Settling);
561    }
562
563    pub(crate) fn finish_success(mut self) {
564        let phase = self
565            .phase
566            .take()
567            .expect("batch phase guard already finished");
568        drop(self.permit.take());
569        self.tracker.complete(phase);
570    }
571
572    pub(crate) fn finish_failure(mut self, failure: PipelineFailure) {
573        let phase = self
574            .phase
575            .take()
576            .expect("batch phase guard already finished");
577        drop(self.permit.take());
578        self.tracker.fail_phase(phase, failure);
579    }
580
581    fn transition(&mut self, next: BatchPhase) {
582        let current = self.phase.expect("batch phase guard already finished");
583        self.tracker.transition(current, next);
584        self.phase = Some(next);
585    }
586}
587
588impl Drop for BatchPhaseGuard {
589    fn drop(&mut self) {
590        let Some(phase) = self.phase.take() else {
591            return;
592        };
593        drop(self.permit.take());
594        let (kind, message) = if std::thread::panicking() {
595            (
596                PipelineFailureKind::TaskPanicked,
597                "transfer task panicked before publishing settlement",
598            )
599        } else {
600            (
601                PipelineFailureKind::TaskAborted,
602                "transfer task ended before publishing settlement",
603            )
604        };
605        self.tracker
606            .fail_phase(phase, PipelineFailure::new(kind, message));
607    }
608}
609
610fn increment_phase(
611    snapshot: &mut PipelineSnapshot,
612    phase: BatchPhase,
613) -> Result<(), PipelineFailure> {
614    match phase {
615        BatchPhase::Starting => increment_usize(&mut snapshot.starting, "starting"),
616        BatchPhase::InFlight => increment_usize(&mut snapshot.in_flight, "in_flight"),
617        BatchPhase::Settling => increment_usize(&mut snapshot.settling, "settling"),
618    }
619}
620
621fn decrement_phase(
622    snapshot: &mut PipelineSnapshot,
623    phase: BatchPhase,
624) -> Result<(), PipelineFailure> {
625    match phase {
626        BatchPhase::Starting => decrement_usize(&mut snapshot.starting, "starting"),
627        BatchPhase::InFlight => decrement_usize(&mut snapshot.in_flight, "in_flight"),
628        BatchPhase::Settling => decrement_usize(&mut snapshot.settling, "settling"),
629    }
630}
631
632fn increment_usize(value: &mut usize, name: &str) -> Result<(), PipelineFailure> {
633    *value = value.checked_add(1).ok_or_else(|| {
634        PipelineFailure::new(
635            PipelineFailureKind::CounterOverflow,
636            format!("{name} counter overflowed"),
637        )
638    })?;
639    Ok(())
640}
641
642fn decrement_usize(value: &mut usize, name: &str) -> Result<(), PipelineFailure> {
643    *value = value.checked_sub(1).ok_or_else(|| {
644        PipelineFailure::new(
645            PipelineFailureKind::Invariant,
646            format!("{name} counter underflowed"),
647        )
648    })?;
649    Ok(())
650}
651
652fn record_failure(snapshot: &mut PipelineSnapshot, failure: PipelineFailure) {
653    if snapshot.failure.is_none() {
654        snapshot.failure = Some(failure);
655    }
656}
657
658#[cfg(test)]
659mod tests {
660    use std::panic::{AssertUnwindSafe, catch_unwind};
661
662    use futures::{FutureExt, poll};
663    use tokio::sync::{Semaphore, oneshot};
664
665    use super::*;
666
667    async fn starting_phase(
668        tracker: &PipelineSettlementTracker,
669        semaphore: &Arc<Semaphore>,
670    ) -> BatchPhaseGuard {
671        let permit = semaphore
672            .clone()
673            .acquire_owned()
674            .await
675            .expect("test semaphore should remain open");
676        BatchPhaseGuard::starting(tracker.clone(), permit)
677    }
678
679    fn waiter(
680        tracker: &PipelineSettlementTracker,
681        checkpoint: PipelineCheckpoint,
682        delta: u64,
683    ) -> SettlementWaiter {
684        SettlementWaiter::new(PipelineLane::G1ToG2, tracker, checkpoint, delta)
685            .expect("test settlement target should be valid")
686    }
687
688    #[tokio::test]
689    async fn concurrency_one_waits_for_queued_successor_to_invoke_worker() {
690        let tracker = PipelineSettlementTracker::new(1);
691        let semaphore = Arc::new(Semaphore::new(1));
692        let checkpoint = tracker.snapshot().checkpoint();
693        tracker.queue_batch();
694        tracker.queue_batch();
695
696        let mut first = starting_phase(&tracker, &semaphore).await;
697        first.mark_in_flight();
698        let mut settlement = wait_for_settlement(vec![waiter(&tracker, checkpoint, 1)]).boxed();
699
700        first.mark_settling();
701        first.finish_success();
702        assert!(poll!(&mut settlement).is_pending());
703
704        let mut successor = starting_phase(&tracker, &semaphore).await;
705        assert!(poll!(&mut settlement).is_pending());
706        successor.mark_in_flight();
707        assert_eq!(settlement.await, Ok(()));
708
709        successor.mark_settling();
710        successor.finish_success();
711    }
712
713    #[tokio::test]
714    async fn simultaneous_completions_fill_every_available_successor_slot() {
715        let tracker = PipelineSettlementTracker::new(2);
716        let semaphore = Arc::new(Semaphore::new(2));
717        let checkpoint = tracker.snapshot().checkpoint();
718        for _ in 0..4 {
719            tracker.queue_batch();
720        }
721
722        let mut first = starting_phase(&tracker, &semaphore).await;
723        let mut second = starting_phase(&tracker, &semaphore).await;
724        first.mark_in_flight();
725        second.mark_in_flight();
726        let mut settlement = wait_for_settlement(vec![waiter(&tracker, checkpoint, 2)]).boxed();
727
728        first.mark_settling();
729        second.mark_settling();
730        first.finish_success();
731        second.finish_success();
732        assert!(poll!(&mut settlement).is_pending());
733
734        let mut third = starting_phase(&tracker, &semaphore).await;
735        let mut fourth = starting_phase(&tracker, &semaphore).await;
736        third.mark_in_flight();
737        assert!(poll!(&mut settlement).is_pending());
738        fourth.mark_in_flight();
739        assert_eq!(settlement.await, Ok(()));
740
741        third.mark_settling();
742        fourth.mark_settling();
743        third.finish_success();
744        fourth.finish_success();
745    }
746
747    #[tokio::test]
748    async fn completion_without_successor_settles() {
749        let tracker = PipelineSettlementTracker::new(1);
750        let semaphore = Arc::new(Semaphore::new(1));
751        let checkpoint = tracker.snapshot().checkpoint();
752        tracker.queue_batch();
753        let mut phase = starting_phase(&tracker, &semaphore).await;
754        phase.mark_in_flight();
755        let settlement = wait_for_settlement(vec![waiter(&tracker, checkpoint, 1)]);
756
757        phase.mark_settling();
758        phase.finish_success();
759        assert_eq!(settlement.await, Ok(()));
760    }
761
762    #[tokio::test]
763    async fn progress_between_subscription_and_initial_check_is_observed() {
764        let tracker = PipelineSettlementTracker::new(1);
765        let semaphore = Arc::new(Semaphore::new(1));
766        let checkpoint = tracker.snapshot().checkpoint();
767        tracker.queue_batch();
768        let mut phase = starting_phase(&tracker, &semaphore).await;
769        phase.mark_in_flight();
770        let subscribed = waiter(&tracker, checkpoint, 1);
771
772        phase.mark_settling();
773        phase.finish_success();
774        assert_eq!(wait_for_settlement(vec![subscribed]).await, Ok(()));
775    }
776
777    #[tokio::test]
778    async fn task_abort_removes_phase_and_publishes_terminal_failure() {
779        let tracker = PipelineSettlementTracker::new(1);
780        tracker.queue_batch();
781        let task_tracker = tracker.clone();
782        let semaphore = Arc::new(Semaphore::new(1));
783        let task_semaphore = semaphore.clone();
784        let (started_tx, started_rx) = oneshot::channel();
785        let task = tokio::spawn(async move {
786            let mut phase = starting_phase(&task_tracker, &task_semaphore).await;
787            phase.mark_in_flight();
788            started_tx.send(()).expect("test should observe task start");
789            std::future::pending::<()>().await;
790            phase.mark_settling();
791        });
792        started_rx.await.expect("transfer task should start");
793
794        task.abort();
795        assert!(
796            task.await
797                .expect_err("task should be cancelled")
798                .is_cancelled()
799        );
800        let snapshot = tracker.snapshot();
801        assert_eq!(snapshot.in_flight, 0);
802        assert_eq!(
803            snapshot.failure.expect("abort must publish failure").kind(),
804            PipelineFailureKind::TaskAborted
805        );
806    }
807
808    #[tokio::test]
809    async fn cancellation_panic_executor_error_and_shutdown_are_structured() {
810        let queued_tracker = PipelineSettlementTracker::new(1);
811        drop(QueuedBatchGuard::new(queued_tracker.clone()));
812        assert_eq!(
813            queued_tracker
814                .snapshot()
815                .failure
816                .expect("queued cancellation must fail")
817                .kind(),
818            PipelineFailureKind::TaskAborted
819        );
820
821        let panic_tracker = PipelineSettlementTracker::new(1);
822        panic_tracker.queue_batch();
823        let panic_semaphore = Arc::new(Semaphore::new(1));
824        let phase = starting_phase(&panic_tracker, &panic_semaphore).await;
825        assert!(
826            catch_unwind(AssertUnwindSafe(move || {
827                let _phase = phase;
828                panic!("intentional transfer panic");
829            }))
830            .is_err()
831        );
832        assert_eq!(
833            panic_tracker
834                .snapshot()
835                .failure
836                .expect("panic must publish failure")
837                .kind(),
838            PipelineFailureKind::TaskPanicked
839        );
840
841        let error_tracker = PipelineSettlementTracker::new(1);
842        error_tracker.queue_batch();
843        let error_semaphore = Arc::new(Semaphore::new(1));
844        let phase = starting_phase(&error_tracker, &error_semaphore).await;
845        phase.finish_failure(PipelineFailure::new(
846            PipelineFailureKind::Executor,
847            "injected executor error",
848        ));
849        assert_eq!(
850            error_tracker
851                .snapshot()
852                .failure
853                .expect("executor error must publish failure")
854                .kind(),
855            PipelineFailureKind::Executor
856        );
857
858        let shutdown_tracker = PipelineSettlementTracker::new(1);
859        PipelineRunGuard::new(shutdown_tracker.clone(), "test pipeline").finish_shutdown();
860        assert_eq!(
861            shutdown_tracker
862                .snapshot()
863                .failure
864                .expect("shutdown must publish failure")
865                .kind(),
866            PipelineFailureKind::Shutdown
867        );
868    }
869
870    #[tokio::test]
871    async fn closed_progress_and_stale_checkpoint_return_errors() {
872        let tracker = PipelineSettlementTracker::new(1);
873        let checkpoint = tracker.snapshot().checkpoint();
874        let closed_waiter = waiter(&tracker, checkpoint, 1);
875        drop(tracker);
876        assert!(matches!(
877            wait_for_settlement(vec![closed_waiter]).await,
878            Err(SettlementError::ProgressClosed {
879                lane: PipelineLane::G1ToG2
880            })
881        ));
882
883        let tracker = PipelineSettlementTracker::new(1);
884        let stale = PipelineCheckpoint {
885            epoch: tracker.snapshot().epoch + 1,
886            completed_batches: 0,
887        };
888        assert!(matches!(
889            wait_for_settlement(vec![waiter(&tracker, stale, 0)]).await,
890            Err(SettlementError::StaleToken {
891                lane: PipelineLane::G1ToG2
892            })
893        ));
894    }
895
896    #[test]
897    fn foreign_unavailable_and_overflow_targets_return_errors() {
898        let engine_id = Uuid::new_v4();
899        let token = SettlementToken {
900            engine_id,
901            checkpoints: [None; 3],
902        };
903        assert_eq!(
904            token.validate_engine(Uuid::new_v4()),
905            Err(SettlementError::ForeignToken)
906        );
907        assert!(matches!(
908            token.checkpoint(PipelineLane::G2ToG3),
909            Err(SettlementError::LaneUnavailable {
910                lane: PipelineLane::G2ToG3
911            })
912        ));
913
914        let mut target = SettlementTarget::new();
915        target
916            .add_completed_batches(PipelineLane::G2ToG4, u64::MAX)
917            .expect("first target addition should fit");
918        assert!(matches!(
919            target.add_completed_batches(PipelineLane::G2ToG4, 1),
920            Err(SettlementError::TargetOverflow {
921                lane: PipelineLane::G2ToG4
922            })
923        ));
924
925        let tracker = PipelineSettlementTracker::new(1);
926        let checkpoint = PipelineCheckpoint {
927            epoch: 0,
928            completed_batches: u64::MAX,
929        };
930        assert!(matches!(
931            SettlementWaiter::new(PipelineLane::G1ToG2, &tracker, checkpoint, 1),
932            Err(SettlementError::TargetOverflow {
933                lane: PipelineLane::G1ToG2
934            })
935        ));
936    }
937}