1use 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#[derive(Debug)]
63#[non_exhaustive]
64pub struct SplitOpening<'a> {
65 pub split: &'a SplitSpec,
67 pub resume: Option<&'a SplitProgress>,
70 pub lane: LaneId,
73 pub partition: PartitionId,
76 pub epoch: LeaseEpoch,
78 pub waker: &'a ControlWaker,
83}
84
85pub trait SplitSource {
92 type Lane: SourceLane;
94
95 fn open_split(&mut self, opening: SplitOpening<'_>) -> Result<Self::Lane, SourceError>;
98
99 fn validate_resume(
111 &self,
112 split: &SplitSpec,
113 progress: &SplitProgress,
114 ) -> Result<(), SourceError> {
115 let _ = (split, progress);
116 Ok(())
117 }
118
119 fn encode_commit(
124 &mut self,
125 split: &SplitId,
126 watermark: i64,
127 ) -> Result<SplitProgress, SourceError>;
128
129 fn sweep(&mut self, split: &SplitId) -> Result<Option<SplitProgress>, SourceError>;
134
135 fn close_split(&mut self, split: &SplitId);
145
146 fn take_finishing(&mut self) -> Vec<SplitId> {
154 Vec::new()
155 }
156
157 fn begin_revoke(&mut self, split: &SplitId) -> bool {
188 let _ = split;
189 false
190 }
191
192 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 Live,
222 Draining,
230 Retired,
234}
235
236#[derive(Debug)]
237struct Tenancy {
238 split: SplitSpec,
239 epoch: LeaseEpoch,
240 lane: Option<LaneId>,
241 state: TenancyState,
242 fenced: bool,
244 progress: Option<SplitProgress>,
249 completed: bool,
251 handed_off: bool,
256}
257
258pub struct CoordinationDriver {
261 coordinator: Box<dyn SplitCoordinator>,
262 wait: crossbeam_channel::Receiver<()>,
265 waker: ControlWaker,
266 tenancies: BTreeMap<PartitionId, Tenancy>,
267 by_split: BTreeMap<SplitId, PartitionId>,
268 pending_lost: Vec<LaneId>,
270 pending_retired: Vec<LaneId>,
273 pending_open: Vec<PartitionId>,
276 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 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 #[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 #[must_use]
335 pub fn stall_drains(mut self, drains: bool) -> CoordinationDriver {
336 self.stall_drains = drains;
337 self
338 }
339
340 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 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 let mut park = timeout;
368 loop {
369 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 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 self.tenancies
387 .retain(|_, t| t.state != TenancyState::Retired);
388
389 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 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 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 let events = self.coordinator.poll().map_err(as_source_error)?;
432 let mut surfaced: Option<SourceError> = None;
438 for event in events {
439 if let Err(e) = self.apply(source, event) {
440 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 self.sweep(source)?;
457
458 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 park = Duration::ZERO;
470 continue;
471 }
472 break;
473 }
474
475 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 if !park.is_zero() {
492 let _ = self.wait.recv_timeout(park);
493 }
494 Ok(SourceEvent::Idle)
495 }
496
497 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 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 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(()); };
533 match self.coordinator.fail(split, reason) {
534 Ok(()) => {}
535 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 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 #[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 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 fn report_poison(&mut self, split: &SplitId, reason: &str) -> bool {
584 match self.coordinator.fail(split, reason) {
585 Ok(()) => true,
586 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 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 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 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 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 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 }
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 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 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 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; };
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 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 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; };
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 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 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 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 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 self.retire(source, partition, false);
926 }
927 }
928 CommitDisposition::Fenced => {}
930 CommitDisposition::Deferred => {
931 let tenancy = self.tenancies.get_mut(&partition).expect("live tenancy");
939 tenancy.progress = Some(progress);
940 }
941 }
942 Ok(())
943 }
944
945 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 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 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 CommitDisposition::Fenced => {}
992 CommitDisposition::Deferred => {
993 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
1005enum CommitDisposition {
1008 Durable,
1010 Fenced,
1013 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 #[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 fails: Vec<(SplitId, String)>,
1056 released: Vec<SplitId>,
1057 released_drained: Vec<SplitId>,
1060 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 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 #[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 reject_resume: HashMap<String, ErrorClass>,
1261 finishing: Vec<String>,
1262 fail_open: Vec<String>,
1265 accept_revoke: HashSet<String>,
1268 begin_revoke_calls: Vec<String>,
1270 drain_ready_calls: Vec<String>,
1273 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 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 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 #[test]
1403 fn a_signal_cuts_the_control_plane_park_short() {
1404 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 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 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 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 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 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 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 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 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 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 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 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 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 let event = poll(&mut d, &mut s);
1618 assert!(
1619 matches!(event, SourceEvent::LanesRevoked { ref lanes, .. } if lanes[..] == [LaneId(0)])
1620 );
1621
1622 s.encoded.clear();
1624 d.commit(&mut s, &[(a_partition, 11)]).unwrap();
1625 assert!(s.encoded.is_empty());
1626
1627 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 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 poll(&mut d, &mut s);
1692 assert!(script.commits().is_empty(), "a tick recommitted nothing");
1693
1694 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 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 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 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 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 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 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 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 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 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 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 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 #[test]
1905 fn a_drain_keeps_the_tenancy_commit_eligible_until_the_final_commit() {
1906 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 let script = Script::default();
2146 let mut d = driver(&script);
2147 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 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 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 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 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 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 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 let script = Script::default();
2231 let mut d = driver(&script);
2232 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 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 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 assert!(script.released_drained().is_empty());
2270 }
2271
2272 #[test]
2273 fn a_completed_progress_during_a_drain_is_never_terminal() {
2274 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 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 s.complete_at.insert("a".into(), 42);
2297 d.commit(&mut s, &[(partition, 42)]).unwrap();
2298
2299 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 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 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}