1#![cfg(feature = "cluster")]
4#![allow(clippy::disallowed_types)] use std::sync::{atomic::AtomicU64, atomic::Ordering, Arc};
7use std::time::Duration;
8
9use laminar_connectors::connector::{SourceDrainOutcome, SourceDrainResolution};
10use laminar_core::checkpoint::{
11 AssignmentDrainId, AssignmentDrainTransition, CheckpointAttempt, CheckpointAttemptRelation,
12 CheckpointScope, CommittedCheckpointIndex, CommittedCheckpointRef,
13};
14use laminar_core::checkpoint_decision::{
15 CheckpointArtifactInventory, CheckpointOutcome, CheckpointVerdict, RecordOutcomeResult,
16};
17use laminar_core::cluster::control::{
18 AssignmentDrainDecision, AssignmentDrainVerdict, AssignmentRecoveryDecision,
19 AssignmentSnapshot, AssignmentSnapshotStore, CheckpointAssignmentAdoption,
20 CheckpointAssignmentFence, CheckpointParticipant, ClusterController, LeaderLeaseStore,
21 LeaderProof, ProcessLeaseFence, RecordAssignmentDrainDecisionResult,
22 RecordAssignmentRecoveryDecisionResult, RotateOutcome, SnapshotError,
23};
24use laminar_core::cluster::discovery::NodeState;
25use laminar_core::state::{
26 owners_per_domain, rendezvous_assignment, Locality, NodeId, VnodeRegistry,
27};
28#[cfg(test)]
29use tokio::sync::Notify;
30use tokio::task::JoinHandle;
31use tokio::time::{Instant, MissedTickBehavior};
32use tokio_util::sync::CancellationToken;
33use tracing::{debug, info, warn};
34
35use crate::db::{DbState, LaminarDB};
36use crate::engine_metrics::EngineMetrics;
37
38mod recovery_adoption;
39use recovery_adoption::{
40 adopt_materialized_recovery_head, ensure_local_recovery_fault, local_recovery_assignment_scope,
41 prepare_recovery_assignment_adoption, prepare_watched_recovery_adoption,
42 LocalRecoveryAssignmentScope,
43};
44
45#[derive(Debug, Clone, Copy)]
47pub struct RebalanceConfig {
48 pub watcher_poll: Duration,
50 pub rebalance_debounce: Duration,
52 pub checkpoint_timeout: Duration,
55 pub retry_delay: Duration,
57 pub drain_ack_timeout: Duration,
60 pub placement_isolation_tier: usize,
62}
63
64impl Default for RebalanceConfig {
65 fn default() -> Self {
66 Self {
67 watcher_poll: Duration::from_secs(2),
68 rebalance_debounce: Duration::from_secs(5),
69 checkpoint_timeout: Duration::from_secs(15),
72 retry_delay: Duration::from_secs(2),
73 drain_ack_timeout: Duration::from_secs(10),
74 placement_isolation_tier: 0,
75 }
76 }
77}
78
79impl RebalanceConfig {
80 #[doc(hidden)]
82 #[must_use]
83 pub fn test_defaults() -> Self {
84 Self {
85 watcher_poll: Duration::from_millis(200),
86 rebalance_debounce: Duration::from_millis(500),
87 checkpoint_timeout: Duration::from_secs(30),
88 retry_delay: Duration::from_millis(500),
89 drain_ack_timeout: Duration::from_secs(5),
90 placement_isolation_tier: 0,
91 }
92 }
93}
94
95async fn close_local_assignment_authority(
96 db: &Arc<LaminarDB>,
97 controller: &ClusterController,
98 recovery_target: &CheckpointAssignmentFence,
99 removed_processes: &[CheckpointParticipant],
100 deadline: tokio::time::Instant,
101) -> Result<Vec<ProcessLeaseFence>, String> {
102 db.set_source_gate(true);
103 controller.publish_checkpoint_assignment_fence(None);
104 controller.publish_checkpoint_drain_transition(None);
105 db.invalidate_shuffle_assignment_fence();
107 let _adoption = tokio::time::timeout_at(deadline, db.assignment_adoption_lock.lock())
108 .await
109 .map_err(|_| "timed out serializing assignment authority closure".to_string())?;
110 db.set_source_gate(true);
114 controller.publish_checkpoint_assignment_fence(None);
115 controller.publish_checkpoint_drain_transition(None);
116 db.invalidate_shuffle_assignment_fence();
117 let fence_results = futures::future::join_all(
118 removed_processes
119 .iter()
120 .copied()
121 .map(|participant| controller.fence_process_incarnation(participant, deadline)),
122 )
123 .await;
124 let mut process_fences = Vec::with_capacity(fence_results.len());
125 for result in fence_results {
126 process_fences.push(result?);
127 }
128 if recovery_target.participant_incarnation(controller.instance_id().0)
132 == Some(controller.recovery_incarnation())
133 {
134 ensure_local_recovery_fault(db, controller).await?;
135 }
136 let _transition = tokio::time::timeout_at(
137 deadline,
138 Arc::clone(&db.rotation_execution_fence).write_owned(),
139 )
140 .await
141 .map_err(|_| "timed out draining assignment execution after closure".to_string())?;
142 Ok(process_fences)
143}
144
145async fn suspend_local_assignment_authority(
148 db: &Arc<LaminarDB>,
149 controller: Option<&ClusterController>,
150 deadline: tokio::time::Instant,
151) -> Result<(), String> {
152 db.set_source_gate(true);
153 if let Some(controller) = controller {
154 controller.publish_checkpoint_assignment_fence(None);
155 controller.publish_checkpoint_drain_transition(None);
156 }
157 db.suspend_shuffle_assignment_fence();
160 let _adoption = tokio::time::timeout_at(deadline, db.assignment_adoption_lock.lock())
161 .await
162 .map_err(|_| "timed out serializing assignment authority suspension".to_string())?;
163 suspend_local_assignment_authority_locked(db, controller, deadline).await
164}
165
166async fn suspend_local_assignment_authority_locked(
168 db: &Arc<LaminarDB>,
169 controller: Option<&ClusterController>,
170 deadline: tokio::time::Instant,
171) -> Result<(), String> {
172 db.set_source_gate(true);
175 if let Some(controller) = controller {
176 controller.publish_checkpoint_assignment_fence(None);
177 controller.publish_checkpoint_drain_transition(None);
178 }
179 db.suspend_shuffle_assignment_fence();
180 let _transition = tokio::time::timeout_at(
181 deadline,
182 Arc::clone(&db.rotation_execution_fence).write_owned(),
183 )
184 .await
185 .map_err(|_| "timed out draining assignment execution after suspension".to_string())?;
186 Ok(())
187}
188
189async fn try_suspend_recovery_assignment_authority(
193 db: &Arc<LaminarDB>,
194 controller: &ClusterController,
195 deadline: tokio::time::Instant,
196) -> Result<bool, String> {
197 let _adoption = tokio::time::timeout_at(deadline, db.assignment_adoption_lock.lock())
198 .await
199 .map_err(|_| "timed out serializing recovery assignment suspension".to_string())?;
200 if db.has_unapplied_vnode_transition() {
201 return Ok(false);
202 }
203 suspend_local_assignment_authority_locked(db, Some(controller), deadline).await?;
204 Ok(true)
205}
206
207async fn hold_terminal_source_resolution(
208 db: &Arc<LaminarDB>,
209 controller: Option<&ClusterController>,
210 round: AssignmentDrainId,
211 deadline: tokio::time::Instant,
212 held: &mut Option<(AssignmentDrainId, u64)>,
213) -> Result<u64, String> {
214 let revision = db
215 .assignment_authority_revision
216 .load(std::sync::atomic::Ordering::Acquire);
217 let controller_closed = controller.is_none_or(|controller| {
218 controller
219 .checkpoint_assignment_fence(round.predecessor_version)
220 .is_none()
221 && controller
222 .checkpoint_assignment_fence(round.target_version)
223 .is_none()
224 });
225 if *held == Some((round, revision)) && db.cluster_intake_fenced() && controller_closed {
226 return Ok(revision);
227 }
228
229 suspend_local_assignment_authority(db, controller, deadline).await?;
230 let revision = db
231 .assignment_authority_revision
232 .load(std::sync::atomic::Ordering::Acquire);
233 *held = Some((round, revision));
234 Ok(revision)
235}
236
237struct SnapshotWatcher {
243 db: Arc<LaminarDB>,
244 store: Arc<AssignmentSnapshotStore>,
245 registry: Arc<VnodeRegistry>,
246 shutdown: CancellationToken,
247 config: RebalanceConfig,
248 controller: Option<Arc<ClusterController>>,
249 ticker: tokio::time::Interval,
250 metrics_version: u64,
251 last_drained: u64,
252 durable_snapshot: Option<AssignmentSnapshot>,
253 durable_drain_transition: Option<AssignmentDrainTransition>,
254 active_local_drain: Option<AssignmentDrainTransition>,
255 terminal_resolution_hold: Option<(AssignmentDrainId, u64)>,
256 installed_fence: Option<(u64, [u8; 32])>,
257 installed_authority_revision: u64,
258 assignment_authority_dirty: bool,
259}
260
261impl SnapshotWatcher {
262 fn new(
263 db: Arc<LaminarDB>,
264 store: Arc<AssignmentSnapshotStore>,
265 registry: Arc<VnodeRegistry>,
266 shutdown: CancellationToken,
267 config: RebalanceConfig,
268 controller: Option<Arc<ClusterController>>,
269 ) -> Self {
270 let mut ticker = tokio::time::interval(config.watcher_poll);
271 ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
272 Self {
273 db,
274 store,
275 registry,
276 shutdown,
277 config,
278 controller,
279 ticker,
280 metrics_version: 0,
281 last_drained: 0,
282 durable_snapshot: None,
283 durable_drain_transition: None,
284 active_local_drain: None,
285 terminal_resolution_hold: None,
286 installed_fence: None,
287 installed_authority_revision: 0,
288 assignment_authority_dirty: false,
289 }
290 }
291
292 async fn install_checkpoint_drain_transition(
297 &self,
298 draining: &AssignmentSnapshot,
299 transition: &AssignmentDrainTransition,
300 predecessor: AssignmentSnapshot,
301 deadline: tokio::time::Instant,
302 ) -> Result<AssignmentSnapshot, String> {
303 let Some(controller) = self.controller.as_deref() else {
304 return Ok(predecessor);
305 };
306 if controller.checkpoint_drain_transition().as_ref() == Some(transition) {
307 return Ok(predecessor);
308 }
309 let _adoption = tokio::time::timeout_at(deadline, self.db.assignment_adoption_lock.lock())
310 .await
311 .map_err(|_| {
312 "snapshot watcher timed out serializing source-drain activation".to_string()
313 })?;
314 let predecessor =
315 audit_exact_drain_head(&self.store, &self.registry, draining, controller, deadline)
316 .await?;
317 self.db
318 .validate_source_drain_snapshot(draining)
319 .map_err(|error| error.to_string())?;
320 match controller.checkpoint_drain_transition() {
321 None => controller.publish_checkpoint_drain_transition(Some(transition.clone())),
322 Some(installed) if installed == *transition => {}
323 Some(_) => {
324 return Err(
325 "snapshot watcher found a conflicting local source-drain transition".into(),
326 );
327 }
328 }
329 Ok(predecessor)
330 }
331
332 async fn ensure_current_assignment_authority_cached(
333 &mut self,
334 audited_successor: &AuditedAssignmentAuthority,
335 deadline: tokio::time::Instant,
336 ) -> Result<(), String> {
337 let assignment = self.registry.versioned_snapshot();
338 let version = assignment.version();
339 let expected_predecessor = audited_successor.retained_predecessor(version)?;
340 let cache_is_exact = self.durable_snapshot.as_ref().is_some_and(|snapshot| {
341 !snapshot.draining
342 && snapshot.version == version
343 && snapshot.has_canonical_participants()
344 && snapshot
345 .assignment_fence()
346 .is_ok_and(|fence| &fence == expected_predecessor)
347 && snapshot
348 .to_vnode_vec(self.registry.vnode_count())
349 .is_ok_and(|owners| owners.as_slice() == assignment.owners())
350 });
351 if cache_is_exact {
352 return Ok(());
353 }
354
355 let snapshot = tokio::time::timeout_at(deadline, self.store.load_version(version))
356 .await
357 .map_err(|_| format!("assignment {version} history read timed out"))?
358 .map_err(|error| error.to_string())?
359 .ok_or_else(|| format!("assignment {version} is absent from durable history"))?;
360 if snapshot.draining
361 || !snapshot.has_canonical_participants()
362 || snapshot
363 .to_vnode_vec(self.registry.vnode_count())
364 .map_err(|error| error.to_string())?
365 .as_slice()
366 != assignment.owners()
367 {
368 return Err(format!(
369 "assignment {version} durable history does not match the current registry"
370 ));
371 }
372 if snapshot
373 .assignment_fence()
374 .map_err(|error| error.to_string())?
375 != *expected_predecessor
376 {
377 return Err(format!(
378 "assignment {version} durable history is not the audited successor's exact predecessor"
379 ));
380 }
381 self.durable_drain_transition = None;
382 self.durable_snapshot = Some(snapshot);
383 Ok(())
384 }
385
386 async fn publish_authority(
387 &mut self,
388 mut authority_revision: u64,
389 operation_timeout: Duration,
390 ) {
391 let (head_deadline, current_authority_revision) = (
393 Instant::now() + operation_timeout,
394 AtomicU64::load(&self.db.assignment_authority_revision, Ordering::Acquire),
395 );
396 if authority_revision != current_authority_revision {
397 self.db.set_source_gate(true);
400 if let Some(ref c) = self.controller {
401 c.publish_checkpoint_drain_transition(None);
402 c.publish_checkpoint_assignment_fence(None);
403 }
404 self.assignment_authority_dirty = true;
405 return;
406 }
407 if self.installed_fence.is_some()
408 && self.installed_authority_revision != current_authority_revision
409 {
410 self.assignment_authority_dirty = true;
411 }
412
413 let Ok(adoption_guard) =
417 tokio::time::timeout_at(head_deadline, self.db.assignment_adoption_lock.lock()).await
418 else {
419 self.assignment_authority_dirty = true;
420 warn!(
421 "vnode-state readiness publication timed out waiting for assignment serialization"
422 );
423 return;
424 };
425 let assignment = self.registry.versioned_snapshot();
426 let version = assignment.version();
427 if let Some(ref c) = self.controller {
428 let participant = CheckpointParticipant {
429 node_id: c.instance_id().0,
430 boot_incarnation: c.recovery_incarnation(),
431 };
432 let exact_assignment_fence = self
433 .durable_snapshot
434 .as_ref()
435 .filter(|snapshot| {
436 !snapshot.draining
437 && snapshot.version == version
438 && snapshot.has_canonical_participants()
439 && snapshot
440 .to_vnode_vec(self.registry.vnode_count())
441 .is_ok_and(|owners| owners.as_slice() == assignment.owners())
442 })
443 .and_then(|snapshot| snapshot.assignment_fence().ok());
444 let local_is_participant = exact_assignment_fence.as_ref().is_some_and(|fence| {
445 fence.participant_incarnation(participant.node_id)
446 == Some(participant.boot_incarnation)
447 });
448 if let Some(exact_assignment_fence) = exact_assignment_fence {
449 let vnode_state_ready = if local_is_participant {
450 match tokio::time::timeout_at(
451 head_deadline,
452 self.db
453 .local_vnode_state_is_ready(&self.registry, &exact_assignment_fence),
454 )
455 .await
456 {
457 Ok(Ok(ready)) => ready,
458 Ok(Err(error)) => {
459 warn!(%error, version, "could not determine local vnode-state readiness");
460 drop(adoption_guard);
461 return;
462 }
463 Err(_) => {
464 warn!(
465 version,
466 "local vnode-state readiness exceeded the head deadline"
467 );
468 drop(adoption_guard);
469 return;
470 }
471 }
472 } else {
473 false
477 };
478 let publication =
482 self.db
483 .publish_local_vnode_state_report(c, &assignment, vnode_state_ready);
484 match tokio::time::timeout_at(head_deadline, publication).await {
485 Ok(Ok(_)) => {}
486 Ok(Err(error)) => {
487 warn!(%error, version, "assignment vnode-state report publication failed; retrying");
488 drop(adoption_guard);
489 return;
490 }
491 Err(_) => {
492 warn!(
493 version,
494 "assignment vnode-state report publication exceeded the head deadline"
495 );
496 drop(adoption_guard);
497 return;
498 }
499 }
500 }
501 }
502 drop(adoption_guard);
503 if version != self.metrics_version {
504 self.metrics_version = version;
505 if let (Some(c), Some(metrics)) = (self.controller.as_ref(), self.db.engine_metrics()) {
506 let nodes = c.assignable_with_locality();
507 publish_placement_metrics(
508 &metrics,
509 &self.registry,
510 &nodes,
511 self.config.placement_isolation_tier,
512 );
513 }
514 }
515
516 if let Some(ref c) = self.controller {
520 let fence = match self.durable_snapshot.as_ref() {
521 Some(snapshot)
522 if !snapshot.draining
523 && snapshot.version == version
524 && snapshot.has_canonical_participants()
525 && snapshot
526 .to_vnode_vec(self.registry.vnode_count())
527 .is_ok_and(|owners| owners.as_slice() == assignment.owners()) =>
528 {
529 if let Ok(fence) = tokio::time::timeout_at(
530 head_deadline,
531 compute_checkpoint_assignment_fence(
532 c,
533 &self.registry,
534 &snapshot.participants,
535 ),
536 )
537 .await
538 {
539 fence.filter(|fence| fence.participants == snapshot.participants)
540 } else {
541 warn!(
542 version,
543 "assignment fence computation exceeded the head deadline"
544 );
545 None
546 }
547 }
548 _ => None,
549 };
550 if authority_revision
551 != self
552 .db
553 .assignment_authority_revision
554 .load(std::sync::atomic::Ordering::Acquire)
555 {
556 self.db.set_source_gate(true);
557 c.publish_checkpoint_drain_transition(None);
558 c.publish_checkpoint_assignment_fence(None);
559 self.assignment_authority_dirty = true;
560 return;
561 }
562 let drain_transition = self
563 .durable_drain_transition
564 .as_ref()
565 .filter(|transition| fence.as_ref() == Some(&transition.predecessor))
566 .cloned();
567 match fence {
568 Some(fence) => {
569 let identity = (fence.assignment_version, fence.digest());
570 let local_has_authority = fence.participant_incarnation(c.instance_id().0)
571 == Some(c.recovery_incarnation());
572 let published_fence = c.checkpoint_assignment_fence(fence.assignment_version);
573 let needs_activation = self.assignment_authority_dirty
574 || self.installed_fence != Some(identity)
575 || published_fence.as_ref() != Some(&fence)
576 || c.checkpoint_drain_transition() != drain_transition
577 || (local_has_authority
578 && !c.is_recovering()
579 && drain_transition.is_none()
580 && self.db.cluster_intake_fenced());
581 if needs_activation {
582 if self.installed_fence.is_some() && self.installed_fence != Some(identity)
583 {
584 self.db.set_source_gate(true);
589 c.publish_checkpoint_drain_transition(None);
590 c.publish_checkpoint_assignment_fence(None);
591 self.db.suspend_shuffle_assignment_fence();
592 authority_revision = self
593 .db
594 .assignment_authority_revision
595 .load(std::sync::atomic::Ordering::Acquire);
596 self.installed_fence = None;
597 }
598 match self
599 .db
600 .activate_watcher_assignment_authority(
601 c,
602 &fence,
603 drain_transition,
604 authority_revision,
605 head_deadline,
606 operation_timeout,
607 )
608 .await
609 {
610 Ok(activation) if activation.installed => {
611 self.installed_fence = Some(identity);
612 self.installed_authority_revision = activation.revision;
613 self.assignment_authority_dirty = false;
614 }
615 Ok(_) => self.assignment_authority_dirty = true,
616 Err(error) => {
617 self.installed_fence = None;
618 self.assignment_authority_dirty = true;
619 warn!(%error, version, "shuffle assignment certificate install failed");
620 }
621 }
622 }
623 }
624 None if !self.assignment_authority_dirty => {
625 self.assignment_authority_dirty = true;
630 match suspend_local_assignment_authority(
631 &self.db,
632 Some(c.as_ref()),
633 head_deadline,
634 )
635 .await
636 {
637 Ok(()) => warn!(
638 version,
639 "owner-complete assignment certificate unavailable; authority suspended"
640 ),
641 Err(error) => warn!(
642 %error,
643 version,
644 "assignment suspension could not drain the prior execution scope"
645 ),
646 }
647 }
648 None => {}
649 }
650 }
651 }
652 async fn settle_terminal_drain(
653 &mut self,
654 snapshot: &AssignmentSnapshot,
655 audited_terminal: Option<&AuditedDrainOutcome>,
656 head_deadline: tokio::time::Instant,
657 authority_revision: &mut u64,
658 ) -> bool {
659 let transition = self.active_local_drain.clone().or_else(|| {
660 audited_terminal.and_then(|audited| {
661 self.controller
662 .as_deref()
663 .and_then(|controller| local_drain_participant(controller, &audited.transition))
664 .map(|_| audited.transition.clone())
665 })
666 });
667 let Some(transition) = transition else {
668 self.terminal_resolution_hold = None;
669 return true;
670 };
671
672 let audited = audited_terminal.filter(|audited| audited.transition == transition);
673 let stopped_terminal_may_defer = match (audited, self.controller.as_deref()) {
678 (Some(audited), Some(controller)) => {
679 let predecessor = audited.transition.predecessor.clone();
680 crate::db::audited_stopped_terminal_round(controller, &predecessor, head_deadline)
681 .await
682 .is_ok_and(|round| round.is_some())
683 && DbState::load(&self.db.state) == DbState::Created
684 && self.db.runtime_shutdown.read().is_cancelled()
685 && self.db.cluster_intake_fenced()
686 && self.db.coordinated_recovery_in_progress()
687 }
688 _ => false,
689 };
690 if stopped_terminal_may_defer {
691 self.active_local_drain = None;
692 self.terminal_resolution_hold = None;
693 return true;
694 }
695 let already_resolved = match audited {
696 Some(audited) => {
697 let resolution = SourceDrainResolution {
698 round: transition.id(),
699 outcome: audited.outcome,
700 };
701 match crate::pipeline::streaming_coordinator::owned_source_drain_resolved(
702 &self.db.owned_source_tasks,
703 resolution,
704 ) {
705 Ok(resolved) => resolved,
706 Err(error) => {
707 self.assignment_authority_dirty = true;
708 let _ = hold_terminal_source_resolution(
709 &self.db,
710 self.controller.as_deref(),
711 transition.id(),
712 head_deadline,
713 &mut self.terminal_resolution_hold,
714 )
715 .await;
716 warn!(%error, version = snapshot.version, "snapshot watcher: source drain status audit failed; assignment authority suspended");
717 return false;
718 }
719 }
720 }
721 None => false,
722 };
723 if already_resolved {
724 self.active_local_drain = None;
725 self.terminal_resolution_hold = None;
726 return true;
727 }
728
729 self.assignment_authority_dirty = true;
730 match hold_terminal_source_resolution(
731 &self.db,
732 self.controller.as_deref(),
733 transition.id(),
734 head_deadline,
735 &mut self.terminal_resolution_hold,
736 )
737 .await
738 {
739 Ok(revision) => *authority_revision = revision,
740 Err(error) => {
741 warn!(%error, version = snapshot.version, "snapshot watcher: could not suspend authority for source drain resolution");
742 return false;
743 }
744 }
745 match settle_observed_local_drain(
746 &self.db,
747 &self.store,
748 &self.registry,
749 self.controller.as_deref(),
750 &transition,
751 snapshot,
752 audited,
753 head_deadline,
754 SourceDrainResolutionDeadline::Fresh(self.config.drain_ack_timeout),
755 )
756 .await
757 {
758 Ok(true) => {
759 self.active_local_drain = None;
760 self.terminal_resolution_hold = None;
761 *authority_revision = self
762 .db
763 .assignment_authority_revision
764 .load(std::sync::atomic::Ordering::Acquire);
765 true
766 }
767 Ok(false) => false,
768 Err(error) => {
769 warn!(%error, version = snapshot.version, "snapshot watcher: local source drain resolution failed; assignment authority suspended");
770 false
771 }
772 }
773 }
774
775 async fn run(mut self) {
776 loop {
777 tokio::select! {
778 biased;
779 () = self.shutdown.cancelled() => return,
780 _ = self.ticker.tick() => {}
781 }
782 let local = self.registry.assignment_version();
783 let head_deadline = tokio::time::Instant::now() + self.config.checkpoint_timeout;
784 let mut authority_revision = self
785 .db
786 .assignment_authority_revision
787 .load(std::sync::atomic::Ordering::Acquire);
788
789 let audit = tokio::select! {
793 biased;
794 () = self.shutdown.cancelled() => return,
795 result = tokio::time::timeout_at(head_deadline, self.store.load()) => result,
796 };
797 let mut audited_target = None;
798 let mut audited_terminal = None;
799 let mut audited_recovery = false;
800 if let Ok(Ok(Some(snapshot))) = &audit {
801 if !snapshot.draining {
802 let materialization_audit = tokio::time::timeout_at(
803 head_deadline,
804 audit_assignment_snapshot_authority_outcome(
805 &self.store,
806 self.controller.as_deref(),
807 snapshot,
808 ),
809 )
810 .await;
811 match materialization_audit {
812 Ok(Ok(outcome)) => {
813 audited_recovery = outcome.is_recovery();
814 audited_terminal = outcome.terminal().cloned();
815 audited_target = Some(outcome);
816 }
817 failed => {
818 self.assignment_authority_dirty = true;
819 let _ = suspend_local_assignment_authority(
820 &self.db,
821 self.controller.as_deref(),
822 head_deadline,
823 )
824 .await;
825 match failed {
826 Ok(Err(error)) => {
827 warn!(%error, version = snapshot.version, "snapshot watcher: drain finalization audit failed; assignment authority suspended");
828 }
829 Err(_) => {
830 warn!(version = snapshot.version, timeout = ?self.config.checkpoint_timeout, "snapshot watcher: drain finalization audit timed out; assignment authority suspended");
831 }
832 Ok(Ok(_)) => unreachable!(),
833 }
834 continue;
835 }
836 }
837 }
838 }
839 if let Ok(Ok(Some(snapshot))) = &audit {
840 if !snapshot.draining
841 && !self
842 .settle_terminal_drain(
843 snapshot,
844 audited_terminal.as_ref(),
845 head_deadline,
846 &mut authority_revision,
847 )
848 .await
849 {
850 continue;
851 }
852 }
853 match audit {
854 Ok(Ok(Some(snap))) if snap.draining && snap.version > local => {
857 let predecessor = tokio::time::timeout_at(
864 head_deadline,
865 audit_drain_predecessor(
866 &self.store,
867 &self.registry,
868 &snap,
869 self.controller.as_deref(),
870 ),
871 )
872 .await;
873 match predecessor {
874 Ok(Ok(predecessor)) => {
875 let transition = snap
876 .drain_transition
877 .as_ref()
878 .expect("validated draining snapshot has a transition")
879 .clone();
880 let predecessor = match self
881 .install_checkpoint_drain_transition(
882 &snap,
883 &transition,
884 predecessor,
885 head_deadline,
886 )
887 .await
888 {
889 Ok(predecessor) => predecessor,
890 Err(error) => {
891 self.durable_snapshot = None;
892 self.durable_drain_transition = None;
893 self.assignment_authority_dirty = true;
894 let _ = suspend_local_assignment_authority(
895 &self.db,
896 self.controller.as_deref(),
897 head_deadline,
898 )
899 .await;
900 warn!(%error, version = snap.version, "snapshot watcher: source-drain activation failed; assignment authority suspended");
901 continue;
902 }
903 };
904 self.durable_drain_transition = Some(transition.clone());
905 self.durable_snapshot = Some(predecessor);
906 if snap.version != self.last_drained {
907 match self.db.validate_source_drain_snapshot(&snap) {
908 Ok(()) => {
909 let acknowledgement = async {
910 let c = self.controller.as_ref().ok_or_else(|| {
911 "assignment drain has no cluster controller"
912 .to_string()
913 })?;
914 let Some(participant) =
915 local_drain_participant(c, &transition)
916 else {
917 return Ok(());
921 };
922 self.active_local_drain = Some(transition.clone());
923 prepare_and_announce_local_drain(
924 &self.db,
925 &self.store,
926 &self.registry,
927 c,
928 &snap,
929 &transition,
930 participant,
931 std::cmp::min(
932 head_deadline,
933 tokio::time::Instant::now()
934 + self.config.drain_ack_timeout,
935 ),
936 )
937 .await
938 }
939 .await;
940 match acknowledgement {
941 Ok(()) => {
942 self.last_drained = snap.version;
943 continue;
947 }
948 Err(error) => {
949 warn!(%error, version = snap.version, "snapshot watcher: durable drain acknowledgement failed; retrying");
950 continue;
951 }
952 }
953 }
954 Err(error) => {
955 warn!(%error, "snapshot watcher: draining adoption failed");
956 }
957 }
958 }
959 }
960 Ok(Err(error)) => {
961 self.durable_snapshot = None;
962 self.durable_drain_transition = None;
963 self.assignment_authority_dirty = true;
964 let _ = suspend_local_assignment_authority(
965 &self.db,
966 self.controller.as_deref(),
967 head_deadline,
968 )
969 .await;
970 warn!(%error, version = snap.version, "snapshot watcher: draining predecessor audit failed; assignment authority suspended");
971 }
972 Err(_) => {
973 self.durable_snapshot = None;
974 self.durable_drain_transition = None;
975 self.assignment_authority_dirty = true;
976 let _ = suspend_local_assignment_authority(
977 &self.db,
978 self.controller.as_deref(),
979 head_deadline,
980 )
981 .await;
982 warn!(version = snap.version, timeout = ?self.config.checkpoint_timeout, "snapshot watcher: draining predecessor audit timed out; assignment authority suspended");
983 }
984 }
985 }
986 Ok(Ok(Some(snap))) if !snap.draining && snap.version > local => {
987 if audited_recovery {
988 let Some(controller) = self.controller.as_deref() else {
989 self.assignment_authority_dirty = true;
990 let _ =
991 suspend_local_assignment_authority(&self.db, None, head_deadline)
992 .await;
993 warn!(
994 version = snap.version,
995 "snapshot watcher: recovery assignment has no cluster controller"
996 );
997 continue;
998 };
999 let recovery_scope = match local_recovery_assignment_scope(
1000 &snap, controller,
1001 ) {
1002 Ok(scope) => scope,
1003 Err(error) => {
1004 self.assignment_authority_dirty = true;
1005 let _ = suspend_local_assignment_authority(
1006 &self.db,
1007 Some(controller),
1008 head_deadline,
1009 )
1010 .await;
1011 warn!(%error, version = snap.version, "snapshot watcher: recovery target does not bind the local process");
1012 continue;
1013 }
1014 };
1015 if recovery_scope == LocalRecoveryAssignmentScope::Participant {
1016 match try_suspend_recovery_assignment_authority(
1017 &self.db,
1018 controller,
1019 head_deadline,
1020 )
1021 .await
1022 {
1023 Ok(true) => {}
1024 Ok(false) => {
1025 debug!(
1026 version = snap.version,
1027 "snapshot watcher: recovery assignment waits for local vnode transition"
1028 );
1029 if let Err(error) = self
1030 .ensure_current_assignment_authority_cached(
1031 audited_target.as_ref().expect(
1032 "stable successor was audited before suspension",
1033 ),
1034 Instant::now() + self.config.checkpoint_timeout,
1035 )
1036 .await
1037 {
1038 warn!(%error, version = snap.version, "snapshot watcher: could not audit predecessor authority for pending vnode transition");
1039 continue;
1040 }
1041 self.publish_authority(
1045 authority_revision,
1046 self.config.checkpoint_timeout,
1047 )
1048 .await;
1049 continue;
1050 }
1051 Err(error) => {
1052 self.assignment_authority_dirty = true;
1053 warn!(%error, version = snap.version, "snapshot watcher: could not suspend recovery assignment");
1054 continue;
1055 }
1056 }
1057 self.durable_snapshot = None;
1061 self.durable_drain_transition = None;
1062 self.installed_fence = None;
1063 match prepare_watched_recovery_adoption(
1064 &self.db,
1065 &self.store,
1066 controller,
1067 &self.registry,
1068 &snap,
1069 head_deadline,
1070 )
1071 .await
1072 {
1073 Ok(revision) => authority_revision = revision,
1074 Err(error) => {
1075 self.assignment_authority_dirty = true;
1076 warn!(%error, version = snap.version, "snapshot watcher: recovery assignment preparation failed");
1077 continue;
1078 }
1079 }
1080 self.assignment_authority_dirty = true;
1081 } else {
1082 debug!(
1083 version = snap.version,
1084 "snapshot watcher: ownerless process follows recovery topology without publishing a fault"
1085 );
1086 }
1087 }
1088 if snap.version > self.registry.assignment_version() {
1089 debug!(remote = snap.version, "adopting newer assignment");
1090 let recovery_timeout = self.config.checkpoint_timeout;
1091 let adoption = if audited_recovery {
1092 self.db
1093 .adopt_recovery_assignment_snapshot(snap.clone(), recovery_timeout)
1094 .await
1095 } else {
1096 self.db
1097 .adopt_assignment_snapshot(snap.clone(), head_deadline)
1098 .await
1099 };
1100 match adoption {
1101 Ok(adoption) => {
1102 self.durable_drain_transition = None;
1103 self.durable_snapshot = Some(snap.clone());
1104 authority_revision = self
1105 .db
1106 .assignment_authority_revision
1107 .load(std::sync::atomic::Ordering::Acquire);
1108 if adoption.recovery_required {
1109 let Some(controller) = self.controller.as_deref() else {
1110 self.assignment_authority_dirty = true;
1111 warn!(version = snap.version, "snapshot watcher: exact-handoff recovery has no cluster controller");
1112 continue;
1113 };
1114 self.installed_fence = None;
1115 self.assignment_authority_dirty = true;
1116 match local_recovery_assignment_scope(&snap, controller) {
1117 Ok(LocalRecoveryAssignmentScope::Participant) => {
1118 if let Err(error) =
1119 ensure_local_recovery_fault(&self.db, controller)
1120 .await
1121 {
1122 warn!(%error, version = snap.version, "snapshot watcher: could not persist exact-handoff recovery fault");
1123 }
1124 }
1125 Ok(LocalRecoveryAssignmentScope::Ownerless) => {
1126 debug!(
1127 version = snap.version,
1128 "snapshot watcher: ownerless process retains the exact-handoff recovery fence without publishing a fault"
1129 );
1130 }
1131 Err(error) => {
1132 warn!(%error, version = snap.version, "snapshot watcher: exact-handoff recovery target does not bind the local process");
1133 }
1134 }
1135 continue;
1136 }
1137 }
1138 Err(error) if error.is_shuffle_not_ready() => {
1139 debug!(
1143 error = %error,
1144 version = snap.version,
1145 "snapshot watcher: successor waits for local vnode transition"
1146 );
1147 if let Err(cache_error) = self
1148 .ensure_current_assignment_authority_cached(
1149 audited_target
1150 .as_ref()
1151 .expect("stable successor was audited before adoption"),
1152 Instant::now() + self.config.checkpoint_timeout,
1153 )
1154 .await
1155 {
1156 warn!(%cache_error, version = snap.version, "snapshot watcher: could not audit predecessor authority for pending vnode transition");
1157 continue;
1158 }
1159 }
1163 Err(e) => warn!(error = %e, "snapshot watcher: adoption failed"),
1164 }
1165 } else {
1166 self.durable_drain_transition = None;
1167 self.durable_snapshot = Some(snap);
1168 }
1169 }
1170 Ok(Ok(Some(snap))) => {
1171 self.durable_drain_transition
1172 .clone_from(&snap.drain_transition);
1173 self.durable_snapshot = Some(snap);
1174 }
1175 Ok(Ok(None)) => {
1176 self.durable_drain_transition = None;
1177 self.durable_snapshot = None;
1178 }
1179 Ok(Err(error)) => {
1180 self.durable_snapshot = None;
1181 self.durable_drain_transition = None;
1182 self.assignment_authority_dirty = true;
1183 let _ = suspend_local_assignment_authority(
1184 &self.db,
1185 self.controller.as_deref(),
1186 head_deadline,
1187 )
1188 .await;
1189 warn!(%error, "snapshot watcher: durable audit failed; assignment authority suspended");
1190 }
1191 Err(_) => {
1192 self.durable_snapshot = None;
1193 self.durable_drain_transition = None;
1194 self.assignment_authority_dirty = true;
1195 let _ = suspend_local_assignment_authority(
1196 &self.db,
1197 self.controller.as_deref(),
1198 head_deadline,
1199 )
1200 .await;
1201 warn!(
1202 timeout = ?self.config.checkpoint_timeout,
1203 "snapshot watcher: durable audit timed out; assignment authority suspended"
1204 );
1205 }
1206 }
1207
1208 self.publish_authority(authority_revision, self.config.checkpoint_timeout)
1209 .await;
1210 }
1211 }
1212}
1213
1214pub fn spawn_snapshot_watcher(
1220 db: Arc<LaminarDB>,
1221 store: Arc<AssignmentSnapshotStore>,
1222 registry: Arc<VnodeRegistry>,
1223 shutdown: CancellationToken,
1224 config: RebalanceConfig,
1225 controller: Option<Arc<ClusterController>>,
1226) -> JoinHandle<()> {
1227 tokio::spawn(SnapshotWatcher::new(db, store, registry, shutdown, config, controller).run())
1228}
1229async fn compute_checkpoint_assignment_fence(
1232 c: &ClusterController,
1233 registry: &VnodeRegistry,
1234 expected_participants: &[CheckpointParticipant],
1235) -> Option<CheckpointAssignmentFence> {
1236 let assignment = registry.versioned_snapshot();
1237 let participant_ids: Vec<u64> = expected_participants
1238 .iter()
1239 .map(|participant| participant.node_id)
1240 .collect();
1241 let participants = c
1242 .recovery_participant_incarnations(&participant_ids)
1243 .await
1244 .ok()?;
1245 if participants != expected_participants {
1246 return None;
1247 }
1248 let reported: rustc_hash::FxHashMap<u64, CheckpointAssignmentAdoption> = c
1249 .read_adopted_assignments()
1250 .await
1251 .ok()?
1252 .into_iter()
1253 .map(|(node, adoption)| (node.0, adoption))
1254 .collect();
1255 checkpoint_assignment_fence(
1256 assignment.version(),
1257 assignment.owners(),
1258 participants,
1259 &reported,
1260 )
1261}
1262
1263fn checkpoint_assignment_fence(
1264 assignment_version: u64,
1265 owners: &[NodeId],
1266 participants: Vec<CheckpointParticipant>,
1267 reported: &rustc_hash::FxHashMap<u64, CheckpointAssignmentAdoption>,
1268) -> Option<CheckpointAssignmentFence> {
1269 let owner_ids: Vec<u64> = owners.iter().map(|owner| owner.0).collect();
1270 let vnode_count = u32::try_from(owners.len()).ok()?;
1271 let assignment_digest = CheckpointAssignmentFence::owner_map_digest(vnode_count, &owner_ids);
1272 if assignment_version == 0
1273 || owners.is_empty()
1274 || participants.is_empty()
1275 || owners.iter().any(|owner| {
1276 owner.is_unassigned()
1277 || participants
1278 .binary_search_by_key(&owner.0, |participant| participant.node_id)
1279 .is_err()
1280 })
1281 || participants.iter().any(|participant| {
1282 reported.get(&participant.node_id).is_none_or(|adoption| {
1283 adoption.participant != *participant
1287 || adoption.assignment_version != assignment_version
1288 || adoption.vnode_count != vnode_count
1289 || adoption.assignment_digest != assignment_digest
1290 })
1291 })
1292 {
1293 return None;
1294 }
1295
1296 CheckpointAssignmentFence::from_owner_map(assignment_version, &owner_ids, participants).ok()
1297}
1298
1299fn assignment_vnode_state_is_ready(
1300 fence: &CheckpointAssignmentFence,
1301 reported: &rustc_hash::FxHashMap<u64, CheckpointAssignmentAdoption>,
1302) -> bool {
1303 fence.participants.iter().all(|participant| {
1304 reported.get(&participant.node_id).is_some_and(|adoption| {
1305 adoption.vnode_state_ready
1306 && adoption.participant == *participant
1307 && adoption.matches_fence(fence)
1308 })
1309 })
1310}
1311
1312async fn read_assignment_vnode_state_readiness(
1313 controller: &ClusterController,
1314 fence: &CheckpointAssignmentFence,
1315 deadline: tokio::time::Instant,
1316) -> Result<bool, String> {
1317 let reported = tokio::time::timeout_at(deadline, controller.read_adopted_assignments())
1318 .await
1319 .map_err(|_| {
1320 format!(
1321 "vnode-state readiness read for assignment {} timed out",
1322 fence.assignment_version
1323 )
1324 })??
1325 .into_iter()
1326 .map(|(node, adoption)| (node.0, adoption))
1327 .collect();
1328 Ok(assignment_vnode_state_is_ready(fence, &reported))
1329}
1330
1331fn local_drain_participant(
1332 controller: &ClusterController,
1333 transition: &AssignmentDrainTransition,
1334) -> Option<CheckpointParticipant> {
1335 let participant = CheckpointParticipant {
1336 node_id: controller.instance_id().0,
1337 boot_incarnation: controller.recovery_incarnation(),
1338 };
1339 (transition
1340 .predecessor
1341 .participant_incarnation(participant.node_id)
1342 == Some(participant.boot_incarnation))
1343 .then_some(participant)
1344}
1345
1346fn finalized_drain_outcome(
1347 transition: &AssignmentDrainTransition,
1348 finalized: &AssignmentSnapshot,
1349) -> Result<SourceDrainOutcome, String> {
1350 if finalized.draining || finalized.version != transition.target.assignment_version {
1351 return Err(format!(
1352 "assignment {} is not a terminal snapshot for drain target {}",
1353 finalized.version, transition.target.assignment_version
1354 ));
1355 }
1356 let fence = finalized
1357 .assignment_fence()
1358 .map_err(|error| error.to_string())?;
1359 if fence == transition.target {
1360 return Ok(SourceDrainOutcome::Commit);
1361 }
1362 if fence.assignment_version == transition.target.assignment_version
1363 && fence.vnode_count == transition.predecessor.vnode_count
1364 && fence.assignment_digest == transition.predecessor.assignment_digest
1365 && fence.participants == transition.predecessor.participants
1366 {
1367 return Ok(SourceDrainOutcome::Abort);
1368 }
1369 Err(format!(
1370 "assignment {} is neither the committed target nor predecessor rollback for drain {:?}",
1371 finalized.version,
1372 transition.id()
1373 ))
1374}
1375
1376#[derive(Debug, Clone, PartialEq, Eq)]
1377pub(crate) struct AuditedDrainOutcome {
1378 transition: AssignmentDrainTransition,
1379 outcome: SourceDrainOutcome,
1380 handoff_checkpoint: Option<CommittedCheckpointRef>,
1381}
1382
1383#[derive(Debug, Clone, PartialEq, Eq)]
1388pub(crate) struct AuditedAssignmentAuthority {
1389 target_version: u64,
1390 predecessor: Option<CheckpointAssignmentFence>,
1391 handoff_checkpoint: Option<CommittedCheckpointRef>,
1392 recovery_checkpoint: Option<AuditedRecoveryCheckpoint>,
1393 recovery_leader_proof: Option<LeaderProof>,
1394 terminal: Option<AuditedDrainOutcome>,
1395}
1396
1397#[derive(Debug, Clone, PartialEq, Eq)]
1398enum AuditedRecoveryCheckpoint {
1399 Active {
1401 checkpoint: CommittedCheckpointRef,
1402 origin: CheckpointAssignmentFence,
1403 },
1404 Propagated {
1407 checkpoint: CommittedCheckpointRef,
1408 origin: CheckpointAssignmentFence,
1409 successor: CheckpointAssignmentFence,
1410 },
1411 Consumed {
1413 checkpoint: CommittedCheckpointRef,
1414 origin: CheckpointAssignmentFence,
1415 },
1416}
1417
1418impl AuditedRecoveryCheckpoint {
1419 fn checkpoint(&self) -> &CommittedCheckpointRef {
1420 match self {
1421 Self::Active { checkpoint, .. }
1422 | Self::Propagated { checkpoint, .. }
1423 | Self::Consumed { checkpoint, .. } => checkpoint,
1424 }
1425 }
1426
1427 fn origin(&self) -> &CheckpointAssignmentFence {
1428 match self {
1429 Self::Active { origin, .. }
1430 | Self::Propagated { origin, .. }
1431 | Self::Consumed { origin, .. } => origin,
1432 }
1433 }
1434
1435 const fn pin_is_active(&self) -> bool {
1436 matches!(self, Self::Active { .. })
1437 }
1438
1439 const fn was_consumed(&self) -> bool {
1440 matches!(self, Self::Consumed { .. })
1441 }
1442
1443 #[cfg(test)]
1444 const fn was_propagated(&self) -> bool {
1445 matches!(self, Self::Propagated { .. })
1446 }
1447
1448 fn propagated_successor(&self) -> Option<&CheckpointAssignmentFence> {
1449 match self {
1450 Self::Propagated { successor, .. } => Some(successor),
1451 Self::Active { .. } | Self::Consumed { .. } => None,
1452 }
1453 }
1454}
1455
1456impl AuditedAssignmentAuthority {
1457 fn without_successor_edge(target_version: u64) -> Self {
1458 Self {
1459 target_version,
1460 predecessor: None,
1461 handoff_checkpoint: None,
1462 recovery_checkpoint: None,
1463 recovery_leader_proof: None,
1464 terminal: None,
1465 }
1466 }
1467
1468 fn recovery(
1469 target_version: u64,
1470 predecessor: CheckpointAssignmentFence,
1471 recovery_checkpoint: AuditedRecoveryCheckpoint,
1472 recovery_leader_proof: LeaderProof,
1473 ) -> Self {
1474 let handoff_checkpoint = recovery_checkpoint.checkpoint().clone();
1475 Self {
1476 target_version,
1477 predecessor: Some(predecessor),
1478 handoff_checkpoint: Some(handoff_checkpoint),
1479 recovery_checkpoint: Some(recovery_checkpoint),
1480 recovery_leader_proof: Some(recovery_leader_proof),
1481 terminal: None,
1482 }
1483 }
1484
1485 fn drain(target_version: u64, terminal: AuditedDrainOutcome) -> Self {
1486 let handoff_checkpoint = terminal.committed_handoff_checkpoint().cloned();
1487 Self {
1488 target_version,
1489 predecessor: Some(terminal.transition.predecessor.clone()),
1490 handoff_checkpoint,
1491 recovery_checkpoint: None,
1492 recovery_leader_proof: None,
1493 terminal: Some(terminal),
1494 }
1495 }
1496
1497 pub(crate) fn is_recovery(&self) -> bool {
1498 self.recovery_checkpoint.is_some()
1499 }
1500
1501 fn terminal(&self) -> Option<&AuditedDrainOutcome> {
1502 self.terminal.as_ref()
1503 }
1504
1505 pub(crate) fn predecessor(&self) -> Option<&CheckpointAssignmentFence> {
1506 self.predecessor.as_ref()
1507 }
1508
1509 pub(crate) fn handoff_checkpoint(&self) -> Option<&CommittedCheckpointRef> {
1510 self.handoff_checkpoint.as_ref()
1511 }
1512
1513 pub(crate) fn recovery_origin(&self) -> Option<&CheckpointAssignmentFence> {
1514 self.recovery_checkpoint
1515 .as_ref()
1516 .map(AuditedRecoveryCheckpoint::origin)
1517 }
1518
1519 pub(crate) fn recovery_pin_is_active(&self) -> bool {
1520 self.recovery_checkpoint
1521 .as_ref()
1522 .is_some_and(AuditedRecoveryCheckpoint::pin_is_active)
1523 }
1524
1525 pub(crate) fn active_recovery_leader_proof(&self) -> Option<&LeaderProof> {
1526 self.recovery_checkpoint
1527 .as_ref()
1528 .filter(|checkpoint| checkpoint.pin_is_active())
1529 .and(self.recovery_leader_proof.as_ref())
1530 }
1531
1532 pub(crate) fn recovery_checkpoint_was_consumed(&self) -> bool {
1533 self.recovery_checkpoint
1534 .as_ref()
1535 .is_some_and(AuditedRecoveryCheckpoint::was_consumed)
1536 }
1537
1538 #[cfg(test)]
1539 pub(crate) fn recovery_checkpoint_was_propagated(&self) -> bool {
1540 self.recovery_checkpoint
1541 .as_ref()
1542 .is_some_and(AuditedRecoveryCheckpoint::was_propagated)
1543 }
1544
1545 pub(crate) fn propagated_recovery_successor(&self) -> Option<&CheckpointAssignmentFence> {
1546 self.recovery_checkpoint
1547 .as_ref()
1548 .and_then(AuditedRecoveryCheckpoint::propagated_successor)
1549 }
1550
1551 pub(crate) fn into_terminal(self) -> Option<AuditedDrainOutcome> {
1552 self.terminal
1553 }
1554
1555 fn retained_predecessor(
1556 &self,
1557 local_version: u64,
1558 ) -> Result<&CheckpointAssignmentFence, String> {
1559 if local_version.checked_add(1) != Some(self.target_version) {
1560 return Err(format!(
1561 "audited assignment {} is not the exact successor of local assignment {local_version}",
1562 self.target_version
1563 ));
1564 }
1565 self.predecessor.as_ref().ok_or_else(|| {
1566 format!(
1567 "audited assignment {} has no predecessor certificate",
1568 self.target_version
1569 )
1570 })
1571 }
1572}
1573
1574pub(crate) struct AuditedCommittedDrainTransition(AssignmentDrainTransition);
1579
1580impl AuditedCommittedDrainTransition {
1581 pub(crate) fn into_transition(self) -> AssignmentDrainTransition {
1582 self.0
1583 }
1584
1585 #[cfg(test)]
1586 pub(crate) fn from_canonical_for_test(
1587 transition: AssignmentDrainTransition,
1588 ) -> Result<Self, String> {
1589 transition
1590 .is_canonical()
1591 .then_some(Self(transition))
1592 .ok_or_else(|| "test drain transition is not canonical".into())
1593 }
1594}
1595
1596impl AuditedDrainOutcome {
1597 #[must_use]
1598 pub(crate) const fn transition(&self) -> &AssignmentDrainTransition {
1599 &self.transition
1600 }
1601
1602 #[must_use]
1603 pub(crate) fn is_abort(&self) -> bool {
1604 self.outcome == SourceDrainOutcome::Abort
1605 }
1606
1607 #[must_use]
1608 pub(crate) fn is_commit(&self) -> bool {
1609 self.outcome == SourceDrainOutcome::Commit
1610 }
1611
1612 #[must_use]
1613 pub(crate) fn committed_handoff_checkpoint(&self) -> Option<&CommittedCheckpointRef> {
1614 if self.outcome == SourceDrainOutcome::Commit {
1615 self.handoff_checkpoint.as_ref()
1616 } else {
1617 None
1618 }
1619 }
1620
1621 #[must_use]
1623 pub(crate) fn into_committed_transition(self) -> Option<AuditedCommittedDrainTransition> {
1624 (self.outcome == SourceDrainOutcome::Commit)
1625 .then_some(AuditedCommittedDrainTransition(self.transition))
1626 }
1627}
1628
1629pub async fn audit_assignment_snapshot_authority(
1636 store: &AssignmentSnapshotStore,
1637 controller: Option<&ClusterController>,
1638 snapshot: &AssignmentSnapshot,
1639) -> Result<(), String> {
1640 audit_assignment_snapshot_authority_outcome(store, controller, snapshot)
1641 .await
1642 .map(|_| ())
1643}
1644
1645pub async fn startup_committed_assignment(
1652 store: &AssignmentSnapshotStore,
1653 controller: Option<&ClusterController>,
1654 head: AssignmentSnapshot,
1655) -> Result<AssignmentSnapshot, String> {
1656 audit_assignment_snapshot_authority(store, controller, &head)
1657 .await
1658 .map_err(|error| {
1659 format!(
1660 "audit startup assignment {} authority: {error}",
1661 head.version
1662 )
1663 })?;
1664 if !head.draining {
1665 return Ok(head);
1666 }
1667
1668 let transition = store
1669 .load_drain_transition(head.version)
1670 .await
1671 .map_err(|error| {
1672 format!(
1673 "load startup assignment {} drain transition: {error}",
1674 head.version
1675 )
1676 })?
1677 .ok_or_else(|| {
1678 format!(
1679 "draining startup assignment {} has no exact transition",
1680 head.version
1681 )
1682 })?;
1683 if head.drain_transition.as_ref() != Some(&transition) {
1684 return Err(format!(
1685 "draining startup assignment {} does not match its durable transition",
1686 head.version
1687 ));
1688 }
1689 let prior_version = head
1690 .version
1691 .checked_sub(1)
1692 .ok_or_else(|| "draining assignment has no retained committed predecessor".to_string())?;
1693 let prior = store
1694 .load_version(prior_version)
1695 .await
1696 .map_err(|error| {
1697 format!(
1698 "load retained assignment {prior_version} before draining head {}: {error}",
1699 head.version
1700 )
1701 })?
1702 .ok_or_else(|| {
1703 format!(
1704 "draining assignment {} has no retained committed predecessor {prior_version}",
1705 head.version
1706 )
1707 })?;
1708 if prior.draining {
1709 return Err(format!(
1710 "draining assignment {} has a draining predecessor {prior_version}",
1711 head.version
1712 ));
1713 }
1714 if prior
1715 .assignment_fence()
1716 .map_err(|error| error.to_string())?
1717 != transition.predecessor
1718 {
1719 return Err(format!(
1720 "draining assignment {} does not bind retained predecessor {prior_version}",
1721 head.version
1722 ));
1723 }
1724 audit_assignment_snapshot_authority(store, controller, &prior)
1725 .await
1726 .map_err(|error| format!("audit retained assignment {prior_version} authority: {error}"))?;
1727 Ok(prior)
1728}
1729
1730pub(crate) fn audit_assignment_snapshot_authority_outcome<'a>(
1731 store: &'a AssignmentSnapshotStore,
1732 controller: Option<&'a ClusterController>,
1733 snapshot: &'a AssignmentSnapshot,
1734) -> futures::future::BoxFuture<'a, Result<AuditedAssignmentAuthority, String>> {
1735 Box::pin(async move {
1736 if snapshot.draining {
1737 return Ok(AuditedAssignmentAuthority::without_successor_edge(
1738 snapshot.version,
1739 ));
1740 }
1741 if snapshot.version == 1 {
1742 return Ok(AuditedAssignmentAuthority::without_successor_edge(
1743 snapshot.version,
1744 ));
1745 }
1746 let authority = match controller {
1747 Some(controller) => Some(
1748 controller
1749 .checkpoint_authority()
1750 .map_err(|error| error.to_string())?,
1751 ),
1752 None => None,
1753 };
1754 let Some(transition) = store
1755 .load_drain_transition(snapshot.version)
1756 .await
1757 .map_err(|error| error.to_string())?
1758 else {
1759 let (predecessor, handoff_checkpoint, leader_proof) =
1760 audit_materialized_recovery_with_authority(store, authority.as_deref(), snapshot)
1761 .await?;
1762 return Ok(AuditedAssignmentAuthority::recovery(
1763 snapshot.version,
1764 predecessor,
1765 handoff_checkpoint,
1766 leader_proof,
1767 ));
1768 };
1769 audit_materialized_drain_transition(store, authority.as_deref(), snapshot, transition)
1770 .await
1771 .map(|terminal| AuditedAssignmentAuthority::drain(snapshot.version, terminal))
1772 })
1773}
1774
1775async fn audit_materialized_drain_with_authority(
1776 store: &AssignmentSnapshotStore,
1777 authority: Option<&LeaderLeaseStore>,
1778 snapshot: &AssignmentSnapshot,
1779) -> Result<Option<AuditedDrainOutcome>, String> {
1780 if snapshot.draining {
1781 return Ok(None);
1782 }
1783 if snapshot.version == 1 {
1784 return Ok(None);
1785 }
1786 let Some(transition) = store
1787 .load_drain_transition(snapshot.version)
1788 .await
1789 .map_err(|error| error.to_string())?
1790 else {
1791 audit_materialized_recovery_with_authority(store, authority, snapshot).await?;
1792 return Ok(None);
1793 };
1794 audit_materialized_drain_transition(store, authority, snapshot, transition)
1795 .await
1796 .map(Some)
1797}
1798
1799fn validate_portable_recovery_cut(
1800 outcome: &CheckpointOutcome,
1801 index: &CommittedCheckpointIndex,
1802 reference: &CommittedCheckpointRef,
1803 predecessor: &CheckpointAssignmentFence,
1804 context: &str,
1805) -> Result<CheckpointAssignmentFence, String> {
1806 index
1807 .validate()
1808 .map_err(|error| format!("{context} committed checkpoint index is invalid: {error}"))?;
1809 let origin = index
1810 .assignment_fence
1811 .as_ref()
1812 .ok_or_else(|| format!("{context} committed checkpoint has no cluster assignment fence"))?;
1813 if !outcome.is_commit()
1814 || outcome.scope != CheckpointScope::Cluster
1815 || outcome.committed_checkpoint.as_ref() != Some(reference)
1816 || outcome.epoch != reference.epoch
1817 || outcome.checkpoint_id != reference.checkpoint_id
1818 || outcome.epoch != index.epoch
1819 || outcome.checkpoint_id != index.checkpoint_id
1820 || outcome.deployment_id != index.deployment_id
1821 || outcome.assignment_fence.as_ref() != Some(origin)
1822 || outcome.leader_proof.is_none()
1823 || index.scope != CheckpointScope::Cluster
1824 || index.epoch != reference.epoch
1825 || index.checkpoint_id != reference.checkpoint_id
1826 || !index.reassignment_portable
1827 {
1828 return Err(format!(
1829 "{context} does not bind one exact reassignment-portable cluster Commit"
1830 ));
1831 }
1832 if !origin.is_canonical()
1833 || !predecessor.is_canonical()
1834 || origin.assignment_version > predecessor.assignment_version
1835 || origin.vnode_count != predecessor.vnode_count
1836 || origin.partitioning_abi_version != predecessor.partitioning_abi_version
1837 || (origin.assignment_version == predecessor.assignment_version && origin != predecessor)
1838 {
1839 return Err(format!(
1840 "{context} origin assignment is newer than or incompatible with the recovery predecessor"
1841 ));
1842 }
1843 Ok(origin.clone())
1844}
1845
1846type AuditedRecoveryAuthority = (
1847 CheckpointAssignmentFence,
1848 AuditedRecoveryCheckpoint,
1849 LeaderProof,
1850);
1851
1852fn audit_materialized_recovery_with_authority<'a>(
1853 store: &'a AssignmentSnapshotStore,
1854 authority: Option<&'a LeaderLeaseStore>,
1855 snapshot: &'a AssignmentSnapshot,
1856) -> futures::future::BoxFuture<'a, Result<AuditedRecoveryAuthority, String>> {
1857 Box::pin(async move {
1858 let version = snapshot.version;
1859 let authority = authority.ok_or_else(|| {
1860 format!("materialized assignment recovery {version} has no cluster authority")
1861 })?;
1862 let decision = authority
1863 .assignment_recovery_decision(version)
1864 .await
1865 .map_err(|error| error.to_string())?
1866 .ok_or_else(|| {
1867 format!(
1868 "assignment {version} has no drain transition or recovery authority decision"
1869 )
1870 })?;
1871 let proposal = store
1872 .load_recovery_proposal(&decision.proposal)
1873 .await
1874 .map_err(|error| error.to_string())?;
1875 if proposal != *snapshot
1876 || proposal
1877 .assignment_fence()
1878 .map_err(|error| error.to_string())?
1879 != decision.target
1880 {
1881 return Err(format!(
1882 "assignment {} does not match its authorized recovery proposal",
1883 snapshot.version
1884 ));
1885 }
1886 let predecessor_version = version
1887 .checked_sub(1)
1888 .ok_or_else(|| "recovery assignment has no predecessor generation".to_string())?;
1889 let predecessor = store
1890 .load_version(predecessor_version)
1891 .await
1892 .map_err(|error| error.to_string())?
1893 .ok_or_else(|| {
1894 format!("recovery assignment {version} lost predecessor {predecessor_version}")
1895 })?;
1896 if predecessor.draining
1897 || predecessor
1898 .assignment_fence()
1899 .map_err(|error| error.to_string())?
1900 != decision.predecessor
1901 {
1902 return Err(format!(
1903 "assignment {} recovery decision does not bind its exact committed predecessor",
1904 snapshot.version
1905 ));
1906 }
1907 let committed_head = authority
1908 .highest_cluster_committed_outcome()
1909 .await
1910 .map_err(|error| error.to_string())?
1911 .ok_or_else(|| {
1912 format!(
1913 "assignment {} recovery has no committed checkpoint head",
1914 snapshot.version
1915 )
1916 })?;
1917 let head_reference = committed_head.committed_checkpoint.clone().ok_or_else(|| {
1918 format!(
1919 "assignment {} recovery committed head has no checkpoint reference",
1920 snapshot.version
1921 )
1922 })?;
1923 let pinned = authority
1924 .assignment_handoff_checkpoint(&decision.target)
1925 .await
1926 .map_err(|error| error.to_string())?;
1927
1928 let recovery_checkpoint = if let Some(pinned) = pinned {
1929 if pinned != decision.recovery_checkpoint || head_reference != pinned {
1930 return Err(format!(
1931 "assignment {} recovery does not retain its exact authoritative checkpoint head",
1932 snapshot.version
1933 ));
1934 }
1935 let (observed_outcome, observed_index) = authority
1936 .cluster_outcome_with_committed_checkpoint(pinned.epoch)
1937 .await
1938 .map_err(|error| error.to_string())?
1939 .ok_or_else(|| {
1940 format!(
1941 "assignment {} pinned recovery Commit disappeared",
1942 snapshot.version
1943 )
1944 })?;
1945 if observed_outcome != committed_head {
1946 return Err(format!(
1947 "assignment {} recovery checkpoint head changed during authority audit",
1948 snapshot.version
1949 ));
1950 }
1951 let observed_index = observed_index.ok_or_else(|| {
1952 format!(
1953 "assignment {} pinned recovery Commit has no checkpoint index",
1954 snapshot.version
1955 )
1956 })?;
1957 let origin = validate_portable_recovery_cut(
1958 &observed_outcome,
1959 &observed_index,
1960 &pinned,
1961 &decision.predecessor,
1962 "materialized assignment recovery",
1963 )?;
1964 AuditedRecoveryCheckpoint::Active {
1965 checkpoint: pinned,
1966 origin,
1967 }
1968 } else if committed_head.assignment_fence.as_ref() == Some(&decision.target) {
1969 let (observed_outcome, observed_index) = authority
1973 .cluster_outcome_with_committed_checkpoint(head_reference.epoch)
1974 .await
1975 .map_err(|error| error.to_string())?
1976 .ok_or_else(|| {
1977 format!(
1978 "assignment {} consumed recovery Commit disappeared",
1979 snapshot.version
1980 )
1981 })?;
1982 if observed_outcome != committed_head {
1983 return Err(format!(
1984 "assignment {} recovery checkpoint head changed during authority audit",
1985 snapshot.version
1986 ));
1987 }
1988 let observed_index = observed_index.ok_or_else(|| {
1989 format!(
1990 "assignment {} consumed recovery Commit has no checkpoint index",
1991 snapshot.version
1992 )
1993 })?;
1994 let origin = validate_portable_recovery_cut(
1995 &observed_outcome,
1996 &observed_index,
1997 &head_reference,
1998 &decision.target,
1999 "consumed assignment recovery",
2000 )?;
2001 AuditedRecoveryCheckpoint::Consumed {
2002 checkpoint: head_reference,
2003 origin,
2004 }
2005 } else {
2006 let successor_version = snapshot
2011 .version
2012 .checked_add(1)
2013 .ok_or_else(|| "assignment version exhausted during recovery audit".to_string())?;
2014 let successor = authority
2015 .assignment_recovery_decision(successor_version)
2016 .await
2017 .map_err(|error| error.to_string())?
2018 .ok_or_else(|| {
2019 format!(
2020 "assignment {} recovery checkpoint was superseded before the target committed",
2021 snapshot.version
2022 )
2023 })?;
2024 if successor.predecessor != decision.target
2025 || successor.recovery_checkpoint != decision.recovery_checkpoint
2026 || successor.recovery_checkpoint != head_reference
2027 {
2028 return Err(format!(
2029 "assignment {} recovery checkpoint does not continue through its exact successor decision",
2030 snapshot.version
2031 ));
2032 }
2033 let successor_pin = authority
2034 .assignment_handoff_checkpoint(&successor.target)
2035 .await
2036 .map_err(|error| error.to_string())?;
2037 if successor_pin.as_ref() != Some(&head_reference) {
2038 return Err(format!(
2039 "assignment {} recovery checkpoint is not pinned to its exact pending successor",
2040 snapshot.version
2041 ));
2042 }
2043 let (observed_outcome, observed_index) = authority
2044 .cluster_outcome_with_committed_checkpoint(head_reference.epoch)
2045 .await
2046 .map_err(|error| error.to_string())?
2047 .ok_or_else(|| {
2048 format!(
2049 "assignment {} propagated recovery Commit disappeared",
2050 snapshot.version
2051 )
2052 })?;
2053 if observed_outcome != committed_head {
2054 return Err(format!(
2055 "assignment {} recovery checkpoint head changed during authority audit",
2056 snapshot.version
2057 ));
2058 }
2059 let observed_index = observed_index.ok_or_else(|| {
2060 format!(
2061 "assignment {} propagated recovery Commit has no checkpoint index",
2062 snapshot.version
2063 )
2064 })?;
2065 let origin = validate_portable_recovery_cut(
2066 &observed_outcome,
2067 &observed_index,
2068 &head_reference,
2069 &decision.predecessor,
2070 "propagated assignment recovery",
2071 )?;
2072 AuditedRecoveryCheckpoint::Propagated {
2073 checkpoint: head_reference,
2074 origin,
2075 successor: successor.target,
2076 }
2077 };
2078 Ok((
2079 decision.predecessor,
2080 recovery_checkpoint,
2081 decision.leader_proof,
2082 ))
2083 })
2084}
2085
2086fn outcome_is_exact_abort_for_inventory(
2087 outcome: &CheckpointOutcome,
2088 inventory: &CheckpointArtifactInventory,
2089 predecessor: &CheckpointAssignmentFence,
2090) -> bool {
2091 outcome.scope == CheckpointScope::Cluster
2092 && outcome.epoch == inventory.attempt.epoch
2093 && outcome.checkpoint_id == inventory.attempt.checkpoint_id
2094 && outcome.deployment_id == inventory.deployment_id
2095 && outcome.assignment_fence.as_ref() == Some(predecessor)
2096 && outcome.verdict == CheckpointVerdict::Abort
2097 && outcome.committed_checkpoint.is_none()
2098 && outcome.leader_proof.is_some()
2099}
2100
2101fn abort_predecessor_checkpoint_for_recovery<'a>(
2106 store: &'a AssignmentSnapshotStore,
2107 controller: &'a ClusterController,
2108 snapshot: &'a AssignmentSnapshot,
2109 deadline: tokio::time::Instant,
2110) -> futures::future::BoxFuture<'a, Result<(), String>> {
2111 Box::pin(async move {
2112 if controller.capture_leader_proof().is_none() {
2115 return Ok(());
2116 }
2117 let authority = controller
2118 .checkpoint_authority()
2119 .map_err(|error| error.to_string())?;
2120 let (predecessor, audited_recovery, _) = tokio::time::timeout_at(
2121 deadline,
2122 audit_materialized_recovery_with_authority(store, Some(authority.as_ref()), snapshot),
2123 )
2124 .await
2125 .map_err(|_| {
2126 format!(
2127 "assignment {} recovery authority audit timed out while settling its predecessor checkpoint",
2128 snapshot.version
2129 )
2130 })??;
2131 settle_audited_recovery_predecessor_checkpoint(
2132 controller,
2133 authority,
2134 snapshot,
2135 predecessor,
2136 audited_recovery,
2137 deadline,
2138 )
2139 .await
2140 })
2141}
2142
2143fn settle_audited_recovery_predecessor_checkpoint<'a>(
2147 controller: &'a ClusterController,
2148 authority: Arc<LeaderLeaseStore>,
2149 snapshot: &'a AssignmentSnapshot,
2150 predecessor: CheckpointAssignmentFence,
2151 audited_recovery: AuditedRecoveryCheckpoint,
2152 deadline: tokio::time::Instant,
2153) -> futures::future::BoxFuture<'a, Result<(), String>> {
2154 Box::pin(async move {
2155 let recovery_checkpoint = audited_recovery.checkpoint().clone();
2156 let target = snapshot
2157 .assignment_fence()
2158 .map_err(|error| error.to_string())?;
2159 let pin_target = audited_recovery
2160 .propagated_successor()
2161 .cloned()
2162 .unwrap_or_else(|| target.clone());
2163 let decision = tokio::time::timeout_at(
2164 deadline,
2165 authority.assignment_recovery_decision(snapshot.version),
2166 )
2167 .await
2168 .map_err(|_| "recovery decision reload timed out before checkpoint settlement".to_string())?
2169 .map_err(|error| error.to_string())?
2170 .ok_or_else(|| {
2171 format!(
2172 "assignment {} lost its recovery decision before checkpoint settlement",
2173 snapshot.version
2174 )
2175 })?;
2176 if decision.predecessor != predecessor
2177 || decision.target != target
2178 || (!audited_recovery.was_consumed()
2179 && decision.recovery_checkpoint != recovery_checkpoint)
2180 {
2181 return Err(format!(
2182 "assignment {} recovery decision changed during checkpoint settlement",
2183 snapshot.version
2184 ));
2185 }
2186
2187 let Some(inventory) =
2188 tokio::time::timeout_at(deadline, authority.cluster_checkpoint_artifacts())
2189 .await
2190 .map_err(|_| "recovery checkpoint artifact inventory read timed out".to_string())?
2191 .map_err(|error| error.to_string())?
2192 else {
2193 return Ok(());
2194 };
2195 inventory.validate()?;
2196 if inventory.assignment_fence.as_ref() == Some(&target) {
2200 return Ok(());
2201 }
2202 if inventory.assignment_fence.as_ref() != Some(&predecessor) {
2203 return Err(format!(
2204 "assignment {} recovery found checkpoint {} artifacts for an unrelated assignment",
2205 snapshot.version, inventory.attempt.checkpoint_id
2206 ));
2207 }
2208
2209 if audited_recovery.was_consumed() {
2210 return Err(format!(
2211 "assignment {} recovery target already consumed its checkpoint pin while predecessor artifacts remain",
2212 snapshot.version
2213 ));
2214 }
2215
2216 let pinned = tokio::time::timeout_at(
2217 deadline,
2218 authority.assignment_handoff_checkpoint(&pin_target),
2219 )
2220 .await
2221 .map_err(|_| "recovery handoff checkpoint lookup timed out".to_string())?
2222 .map_err(|error| error.to_string())?;
2223 if pinned.as_ref() != Some(&recovery_checkpoint) {
2224 return Err(format!(
2225 "assignment {} recovery lost its exact target-keyed checkpoint pin",
2226 snapshot.version
2227 ));
2228 }
2229 let committed_head =
2230 tokio::time::timeout_at(deadline, authority.highest_cluster_committed_outcome())
2231 .await
2232 .map_err(|_| "recovery committed-head lookup timed out".to_string())?
2233 .map_err(|error| error.to_string())?
2234 .ok_or_else(|| {
2235 "recovery assignment has no committed checkpoint head".to_string()
2236 })?;
2237 if committed_head.committed_checkpoint.as_ref() != Some(&recovery_checkpoint) {
2238 return Err(format!(
2239 "assignment {} recovery pin is not the authoritative committed checkpoint head",
2240 snapshot.version
2241 ));
2242 }
2243 let (observed_outcome, observed_index) = tokio::time::timeout_at(
2244 deadline,
2245 authority.cluster_outcome_with_committed_checkpoint(recovery_checkpoint.epoch),
2246 )
2247 .await
2248 .map_err(|_| "recovery checkpoint head load timed out".to_string())?
2249 .map_err(|error| error.to_string())?
2250 .ok_or_else(|| "recovery checkpoint head disappeared".to_string())?;
2251 if observed_outcome != committed_head {
2252 return Err(format!(
2253 "assignment {} recovery checkpoint head changed during predecessor settlement",
2254 snapshot.version
2255 ));
2256 }
2257 let recovery_index = observed_index
2258 .ok_or_else(|| "recovery checkpoint Commit has no committed index".to_string())?;
2259 let observed_origin = validate_portable_recovery_cut(
2260 &observed_outcome,
2261 &recovery_index,
2262 &recovery_checkpoint,
2263 &predecessor,
2264 "predecessor checkpoint settlement",
2265 )?;
2266 if &observed_origin != audited_recovery.origin() {
2267 return Err(format!(
2268 "assignment {} recovery checkpoint origin changed during predecessor settlement",
2269 snapshot.version
2270 ));
2271 }
2272 if inventory.deployment_id != recovery_index.deployment_id
2273 || inventory.pipeline_identity != recovery_index.pipeline_identity
2274 {
2275 return Err(format!(
2276 "assignment {} recovery checkpoint and active predecessor artifacts are not one pipeline cut",
2277 snapshot.version
2278 ));
2279 }
2280 let recovery_attempt =
2281 CheckpointAttempt::new(recovery_checkpoint.epoch, recovery_checkpoint.checkpoint_id);
2282 if inventory.attempt.relation_to(recovery_attempt) != CheckpointAttemptRelation::Newer {
2283 return Err(format!(
2284 "assignment {} recovery cannot abort checkpoint {} because it does not follow pinned checkpoint {}",
2285 snapshot.version,
2286 inventory.attempt.checkpoint_id,
2287 recovery_checkpoint.checkpoint_id
2288 ));
2289 }
2290
2291 let existing = tokio::time::timeout_at(
2292 deadline,
2293 authority.cluster_attempt_settlement(inventory.attempt),
2294 )
2295 .await
2296 .map_err(|_| "predecessor checkpoint settlement read timed out".to_string())?
2297 .map_err(|error| error.to_string())?;
2298 let winner = if let Some(outcome) = existing {
2299 outcome
2300 } else {
2301 let current_proof = controller.capture_leader_proof().ok_or_else(|| {
2302 "leadership changed before predecessor checkpoint Abort publication".to_string()
2303 })?;
2304 let recorded = tokio::time::timeout_at(
2305 deadline,
2306 authority.record_cluster_outcome(
2307 ¤t_proof,
2308 inventory.attempt.epoch,
2309 inventory.attempt.checkpoint_id,
2310 predecessor.clone(),
2311 CheckpointVerdict::Abort,
2312 None,
2313 ),
2314 )
2315 .await
2316 .map_err(|_| "predecessor checkpoint Abort publication timed out".to_string())?
2317 .map_err(|error| error.to_string())?;
2318 match recorded {
2319 RecordOutcomeResult::Created(outcome)
2320 | RecordOutcomeResult::Unchanged(outcome)
2321 | RecordOutcomeResult::Conflict { winner: outcome } => outcome,
2322 }
2323 };
2324 if !outcome_is_exact_abort_for_inventory(&winner, &inventory, &predecessor) {
2325 return Err(format!(
2326 "checkpoint {} has incompatible terminal authority after assignment {} recovery",
2327 inventory.attempt.checkpoint_id, snapshot.version
2328 ));
2329 }
2330 let settled = tokio::time::timeout_at(
2331 deadline,
2332 authority.cluster_attempt_settlement(inventory.attempt),
2333 )
2334 .await
2335 .map_err(|_| "predecessor checkpoint Abort verification timed out".to_string())?
2336 .map_err(|error| error.to_string())?
2337 .ok_or_else(|| "predecessor checkpoint Abort was not durably visible".to_string())?;
2338 if !outcome_is_exact_abort_for_inventory(&settled, &inventory, &predecessor) {
2339 return Err(format!(
2340 "checkpoint {} Abort verification observed incompatible terminal authority",
2341 inventory.attempt.checkpoint_id
2342 ));
2343 }
2344 let final_pin = tokio::time::timeout_at(
2345 deadline,
2346 authority.assignment_handoff_checkpoint(&pin_target),
2347 )
2348 .await
2349 .map_err(|_| "post-Abort recovery handoff checkpoint lookup timed out".to_string())?
2350 .map_err(|error| error.to_string())?;
2351 if final_pin.as_ref() != Some(&recovery_checkpoint) {
2352 return Err(format!(
2353 "assignment {} recovery checkpoint pin changed while publishing predecessor Abort",
2354 snapshot.version
2355 ));
2356 }
2357 let remaining = tokio::time::timeout_at(deadline, authority.cluster_checkpoint_artifacts())
2358 .await
2359 .map_err(|_| "post-Abort checkpoint artifact inventory read timed out".to_string())?
2360 .map_err(|error| error.to_string())?;
2361 if remaining
2362 .as_ref()
2363 .is_some_and(|current| current != &inventory)
2364 {
2365 return Err(format!(
2366 "checkpoint {} artifact inventory changed while publishing its recovery Abort",
2367 inventory.attempt.checkpoint_id
2368 ));
2369 }
2370 info!(
2371 checkpoint_id = inventory.attempt.checkpoint_id,
2372 predecessor_version = predecessor.assignment_version,
2373 target_version = snapshot.version,
2374 "recovery assignment published the predecessor checkpoint Abort"
2375 );
2376 Ok(())
2377 })
2378}
2379
2380async fn audit_materialized_drain_transition(
2381 store: &AssignmentSnapshotStore,
2382 authority: Option<&LeaderLeaseStore>,
2383 snapshot: &AssignmentSnapshot,
2384 transition: AssignmentDrainTransition,
2385) -> Result<AuditedDrainOutcome, String> {
2386 let authority = authority
2387 .ok_or_else(|| "materialized assignment drain has no cluster authority".to_string())?;
2388 let decision = authority
2389 .assignment_drain_decision(snapshot.version)
2390 .await
2391 .map_err(|error| error.to_string())?
2392 .ok_or_else(|| {
2393 format!(
2394 "assignment {} has a materialized drain outcome without an authority decision",
2395 snapshot.version
2396 )
2397 })?;
2398 if decision.transition != transition {
2399 return Err(format!(
2400 "assignment {} materialization binds a different authority transition",
2401 snapshot.version
2402 ));
2403 }
2404 let predecessor_version = snapshot
2405 .version
2406 .checked_sub(1)
2407 .ok_or_else(|| "materialized drain has no predecessor generation".to_string())?;
2408 let predecessor = store
2409 .load_version(predecessor_version)
2410 .await
2411 .map_err(|error| error.to_string())?
2412 .ok_or_else(|| {
2413 format!(
2414 "materialized drain assignment {} lost predecessor {predecessor_version}",
2415 snapshot.version
2416 )
2417 })?;
2418 if predecessor.draining
2419 || predecessor
2420 .assignment_fence()
2421 .map_err(|error| error.to_string())?
2422 != transition.predecessor
2423 {
2424 return Err(format!(
2425 "assignment {} drain transition does not bind its exact committed predecessor",
2426 snapshot.version
2427 ));
2428 }
2429 let observed = finalized_drain_outcome(&transition, snapshot)?;
2430 let expected = match decision.verdict {
2431 AssignmentDrainVerdict::Commit => SourceDrainOutcome::Commit,
2432 AssignmentDrainVerdict::Abort => SourceDrainOutcome::Abort,
2433 };
2434 if observed != expected {
2435 return Err(format!(
2436 "assignment {} materialization conflicts with its authority decision",
2437 snapshot.version
2438 ));
2439 }
2440 Ok(AuditedDrainOutcome {
2441 transition,
2442 outcome: observed,
2443 handoff_checkpoint: decision.handoff_checkpoint,
2444 })
2445}
2446
2447#[derive(Clone, Copy)]
2451enum SourceDrainResolutionDeadline {
2452 Fresh(Duration),
2453 Absolute(tokio::time::Instant),
2454}
2455
2456impl SourceDrainResolutionDeadline {
2457 fn resolve(self) -> tokio::time::Instant {
2458 match self {
2459 Self::Fresh(timeout) => tokio::time::Instant::now() + timeout,
2460 Self::Absolute(deadline) => deadline,
2461 }
2462 }
2463}
2464
2465async fn settle_observed_local_drain(
2466 db: &Arc<LaminarDB>,
2467 store: &AssignmentSnapshotStore,
2468 registry: &VnodeRegistry,
2469 controller: Option<&ClusterController>,
2470 transition: &AssignmentDrainTransition,
2471 observed: &AssignmentSnapshot,
2472 audited_observed: Option<&AuditedDrainOutcome>,
2473 adoption_deadline: tokio::time::Instant,
2474 drain_deadline: SourceDrainResolutionDeadline,
2475) -> Result<bool, String> {
2476 if observed.draining {
2477 if observed.drain_transition.as_ref() == Some(transition) {
2478 return Ok(false);
2479 }
2480 if observed.version >= transition.target.assignment_version {
2481 return Err(format!(
2482 "drain {:?} was superseded by a different draining assignment {}",
2483 transition.id(),
2484 observed.version
2485 ));
2486 }
2487 return Ok(false);
2488 }
2489 if observed.version < transition.target.assignment_version {
2490 return Ok(false);
2491 }
2492
2493 let finalized = if observed.version == transition.target.assignment_version {
2494 observed.clone()
2495 } else {
2496 tokio::time::timeout_at(
2497 adoption_deadline,
2498 store.load_version(transition.target.assignment_version),
2499 )
2500 .await
2501 .map_err(|_| {
2502 format!(
2503 "timed out loading terminal assignment {} for drain resolution",
2504 transition.target.assignment_version
2505 )
2506 })?
2507 .map_err(|error| error.to_string())?
2508 .ok_or_else(|| {
2509 format!(
2510 "terminal assignment {} was pruned before local drain resolution",
2511 transition.target.assignment_version
2512 )
2513 })?
2514 };
2515 let audited = match audited_observed
2516 .filter(|audited| {
2517 finalized.version == observed.version && audited.transition == *transition
2518 })
2519 .cloned()
2520 {
2521 Some(audited) => audited,
2522 None => tokio::time::timeout_at(
2523 adoption_deadline,
2524 audit_assignment_snapshot_authority_outcome(store, controller, &finalized),
2525 )
2526 .await
2527 .map_err(|_| {
2528 format!(
2529 "timed out auditing terminal assignment {} for drain resolution",
2530 finalized.version
2531 )
2532 })??
2533 .into_terminal()
2534 .ok_or_else(|| {
2535 format!(
2536 "terminal assignment {} lost drain transition history",
2537 finalized.version
2538 )
2539 })?,
2540 };
2541 if audited.transition != *transition {
2542 return Err(format!(
2543 "terminal assignment {} binds a different drain transition",
2544 finalized.version
2545 ));
2546 }
2547 let outcome = audited.outcome;
2548 let target_version = transition.target.assignment_version;
2549 let local_version = registry.assignment_version();
2550 if local_version < target_version {
2551 db.adopt_assignment_snapshot(finalized.clone(), adoption_deadline)
2552 .await
2553 .map_err(|error| error.to_string())?;
2554 }
2555 if registry.assignment_version() != target_version {
2556 return Err(format!(
2557 "cannot resolve drain {:?} at local assignment {}",
2558 transition.id(),
2559 registry.assignment_version()
2560 ));
2561 }
2562 let expected = finalized
2563 .to_vnode_vec(registry.vnode_count())
2564 .map_err(|error| error.to_string())?;
2565 if registry.snapshot().as_ref() != expected.as_slice() {
2566 return Err(format!(
2567 "local assignment {target_version} does not match the durable drain outcome"
2568 ));
2569 }
2570 db.resolve_local_source_drain(transition.id(), outcome, drain_deadline.resolve())
2571 .await
2572 .map_err(|error| error.to_string())?;
2573 Ok(true)
2574}
2575
2576fn clear_settled_source_drain(
2577 controller: &ClusterController,
2578 transition: &AssignmentDrainTransition,
2579) -> Result<(), String> {
2580 match controller.checkpoint_drain_transition() {
2581 Some(active) if active == *transition => controller
2582 .clear_checkpoint_drain_transition_if_matches(transition)
2583 .then_some(())
2584 .ok_or_else(|| "process-local source drain changed during settlement".into()),
2585 Some(_) => Err("process-local source drain changed during settlement".into()),
2586 None => Ok(()),
2587 }
2588}
2589
2590pub(crate) async fn settle_source_drain_before_recovery_release(
2593 db: &Arc<LaminarDB>,
2594 controller: &ClusterController,
2595 expected_fence: &CheckpointAssignmentFence,
2596 deadline: tokio::time::Instant,
2597) -> Result<Option<u64>, String> {
2598 if !controller.is_recovering() || !db.cluster_intake_fenced() {
2599 return Err("recovery source-drain settlement requires closed intake".into());
2600 }
2601 let published_transition = controller.checkpoint_drain_transition();
2602 let store =
2603 db.assignment_snapshot_store.lock().clone().ok_or_else(|| {
2604 "recovery source-drain settlement has no assignment store".to_string()
2605 })?;
2606 let snapshot = tokio::time::timeout_at(deadline, store.load())
2607 .await
2608 .map_err(|_| "recovery source-drain head read timed out".to_string())?
2609 .map_err(|error| error.to_string())?
2610 .ok_or_else(|| "recovery source-drain settlement has no assignment head".to_string())?;
2611 if snapshot.draining {
2612 return Err(format!(
2613 "recovery Release reached unresolved draining assignment {}",
2614 snapshot.version
2615 ));
2616 }
2617 let fence = snapshot
2618 .assignment_fence()
2619 .map_err(|error| error.to_string())?;
2620 if &fence != expected_fence {
2621 return Err(format!(
2622 "recovery assignment {} changed before source-drain settlement",
2623 snapshot.version
2624 ));
2625 }
2626 let audited = tokio::time::timeout_at(
2627 deadline,
2628 audit_assignment_snapshot_authority_outcome(&store, Some(controller), &snapshot),
2629 )
2630 .await
2631 .map_err(|_| "recovery source-drain terminal audit timed out".to_string())??;
2632 let Some(audited) = audited.into_terminal() else {
2633 if published_transition.is_some() {
2634 return Err("process-local source drain has no durable terminal".into());
2635 }
2636 return Ok(None);
2637 };
2638 if published_transition.is_some_and(|active| active != audited.transition) {
2639 return Err("process-local source drain conflicts with the durable terminal".into());
2640 }
2641 if local_drain_participant(controller, &audited.transition).is_none() {
2642 clear_settled_source_drain(controller, &audited.transition)?;
2643 return Ok(None);
2644 }
2645 if tokio::time::Instant::now() >= deadline {
2646 return Err("recovery source-drain settlement deadline expired".into());
2647 }
2648 let registry = db
2649 .vnode_registry
2650 .lock()
2651 .clone()
2652 .ok_or_else(|| "recovery source-drain settlement has no vnode registry".to_string())?;
2653 if !settle_observed_local_drain(
2654 db,
2655 &store,
2656 ®istry,
2657 Some(controller),
2658 &audited.transition,
2659 &snapshot,
2660 Some(&audited),
2661 deadline,
2662 SourceDrainResolutionDeadline::Absolute(deadline),
2663 )
2664 .await?
2665 {
2666 return Err("materialized source drain was not ready for recovery Release".into());
2667 }
2668 let confirmed = tokio::time::timeout_at(deadline, store.load())
2669 .await
2670 .map_err(|_| "recovery source-drain head recheck timed out".to_string())?
2671 .map_err(|error| error.to_string())?
2672 .ok_or_else(|| "recovery source-drain assignment disappeared".to_string())?;
2673 if confirmed != snapshot {
2674 return Err("recovery source-drain assignment changed during settlement".into());
2675 }
2676 let confirmed_terminal = tokio::time::timeout_at(
2677 deadline,
2678 audit_assignment_snapshot_authority_outcome(&store, Some(controller), &confirmed),
2679 )
2680 .await
2681 .map_err(|_| "recovery source-drain terminal recheck timed out".to_string())??;
2682 if confirmed_terminal.terminal() != Some(&audited) {
2683 return Err("recovery source-drain terminal changed during settlement".into());
2684 }
2685 clear_settled_source_drain(controller, &audited.transition)?;
2686 Ok(Some(snapshot.version))
2687}
2688
2689pub(crate) async fn audit_drain_predecessor(
2690 store: &AssignmentSnapshotStore,
2691 registry: &VnodeRegistry,
2692 draining: &AssignmentSnapshot,
2693 controller: Option<&ClusterController>,
2694) -> Result<AssignmentSnapshot, String> {
2695 if !draining.draining {
2696 return Err("drain predecessor audit requires a draining head".into());
2697 }
2698 let local_version = registry.assignment_version();
2699 let expected_target = local_version
2700 .checked_add(1)
2701 .ok_or_else(|| "assignment version overflow during drain audit".to_string())?;
2702 if draining.version != expected_target {
2703 return Err(format!(
2704 "draining assignment {} is not the exact successor of local assignment {local_version}",
2705 draining.version
2706 ));
2707 }
2708 draining
2709 .to_vnode_vec(registry.vnode_count())
2710 .map_err(|error| error.to_string())?;
2711
2712 let predecessor = store
2713 .load_version(local_version)
2714 .await
2715 .map_err(|error| error.to_string())?
2716 .ok_or_else(|| {
2717 format!(
2718 "draining assignment {} lost committed predecessor {local_version}",
2719 draining.version
2720 )
2721 })?;
2722 if predecessor.draining || predecessor.version != local_version {
2723 return Err(format!(
2724 "assignment {local_version} is not the committed predecessor of draining assignment {}",
2725 draining.version
2726 ));
2727 }
2728 let transition = draining
2729 .drain_transition
2730 .as_ref()
2731 .ok_or_else(|| "draining assignment has no exact transition".to_string())?;
2732 if predecessor
2733 .assignment_fence()
2734 .map_err(|error| error.to_string())?
2735 != transition.predecessor
2736 {
2737 return Err(format!(
2738 "draining assignment {} does not bind committed predecessor {local_version}",
2739 draining.version
2740 ));
2741 }
2742 let controller = controller
2743 .ok_or_else(|| "draining assignment requires a cluster controller".to_string())?;
2744 if !controller.process_lease_is_live() {
2745 return Err("local process lease expired while auditing assignment drain".into());
2746 }
2747 let authority = controller
2748 .checkpoint_authority()
2749 .map_err(|error| error.to_string())?;
2750 if authority
2751 .assignment_drain_decision(transition.target.assignment_version)
2752 .await
2753 .map_err(|error| error.to_string())?
2754 .is_some()
2755 {
2756 return Err("assignment drain already has a terminal authority decision".into());
2757 }
2758 let current_leader = authority
2759 .load()
2760 .await
2761 .map_err(|error| error.to_string())?
2762 .ok_or_else(|| "assignment drain leader authority is missing".to_string())?;
2763 if !current_leader.matches_proof(&transition.leader) {
2764 return Err("assignment drain leader term was superseded".into());
2765 }
2766 let owners = predecessor
2767 .to_vnode_vec(registry.vnode_count())
2768 .map_err(|error| error.to_string())?;
2769 let local = registry.snapshot();
2770 if owners.as_slice() != local.as_ref() {
2771 return Err(format!(
2772 "committed predecessor {local_version} does not match the local owner map"
2773 ));
2774 }
2775 Ok(predecessor)
2776}
2777
2778async fn audit_exact_drain_head(
2779 store: &AssignmentSnapshotStore,
2780 registry: &VnodeRegistry,
2781 draining: &AssignmentSnapshot,
2782 controller: &ClusterController,
2783 deadline: tokio::time::Instant,
2784) -> Result<AssignmentSnapshot, String> {
2785 let refreshed = tokio::time::timeout_at(deadline, store.load())
2786 .await
2787 .map_err(|_| "durable drain re-audit timed out".to_string())?
2788 .map_err(|error| error.to_string())?;
2789 if refreshed.as_ref() != Some(draining) {
2790 return Err("durable assignment drain is no longer the unfinalized head".into());
2791 }
2792 tokio::time::timeout_at(
2793 deadline,
2794 audit_drain_predecessor(store, registry, draining, Some(controller)),
2795 )
2796 .await
2797 .map_err(|_| "durable drain predecessor audit timed out".to_string())?
2798}
2799
2800async fn prepare_and_announce_local_drain(
2801 db: &LaminarDB,
2802 store: &AssignmentSnapshotStore,
2803 registry: &VnodeRegistry,
2804 controller: &ClusterController,
2805 draining: &AssignmentSnapshot,
2806 transition: &AssignmentDrainTransition,
2807 participant: CheckpointParticipant,
2808 deadline: tokio::time::Instant,
2809) -> Result<(), String> {
2810 db.prepare_local_source_drain(transition, participant, deadline)
2811 .await
2812 .map_err(|error| error.to_string())?;
2813
2814 audit_exact_drain_head(store, registry, draining, controller, deadline).await?;
2817 if controller.checkpoint_drain_transition().as_ref() != Some(transition) {
2818 return Err(
2819 "source-drain transition changed before durable acknowledgement publication".into(),
2820 );
2821 }
2822 tokio::time::timeout_at(deadline, controller.announce_drain_ack(transition))
2823 .await
2824 .map_err(|_| "durable drain acknowledgement publication timed out".to_string())?
2825}
2826
2827fn publish_placement_metrics(
2829 metrics: &EngineMetrics,
2830 registry: &VnodeRegistry,
2831 nodes: &[(NodeId, Locality)],
2832 isolation_tier: usize,
2833) {
2834 let owners = registry.snapshot();
2835 let total = u32::try_from(owners.len().max(1)).unwrap_or(u32::MAX);
2836 let counts = owners_per_domain(&owners, nodes, isolation_tier);
2837
2838 metrics.placement_vnodes_per_domain.reset();
2839 let mut max = 0u32;
2840 for (domain, &count) in &counts {
2841 let label = if domain.is_empty() {
2842 "unknown"
2843 } else {
2844 domain.as_str()
2845 };
2846 metrics
2847 .placement_vnodes_per_domain
2848 .with_label_values(&[label])
2849 .set(i64::from(count));
2850 max = max.max(count);
2851 }
2852 metrics
2853 .placement_blast_radius_ratio
2854 .set(f64::from(max) / f64::from(total));
2855}
2856
2857pub fn spawn_rebalance_controller(
2860 db: Arc<LaminarDB>,
2861 controller: Arc<ClusterController>,
2862 store: Arc<AssignmentSnapshotStore>,
2863 registry: Arc<VnodeRegistry>,
2864 shutdown: CancellationToken,
2865 config: RebalanceConfig,
2866) -> JoinHandle<()> {
2867 tokio::spawn(async move {
2868 let mut members = controller.members_watch();
2869 let mut audit = tokio::time::interval(config.watcher_poll);
2870 audit.set_missed_tick_behavior(MissedTickBehavior::Skip);
2871 loop {
2872 let membership_changed = tokio::select! {
2873 biased;
2874 () = shutdown.cancelled() => return,
2875 res = members.changed() => {
2876 if res.is_err() {
2877 warn!("membership watch sender dropped; rebalance controller exiting");
2878 return;
2879 }
2880 true
2881 }
2882 _ = audit.tick() => false,
2883 };
2884
2885 if membership_changed {
2886 debug!("membership change observed; debouncing");
2887 loop {
2888 tokio::select! {
2889 biased;
2890 () = shutdown.cancelled() => return,
2891 res = tokio::time::timeout(
2892 config.rebalance_debounce, members.changed()
2893 ) => {
2894 match res {
2895 Ok(Ok(())) => {} Ok(Err(_)) => return, Err(_) => break, }
2899 }
2900 }
2901 }
2902 }
2903
2904 loop {
2905 if !controller.is_leader() {
2906 debug!("membership changed; not the leader — skipping rotation check");
2907 break;
2908 }
2909 let live = controller.assignable_instances();
2911 match try_rebalance_owned(
2912 Arc::clone(&db),
2913 Arc::clone(&controller),
2914 Arc::clone(&store),
2915 Arc::clone(®istry),
2916 live,
2917 config,
2918 )
2919 .await
2920 {
2921 Ok(Some(v)) => {
2922 info!(version = v, "rotated assignment");
2923 break;
2924 }
2925 Ok(None) => {
2926 debug!("live set matches current snapshot; no rotation");
2927 break;
2928 }
2929 Err(e) => {
2930 warn!(error = %e, "rebalance failed; retrying after backoff");
2931 tokio::select! {
2932 biased;
2933 () = shutdown.cancelled() => return,
2934 () = tokio::time::sleep(config.retry_delay) => {}
2935 }
2936 }
2937 }
2938 }
2939 }
2940 })
2941}
2942
2943pub async fn wait_until_drained(
2949 store: &AssignmentSnapshotStore,
2950 authority: Option<&LeaderLeaseStore>,
2951 me: NodeId,
2952 vnode_count: u32,
2953 poll: Duration,
2954 deadline: Duration,
2955) -> bool {
2956 let start = tokio::time::Instant::now();
2957 loop {
2958 if start.elapsed() >= deadline {
2959 return false;
2960 }
2961 let remaining = deadline.saturating_sub(start.elapsed());
2962 match tokio::time::timeout(remaining, store.load()).await {
2963 Ok(Ok(None)) => {
2964 warn!(
2965 "wait_until_drained: durable assignment head is missing; ownership is unknown"
2966 );
2967 }
2968 Ok(Ok(Some(snap))) => {
2969 match audit_materialized_drain_with_authority(store, authority, &snap).await {
2970 Ok(_) => match snap.to_vnode_vec(vnode_count) {
2971 Ok(owners) if !snap.draining && !owners.contains(&me) => return true,
2972 Ok(_) => {}
2973 Err(error) => {
2974 warn!(%error, "wait_until_drained: snapshot cardinality mismatch");
2975 }
2976 },
2977 Err(error) => {
2978 warn!(%error, version = snap.version, "wait_until_drained: assignment authority audit failed");
2979 }
2980 }
2981 }
2982 Ok(Err(e)) => warn!(error = %e, "wait_until_drained: snapshot load failed"),
2983 Err(_) => {
2984 warn!("wait_until_drained: snapshot load exceeded the shutdown deadline");
2985 return false;
2986 }
2987 }
2988 let remaining = deadline.saturating_sub(start.elapsed());
2989 tokio::time::sleep(poll.min(remaining)).await;
2990 if start.elapsed() >= deadline {
2991 return false;
2992 }
2993 }
2994}
2995
2996fn materialize_recovery_decision<'a>(
2997 db: &'a Arc<LaminarDB>,
2998 store: &'a Arc<AssignmentSnapshotStore>,
2999 controller: &'a ClusterController,
3000 registry: &'a VnodeRegistry,
3001 decision: AssignmentRecoveryDecision,
3002 operation_timeout: Duration,
3003) -> futures::future::BoxFuture<'a, Result<Option<u64>, String>> {
3004 Box::pin(async move {
3005 let deadline = tokio::time::Instant::now() + operation_timeout;
3006 let _process_fences =
3007 close_local_assignment_authority(db, controller, &decision.target, &[], deadline)
3008 .await?;
3009 let proposal =
3010 tokio::time::timeout_at(deadline, store.load_recovery_proposal(&decision.proposal))
3011 .await
3012 .map_err(|_| {
3013 "recovery proposal load exceeded the materialization deadline".to_string()
3014 })?
3015 .map_err(|error| error.to_string())?;
3016 if proposal
3017 .assignment_fence()
3018 .map_err(|error| error.to_string())?
3019 != decision.target
3020 {
3021 return Err("recovery authority winner does not match its staged proposal".into());
3022 }
3023 local_recovery_assignment_scope(&proposal, controller)?;
3024 let authority = controller
3025 .checkpoint_authority()
3026 .map_err(|error| error.to_string())?;
3027 let durable = match tokio::time::timeout_at(
3028 deadline,
3029 authority.materialize_assignment_recovery(decision.target_version()),
3030 )
3031 .await
3032 .map_err(|_| "recovery assignment materialization exceeded its deadline".to_string())?
3033 .map_err(|error| error.to_string())?
3034 {
3035 RotateOutcome::Rotated => proposal.clone(),
3036 RotateOutcome::Conflict(winner) => *winner,
3037 };
3038 if durable != proposal {
3039 return Err(format!(
3040 "assignment {} materialization conflicts with the authority recovery winner",
3041 decision.target_version()
3042 ));
3043 }
3044 prepare_recovery_assignment_adoption(db, store, controller, registry, &durable, deadline)
3045 .await?;
3046 let version = durable.version;
3047 db.adopt_recovery_assignment_snapshot(durable, operation_timeout)
3048 .await
3049 .map_err(|error| error.to_string())?;
3050 let oldest_retained = version.saturating_sub(1);
3051 let maintenance_store = Arc::clone(store);
3052 let maintenance_authority = Arc::clone(&authority);
3053 let maintenance_proof = controller.capture_leader_proof();
3054 tokio::spawn(async move {
3055 match maintenance_store.prune_before(oldest_retained).await {
3056 Ok(()) => {
3057 if let Some(proof) = maintenance_proof {
3058 if let Err(error) = maintenance_authority
3059 .prune_assignment_drain_decisions_before(&proof, oldest_retained)
3060 .await
3061 {
3062 warn!(%error, "assignment recovery authority prune failed after snapshot prune");
3063 }
3064 }
3065 }
3066 Err(error) => warn!(%error, "snapshot prune failed after assignment recovery"),
3067 }
3068 });
3069 Ok(Some(version))
3070 })
3071}
3072
3073async fn reconcile_pending_recovery_decision(
3074 db: &Arc<LaminarDB>,
3075 store: &Arc<AssignmentSnapshotStore>,
3076 controller: &ClusterController,
3077 registry: &VnodeRegistry,
3078 current: &AssignmentSnapshot,
3079 operation_timeout: Duration,
3080) -> Result<Option<u64>, String> {
3081 let target_version = current
3082 .version
3083 .checked_add(1)
3084 .ok_or_else(|| "assignment version exhausted".to_string())?;
3085 let authority = controller
3086 .checkpoint_authority()
3087 .map_err(|error| error.to_string())?;
3088 let Some(decision) = tokio::time::timeout(
3089 operation_timeout,
3090 authority.assignment_recovery_decision(target_version),
3091 )
3092 .await
3093 .map_err(|_| "recovery authority lookup timed out".to_string())?
3094 .map_err(|error| error.to_string())?
3095 else {
3096 return Ok(None);
3097 };
3098 if decision.predecessor
3099 != current
3100 .assignment_fence()
3101 .map_err(|error| error.to_string())?
3102 {
3103 return Err(format!(
3104 "pending recovery decision for assignment {target_version} has the wrong predecessor"
3105 ));
3106 }
3107 materialize_recovery_decision(db, store, controller, registry, decision, operation_timeout)
3108 .await
3109}
3110
3111fn replaced_predecessor_processes(
3112 predecessor: &CheckpointAssignmentFence,
3113 target: &CheckpointAssignmentFence,
3114) -> Vec<CheckpointParticipant> {
3115 predecessor
3116 .participants
3117 .iter()
3118 .copied()
3119 .filter(|participant| {
3120 target.participant_incarnation(participant.node_id)
3121 != Some(participant.boot_incarnation)
3122 })
3123 .collect()
3124}
3125
3126fn authorize_recovery_successor<'a>(
3127 db: &'a Arc<LaminarDB>,
3128 store: &'a Arc<AssignmentSnapshotStore>,
3129 controller: &'a ClusterController,
3130 current: &'a AssignmentSnapshot,
3131 proposal: AssignmentSnapshot,
3132 operation_timeout: Duration,
3133 reason: &'a str,
3134) -> futures::future::BoxFuture<'a, Result<AssignmentRecoveryDecision, String>> {
3135 Box::pin(async move {
3136 let predecessor = current
3137 .assignment_fence()
3138 .map_err(|error| error.to_string())?;
3139 let target = proposal
3140 .assignment_fence()
3141 .map_err(|error| error.to_string())?;
3142 let deadline = controller.process_fencing_deadline(operation_timeout)?;
3143 let removed = replaced_predecessor_processes(&predecessor, &target);
3144 let process_fences =
3145 close_local_assignment_authority(db, controller, &target, &removed, deadline).await?;
3146 let proposal_ref =
3147 tokio::time::timeout_at(deadline, store.stage_recovery_proposal(&proposal))
3148 .await
3149 .map_err(|_| "recovery proposal staging exceeded the fencing deadline".to_string())?
3150 .map_err(|error| error.to_string())?;
3151 let observed = tokio::time::timeout_at(deadline, store.load())
3152 .await
3153 .map_err(|_| "assignment head revalidation exceeded the fencing deadline".to_string())?
3154 .map_err(|error| error.to_string())?
3155 .ok_or_else(|| "assignment head disappeared during recovery fencing".to_string())?;
3156 if observed != *current {
3157 return Err(format!(
3158 "assignment head advanced from {} while process fencing was in progress",
3159 current.version
3160 ));
3161 }
3162 tokio::time::timeout_at(
3163 deadline,
3164 audit_assignment_snapshot_authority(store, Some(controller), &observed),
3165 )
3166 .await
3167 .map_err(|_| "predecessor authority audit exceeded the fencing deadline".to_string())??;
3168
3169 let target_checks =
3170 futures::future::join_all(target.participants.iter().copied().map(|participant| {
3171 controller.verify_current_process_incarnation(participant, deadline)
3172 }));
3173 for (participant, check) in target.participants.iter().zip(target_checks.await) {
3174 if !check? {
3175 return Err(format!(
3176 "successor process {} is not the current durable lease owner",
3177 participant.node_id
3178 ));
3179 }
3180 }
3181
3182 let authority = controller
3183 .checkpoint_authority()
3184 .map_err(|error| error.to_string())?;
3185 let committed_head =
3186 tokio::time::timeout_at(deadline, authority.highest_cluster_committed_outcome())
3187 .await
3188 .map_err(|_| {
3189 "recovery checkpoint lookup exceeded the fencing deadline".to_string()
3190 })?
3191 .map_err(|error| error.to_string())?
3192 .ok_or_else(|| {
3193 "assignment recovery requires a committed cluster checkpoint".to_string()
3194 })?;
3195 let head_reference = committed_head.committed_checkpoint.clone().ok_or_else(|| {
3196 "latest committed cluster checkpoint has no committed index reference".to_string()
3197 })?;
3198 let pinned_checkpoint = tokio::time::timeout_at(
3199 deadline,
3200 authority.assignment_handoff_checkpoint(&predecessor),
3201 )
3202 .await
3203 .map_err(|_| "recovery handoff lookup exceeded the fencing deadline".to_string())?
3204 .map_err(|error| error.to_string())?;
3205 let recovery_checkpoint = if let Some(reference) = pinned_checkpoint {
3206 if reference != head_reference {
3207 return Err(
3208 "recovery handoff pin is not the exact latest committed checkpoint head".into(),
3209 );
3210 }
3211 reference
3212 } else {
3213 if committed_head.assignment_fence.as_ref() != Some(&predecessor) {
3214 return Err(
3215 "latest committed cluster checkpoint does not bind the recovery predecessor"
3216 .into(),
3217 );
3218 }
3219 head_reference
3220 };
3221 let (observed_outcome, observed_index) = tokio::time::timeout_at(
3222 deadline,
3223 authority.cluster_outcome_with_committed_checkpoint(recovery_checkpoint.epoch),
3224 )
3225 .await
3226 .map_err(|_| "recovery checkpoint head load exceeded the fencing deadline".to_string())?
3227 .map_err(|error| error.to_string())?
3228 .ok_or_else(|| "recovery checkpoint Commit disappeared during fencing".to_string())?;
3229 if observed_outcome != committed_head {
3230 return Err("recovery checkpoint head changed during successor authorization".into());
3231 }
3232 let recovery_index = observed_index
3233 .ok_or_else(|| "recovery checkpoint Commit has no committed index".to_string())?;
3234 validate_portable_recovery_cut(
3235 &observed_outcome,
3236 &recovery_index,
3237 &recovery_checkpoint,
3238 &predecessor,
3239 "assignment recovery authorization",
3240 )?;
3241 let leader_proof = controller.capture_leader_proof().ok_or_else(|| {
3242 "assignment recovery lost the current durable leader proof".to_string()
3243 })?;
3244 let decision = AssignmentRecoveryDecision::new(
3245 predecessor,
3246 target,
3247 proposal_ref,
3248 process_fences,
3249 recovery_checkpoint,
3250 leader_proof.clone(),
3251 )?;
3252 let decision = match tokio::time::timeout_at(
3253 deadline,
3254 controller.record_assignment_recovery_decision(&leader_proof, decision, deadline),
3255 )
3256 .await
3257 .map_err(|_| "recovery authority admission exceeded the fencing deadline".to_string())?
3258 .map_err(|error| error.clone())?
3259 {
3260 RecordAssignmentRecoveryDecisionResult::Created(decision)
3261 | RecordAssignmentRecoveryDecisionResult::Unchanged(decision) => decision,
3262 RecordAssignmentRecoveryDecisionResult::Conflict { winner } => winner,
3263 };
3264 warn!(
3265 predecessor_version = current.version,
3266 target_version = decision.target_version(),
3267 %reason,
3268 "authorized successor assignment from the last committed cluster cut"
3269 );
3270 Ok(decision)
3271 })
3272}
3273
3274enum DrainPublicationReconciliation {
3275 Deferred,
3276 Resolved {
3277 outcome: RotateOutcome,
3278 authority_changed: bool,
3279 },
3280}
3281
3282async fn reconcile_drain_publication(
3283 store: &AssignmentSnapshotStore,
3284 controller: &ClusterController,
3285 drain: &AssignmentSnapshot,
3286 prior_version: u64,
3287 retry_delay: Duration,
3288) -> Result<DrainPublicationReconciliation, SnapshotError> {
3289 let transition = drain.drain_transition.as_ref().ok_or_else(|| {
3290 SnapshotError::Invalid("drain publication has no exact transition".into())
3291 })?;
3292 let mut had_ambiguous_attempt = false;
3293 let mut authority_changed = false;
3294
3295 loop {
3296 let changed = controller.is_recovering() || !controller.proof_is_live(&transition.leader);
3297 authority_changed |= changed;
3298 if changed && !had_ambiguous_attempt {
3299 return Ok(DrainPublicationReconciliation::Deferred);
3300 }
3301
3302 match store.save_if_version(drain, prior_version).await {
3303 Ok(outcome) => {
3304 authority_changed |=
3305 controller.is_recovering() || !controller.proof_is_live(&transition.leader);
3306 let outcome = match outcome {
3307 RotateOutcome::Conflict(winner) if winner.as_ref() == drain => {
3308 RotateOutcome::Rotated
3309 }
3310 outcome => outcome,
3311 };
3312 return Ok(DrainPublicationReconciliation::Resolved {
3313 outcome,
3314 authority_changed,
3315 });
3316 }
3317 Err(SnapshotError::Io(error)) => {
3318 had_ambiguous_attempt = true;
3319 authority_changed |=
3320 controller.is_recovering() || !controller.proof_is_live(&transition.leader);
3321 warn!(
3322 version = drain.version,
3323 %error,
3324 "ambiguous assignment drain publication; reconciling the exact create"
3325 );
3326 tokio::time::sleep(retry_delay).await;
3327 }
3328 Err(error @ (SnapshotError::Invalid(_) | SnapshotError::Json(_))) => {
3329 return Err(error);
3330 }
3331 }
3332 }
3333}
3334
3335fn execute_graceful_rotation_owned(
3336 db: Arc<LaminarDB>,
3337 controller: Arc<ClusterController>,
3338 store: Arc<AssignmentSnapshotStore>,
3339 registry: Arc<VnodeRegistry>,
3340 current: AssignmentSnapshot,
3341 new_vnodes: std::collections::BTreeMap<u32, NodeId>,
3342 participants: Vec<CheckpointParticipant>,
3343 config: RebalanceConfig,
3344) -> futures::future::BoxFuture<'static, Result<Option<u64>, String>> {
3345 Box::pin(async move {
3346 let leader = controller.capture_leader_proof().ok_or_else(|| {
3347 "assignment drain requires the current durable leader proof".to_string()
3348 })?;
3349 let drain = current
3350 .next_draining(new_vnodes.clone(), participants.clone(), leader)
3351 .map_err(|error| error.to_string())?;
3352 let transition = drain
3353 .drain_transition
3354 .as_ref()
3355 .expect("validated draining snapshot has a transition")
3356 .clone();
3357 let current_owners = current
3358 .to_vnode_vec(registry.vnode_count())
3359 .map_err(|error| error.to_string())?;
3360 let publication_deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
3361 let rotate_outcome = {
3362 let _adoption = tokio::time::timeout_at(
3370 publication_deadline,
3371 db.assignment_adoption_lock.lock(),
3372 )
3373 .await
3374 .map_err(|_| {
3375 format!(
3376 "graceful rotation from assignment {} timed out waiting for assignment serialization",
3377 current.version
3378 )
3379 })?;
3380 let local = registry.versioned_snapshot();
3381 if local.version() != current.version || local.owners() != current_owners.as_slice() {
3382 return Err(format!(
3383 "graceful rotation predecessor assignment {} does not match local assignment {}",
3384 current.version,
3385 local.version()
3386 ));
3387 }
3388 let local_vnode_state_ready = tokio::time::timeout_at(
3389 publication_deadline,
3390 db.local_vnode_state_is_ready(®istry, &transition.predecessor),
3391 )
3392 .await
3393 .map_err(|_| {
3394 format!(
3395 "graceful rotation from assignment {} timed out checking local vnode-state readiness",
3396 current.version
3397 )
3398 })?
3399 .map_err(|error| error.to_string())?;
3400 if !local_vnode_state_ready {
3401 debug!(
3402 version = current.version,
3403 "graceful rotation deferred until local vnode state matches the predecessor"
3404 );
3405 return Ok(None);
3406 }
3407 let Some(published_fence) = controller.checkpoint_assignment_fence(current.version)
3408 else {
3409 debug!(
3410 version = current.version,
3411 "graceful rotation deferred until the predecessor assignment is adopted"
3412 );
3413 return Ok(None);
3414 };
3415 if published_fence != transition.predecessor {
3416 return Err(format!(
3417 "graceful rotation predecessor assignment {} does not match the exact assignment-adoption certificate",
3418 current.version
3419 ));
3420 }
3421 if !read_assignment_vnode_state_readiness(
3425 &controller,
3426 &transition.predecessor,
3427 publication_deadline,
3428 )
3429 .await?
3430 {
3431 debug!(
3432 version = current.version,
3433 "graceful rotation deferred until every participant installs vnode state"
3434 );
3435 return Ok(None);
3436 }
3437 let outcome = match reconcile_drain_publication(
3438 &store,
3439 &controller,
3440 &drain,
3441 current.version,
3442 config.retry_delay,
3443 )
3444 .await
3445 .map_err(|error| {
3446 format!(
3447 "assignment drain {} conditional publication failed deterministically: {error}",
3448 drain.version
3449 )
3450 })? {
3451 DrainPublicationReconciliation::Deferred => {
3452 debug!(
3453 version = current.version,
3454 "graceful rotation deferred after publication authority changed"
3455 );
3456 return Ok(None);
3457 }
3458 DrainPublicationReconciliation::Resolved {
3459 authority_changed: true,
3460 ..
3461 } => {
3462 return Err(format!(
3463 "assignment drain {} became durable after publication authority changed; yielding to durable drain recovery",
3464 drain.version
3465 ));
3466 }
3467 DrainPublicationReconciliation::Resolved {
3468 outcome,
3469 authority_changed: false,
3470 } => outcome,
3471 };
3472 if matches!(&outcome, RotateOutcome::Rotated) {
3473 db.validate_source_drain_snapshot(&drain)
3478 .map_err(|error| error.to_string())?;
3479 match controller.checkpoint_drain_transition() {
3480 None => {
3481 controller.publish_checkpoint_drain_transition(Some(transition.clone()));
3482 }
3483 Some(installed) if installed == transition => {}
3484 Some(_) => {
3485 return Err(
3486 "a different source-drain transition is already installed locally"
3487 .into(),
3488 );
3489 }
3490 }
3491 }
3492 outcome
3493 };
3494 match rotate_outcome {
3495 RotateOutcome::Rotated => {
3496 let drain_deadline = tokio::time::Instant::now() + config.drain_ack_timeout;
3497 let local_receipt = match local_drain_participant(&controller, &transition) {
3501 Some(participant) => {
3502 prepare_and_announce_local_drain(
3503 &db,
3504 &store,
3505 ®istry,
3506 &controller,
3507 &drain,
3508 &transition,
3509 participant,
3510 drain_deadline,
3511 )
3512 .await
3513 }
3514 None => Ok(()),
3515 };
3516 let acked = match local_receipt {
3517 Ok(()) => await_drain_quorum(&controller, &transition, drain_deadline).await,
3518 Err(error) => {
3519 warn!(%error, version = drain.version, "leader durable drain acknowledgement failed");
3520 false
3521 }
3522 };
3523 if !acked {
3524 let failure = "drain ack quorum not reached before timeout";
3525 if let Err(error) = audit_exact_drain_head(
3526 &store,
3527 ®istry,
3528 &drain,
3529 &controller,
3530 tokio::time::Instant::now() + config.checkpoint_timeout,
3531 )
3532 .await
3533 {
3534 return Err(format!(
3535 "{failure}; drain is no longer authoritative: {error}"
3536 ));
3537 }
3538 if let Err(recovery_error) = ensure_local_recovery_fault(&db, &controller).await
3544 {
3545 return Err(format!(
3546 "{failure}; coordinated recovery request failed: {recovery_error}"
3547 ));
3548 }
3549 return Err(format!("{failure}; coordinated recovery requested"));
3550 }
3551 audit_exact_drain_head(
3552 &store,
3553 ®istry,
3554 &drain,
3555 &controller,
3556 tokio::time::Instant::now() + config.checkpoint_timeout,
3557 )
3558 .await?;
3559 let handoff_checkpoint =
3562 match pre_rotation_checkpoint(&db, &controller, &transition, config).await {
3563 Ok(reference) => reference,
3564 Err(failure) => {
3565 if let Err(error) = audit_exact_drain_head(
3566 &store,
3567 ®istry,
3568 &drain,
3569 &controller,
3570 tokio::time::Instant::now() + config.checkpoint_timeout,
3571 )
3572 .await
3573 {
3574 return Err(format!(
3575 "{failure}; drain is no longer authoritative: {error}"
3576 ));
3577 }
3578 if let Err(abort_error) = finalize_drain_snapshot(
3579 &db,
3580 &store,
3581 &controller,
3582 &drain,
3583 ¤t,
3584 AssignmentDrainVerdict::Abort,
3585 None,
3586 config,
3587 DrainFinalizationMode::Live,
3588 )
3589 .await
3590 {
3591 return Err(format!(
3592 "{failure}; assignment drain abort failed: {abort_error}"
3593 ));
3594 }
3595 return Err(failure);
3596 }
3597 };
3598 audit_exact_drain_head(
3599 &store,
3600 ®istry,
3601 &drain,
3602 &controller,
3603 tokio::time::Instant::now() + config.checkpoint_timeout,
3604 )
3605 .await?;
3606 return finalize_drain_snapshot(
3607 &db,
3608 &store,
3609 &controller,
3610 &drain,
3611 ¤t,
3612 AssignmentDrainVerdict::Commit,
3613 Some(handoff_checkpoint),
3614 config,
3615 DrainFinalizationMode::Live,
3616 )
3617 .await;
3618 }
3619 RotateOutcome::Conflict(winner) => {
3620 let v = winner.version;
3621 adopt_any(
3622 &db,
3623 &store,
3624 &controller,
3625 ®istry,
3626 *winner,
3627 config.checkpoint_timeout,
3628 )
3629 .await?;
3630 Ok(Some(v))
3631 }
3632 }
3633 })
3634}
3635fn try_rebalance_owned(
3636 db: Arc<LaminarDB>,
3637 controller: Arc<ClusterController>,
3638 store: Arc<AssignmentSnapshotStore>,
3639 registry: Arc<VnodeRegistry>,
3640 live: Vec<NodeId>,
3641 config: RebalanceConfig,
3642) -> futures::future::BoxFuture<'static, Result<Option<u64>, String>> {
3643 Box::pin(async move {
3644 let head_deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
3645 let current = tokio::time::timeout_at(head_deadline, store.load())
3646 .await
3647 .map_err(|_| "durable assignment head audit timed out".to_string())?
3648 .map_err(|e| e.to_string())?
3649 .ok_or_else(|| "no snapshot on store — boot seed missing".to_string())?;
3650 let current_authority = tokio::time::timeout_at(
3651 head_deadline,
3652 audit_assignment_snapshot_authority_outcome_owned(
3653 Arc::clone(&store),
3654 Arc::clone(&controller),
3655 current.clone(),
3656 ),
3657 )
3658 .await
3659 .map_err(|_| "durable assignment authority audit timed out".to_string())??;
3660 if current_authority.is_recovery() {
3664 abort_predecessor_checkpoint_for_recovery(&store, &controller, ¤t, head_deadline)
3665 .await?;
3666 }
3667 let current_owners = current
3668 .to_vnode_vec(registry.vnode_count())
3669 .map_err(|error| error.to_string())?;
3670
3671 let local_assignment = registry.versioned_snapshot();
3672 validate_local_assignment_head(
3673 ¤t,
3674 ¤t_owners,
3675 local_assignment.version(),
3676 local_assignment.owners(),
3677 )?;
3678
3679 let may_reconcile_propagated_successor = local_assignment.version() >= current.version
3685 || DbState::load(&db.state) == DbState::Faulted;
3686 if let Some(successor) = current_authority
3687 .propagated_recovery_successor()
3688 .filter(|_| may_reconcile_propagated_successor)
3689 {
3690 let expected_version = successor.assignment_version;
3691 let materialized = reconcile_pending_recovery_decision_owned(
3692 Arc::clone(&db),
3693 Arc::clone(&store),
3694 Arc::clone(&controller),
3695 Arc::clone(®istry),
3696 current.clone(),
3697 config.checkpoint_timeout,
3698 )
3699 .await?
3700 .ok_or_else(|| {
3701 format!(
3702 "assignment {} propagated its checkpoint pin to missing recovery successor {expected_version}",
3703 current.version
3704 )
3705 })?;
3706 if materialized != expected_version {
3707 return Err(format!(
3708 "assignment {} propagated recovery materialized {materialized}, expected {expected_version}",
3709 current.version
3710 ));
3711 }
3712 return Ok(Some(materialized));
3713 }
3714 if current.draining {
3715 let prior_version = current.version.checked_sub(1).ok_or_else(|| {
3716 "draining assignment has no prior committed generation".to_string()
3717 })?;
3718 let prior = tokio::time::timeout_at(head_deadline, store.load_version(prior_version))
3719 .await
3720 .map_err(|_| {
3721 format!("committed predecessor {prior_version} load exceeded the head deadline")
3722 })?
3723 .map_err(|error| error.to_string())?
3724 .ok_or_else(|| {
3725 format!(
3726 "draining assignment {} lost committed predecessor {prior_version}",
3727 current.version
3728 )
3729 })?;
3730 if prior.draining {
3731 return Err(format!(
3732 "draining assignment {} has a non-committed predecessor",
3733 current.version
3734 ));
3735 }
3736 tokio::time::timeout_at(
3737 head_deadline,
3738 audit_assignment_snapshot_authority_owned(
3739 Arc::clone(&store),
3740 Arc::clone(&controller),
3741 prior.clone(),
3742 ),
3743 )
3744 .await
3745 .map_err(|_| {
3746 format!(
3747 "committed predecessor {prior_version} authority audit exceeded the head deadline"
3748 )
3749 })??;
3750 prior
3751 .to_vnode_vec(registry.vnode_count())
3752 .map_err(|error| error.to_string())?;
3753 suspend_local_assignment_authority(&db, Some(controller.as_ref()), head_deadline)
3759 .await?;
3760 return finalize_drain_snapshot(
3761 &db,
3762 &store,
3763 &controller,
3764 ¤t,
3765 &prior,
3766 AssignmentDrainVerdict::Abort,
3767 None,
3768 config,
3769 DrainFinalizationMode::Live,
3770 )
3771 .await;
3772 }
3773
3774 if live.is_empty() {
3777 return Ok(None);
3778 }
3779 let live_ids = successor_participant_ids(&live);
3780 let observed_processes = controller
3781 .available_recovery_participant_incarnations(&live_ids)
3782 .await?;
3783 let process_read_deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
3784 let process_checks =
3785 futures::future::join_all(observed_processes.iter().copied().map(|participant| {
3786 controller.verify_current_process_incarnation(participant, process_read_deadline)
3787 }))
3788 .await;
3789 let mut available_processes = Vec::with_capacity(observed_processes.len());
3790 for (participant, check) in observed_processes.into_iter().zip(process_checks) {
3791 if check? {
3792 available_processes.push(participant);
3793 }
3794 }
3795 let successor_processes = available_processes
3796 .iter()
3797 .copied()
3798 .filter(|participant| controller.admit_successor_process(*participant))
3799 .collect::<Vec<_>>();
3800 let successors: Vec<NodeId> = successor_processes
3801 .iter()
3802 .map(|participant| NodeId(participant.node_id))
3803 .collect();
3804 if successors.is_empty() {
3805 return Ok(None);
3806 }
3807 let current_roster_is_live = successor_participants(¤t_owners, &successor_processes)
3808 .is_ok_and(|participants| participants == current.participants);
3809 if current.version > local_assignment.version()
3810 && current_authority.is_recovery()
3811 && assignment_binds_local_process(¤t, &controller)?
3812 {
3813 return adopt_materialized_recovery_head(
3814 &db,
3815 &store,
3816 &controller,
3817 ®istry,
3818 ¤t,
3819 head_deadline,
3820 config.checkpoint_timeout,
3821 )
3822 .await
3823 .map(Some);
3824 }
3825 if current.version > local_assignment.version() && current_roster_is_live {
3826 db.adopt_assignment_snapshot(current.clone(), head_deadline)
3831 .await
3832 .map_err(|error| error.to_string())?;
3833 let reconciled_version = registry.assignment_version();
3834 if reconciled_version < current.version {
3835 return Err(format!(
3836 "durable assignment {} was not adopted; local assignment remains {}",
3837 current.version, reconciled_version
3838 ));
3839 }
3840 return Ok(Some(reconciled_version));
3841 }
3842 if let Some(version) = reconcile_pending_recovery_decision_owned(
3846 Arc::clone(&db),
3847 Arc::clone(&store),
3848 Arc::clone(&controller),
3849 Arc::clone(®istry),
3850 current.clone(),
3851 config.checkpoint_timeout,
3852 )
3853 .await?
3854 {
3855 return Ok(Some(version));
3856 }
3857 let mut new_assignment =
3858 rendezvous_assignment(registry.vnode_count(), &successors).to_vec();
3859 let mut new_vnodes = AssignmentSnapshot::vnodes_from_vec(&new_assignment);
3860 let mut participants = successor_participants(&new_assignment, &successor_processes)?;
3866 let roster_changed = current.participants != participants;
3867
3868 if new_vnodes == current.vnodes && !roster_changed {
3869 return Ok(None);
3870 }
3871
3872 let recovery_reason =
3873 predecessor_cut_unavailability(&controller, ¤t, ¤t_owners, &successors)
3874 .await;
3875 if let Some(reason) = recovery_reason {
3876 retain_recovery_predecessors(
3877 &mut new_assignment,
3878 ¤t.participants,
3879 &successor_processes,
3880 )?;
3881 new_vnodes = AssignmentSnapshot::vnodes_from_vec(&new_assignment);
3882 participants = successor_participants(&new_assignment, &successor_processes)?;
3883 let proposal = current
3884 .next_for_participants(new_vnodes, participants)
3885 .map_err(|error| error.to_string())?;
3886 return authorize_recovery_successor_owned(
3887 Arc::clone(&db),
3888 Arc::clone(&store),
3889 Arc::clone(&controller),
3890 Arc::clone(®istry),
3891 current.clone(),
3892 proposal,
3893 config.checkpoint_timeout,
3894 reason,
3895 )
3896 .await;
3897 }
3898
3899 execute_graceful_rotation_owned(
3900 db,
3901 controller,
3902 store,
3903 registry,
3904 current,
3905 new_vnodes,
3906 participants,
3907 config,
3908 )
3909 .await
3910 })
3911}
3912
3913fn validate_local_assignment_head(
3914 current: &AssignmentSnapshot,
3915 current_owners: &[NodeId],
3916 local_version: u64,
3917 local_owners: &[NodeId],
3918) -> Result<(), String> {
3919 if current.version < local_version {
3920 return Err(format!(
3921 "durable assignment head {} regressed behind local assignment {local_version}",
3922 current.version
3923 ));
3924 }
3925 if current.version == local_version && current_owners != local_owners {
3926 return Err(format!(
3927 "durable and local assignment {} have different owner maps",
3928 current.version
3929 ));
3930 }
3931 Ok(())
3932}
3933
3934#[cfg(test)]
3935async fn try_rebalance(
3936 db: &Arc<LaminarDB>,
3937 controller: &Arc<ClusterController>,
3938 store: &Arc<AssignmentSnapshotStore>,
3939 registry: &Arc<VnodeRegistry>,
3940 live: &[NodeId],
3941 config: RebalanceConfig,
3942) -> Result<Option<u64>, String> {
3943 try_rebalance_owned(
3944 Arc::clone(db),
3945 Arc::clone(controller),
3946 Arc::clone(store),
3947 Arc::clone(registry),
3948 live.to_vec(),
3949 config,
3950 )
3951 .await
3952}
3953
3954fn audit_assignment_snapshot_authority_owned(
3955 store: Arc<AssignmentSnapshotStore>,
3956 controller: Arc<ClusterController>,
3957 snapshot: AssignmentSnapshot,
3958) -> futures::future::BoxFuture<'static, Result<(), String>> {
3959 Box::pin(async move {
3960 audit_assignment_snapshot_authority(&store, Some(&controller), &snapshot).await
3961 })
3962}
3963
3964fn audit_assignment_snapshot_authority_outcome_owned(
3965 store: Arc<AssignmentSnapshotStore>,
3966 controller: Arc<ClusterController>,
3967 snapshot: AssignmentSnapshot,
3968) -> futures::future::BoxFuture<'static, Result<AuditedAssignmentAuthority, String>> {
3969 Box::pin(async move {
3970 audit_assignment_snapshot_authority_outcome(&store, Some(&controller), &snapshot).await
3971 })
3972}
3973
3974fn reconcile_pending_recovery_decision_owned(
3975 db: Arc<LaminarDB>,
3976 store: Arc<AssignmentSnapshotStore>,
3977 controller: Arc<ClusterController>,
3978 registry: Arc<VnodeRegistry>,
3979 current: AssignmentSnapshot,
3980 operation_timeout: Duration,
3981) -> futures::future::BoxFuture<'static, Result<Option<u64>, String>> {
3982 Box::pin(async move {
3983 reconcile_pending_recovery_decision(
3984 &db,
3985 &store,
3986 &controller,
3987 ®istry,
3988 ¤t,
3989 operation_timeout,
3990 )
3991 .await
3992 })
3993}
3994
3995fn authorize_recovery_successor_owned(
3996 db: Arc<LaminarDB>,
3997 store: Arc<AssignmentSnapshotStore>,
3998 controller: Arc<ClusterController>,
3999 registry: Arc<VnodeRegistry>,
4000 current: AssignmentSnapshot,
4001 proposal: AssignmentSnapshot,
4002 operation_timeout: Duration,
4003 reason: String,
4004) -> futures::future::BoxFuture<'static, Result<Option<u64>, String>> {
4005 Box::pin(async move {
4006 let decision = authorize_recovery_successor(
4007 &db,
4008 &store,
4009 &controller,
4010 ¤t,
4011 proposal,
4012 operation_timeout,
4013 &reason,
4014 )
4015 .await?;
4016 materialize_recovery_decision(
4017 &db,
4018 &store,
4019 &controller,
4020 ®istry,
4021 decision,
4022 operation_timeout,
4023 )
4024 .await
4025 })
4026}
4027
4028async fn predecessor_cut_unavailability(
4029 controller: &ClusterController,
4030 current: &AssignmentSnapshot,
4031 current_owners: &[NodeId],
4032 live: &[NodeId],
4033) -> Option<String> {
4034 let participant_ids: Vec<u64> = current
4035 .participants
4036 .iter()
4037 .map(|participant| participant.node_id)
4038 .collect();
4039 match controller
4040 .recovery_participant_incarnations(&participant_ids)
4041 .await
4042 {
4043 Ok(observed) if observed == current.participants => {}
4044 Ok(_) => {
4045 return Some("the predecessor process incarnation roster changed".into());
4046 }
4047 Err(error) => {
4048 return Some(format!(
4049 "the predecessor process incarnation roster is unavailable: {error}"
4050 ));
4051 }
4052 }
4053
4054 let members = controller.members_watch().borrow().clone();
4055 let owners: std::collections::BTreeSet<NodeId> = current_owners.iter().copied().collect();
4056 for owner in owners {
4057 if controller.is_unresponsive(owner) {
4058 return Some(format!(
4059 "predecessor owner {owner} cannot certify the source cut"
4060 ));
4061 }
4062 if live.contains(&owner) {
4063 continue;
4064 }
4065 let can_drain = members.iter().any(|member| {
4066 member.id == owner && matches!(member.state, NodeState::Active | NodeState::Draining)
4067 });
4068 if !can_drain {
4069 return Some(format!(
4070 "predecessor owner {owner} cannot certify the source cut"
4071 ));
4072 }
4073 }
4074 None
4075}
4076
4077fn successor_participant_ids(owners: &[NodeId]) -> Vec<u64> {
4078 owners
4079 .iter()
4080 .map(|owner| owner.0)
4081 .collect::<std::collections::BTreeSet<_>>()
4082 .into_iter()
4083 .collect()
4084}
4085
4086fn successor_participants(
4087 owners: &[NodeId],
4088 available: &[CheckpointParticipant],
4089) -> Result<Vec<CheckpointParticipant>, String> {
4090 let participant_ids = successor_participant_ids(owners);
4091 let participant_set: std::collections::BTreeSet<u64> =
4092 participant_ids.iter().copied().collect();
4093 let participants = available
4094 .iter()
4095 .copied()
4096 .filter(|participant| participant_set.contains(&participant.node_id))
4097 .collect::<Vec<_>>();
4098 if participants.len() != participant_ids.len() {
4099 return Err("successor owner roster lost a lease-validated process identity".into());
4100 }
4101 Ok(participants)
4102}
4103
4104fn retain_recovery_predecessors(
4105 assignment: &mut [NodeId],
4106 predecessor: &[CheckpointParticipant],
4107 successors: &[CheckpointParticipant],
4108) -> Result<(), String> {
4109 let successor_processes = successors
4110 .iter()
4111 .map(|participant| (participant.node_id, participant.boot_incarnation))
4112 .collect::<std::collections::BTreeMap<_, _>>();
4113 let retained = predecessor
4114 .iter()
4115 .filter(|participant| {
4116 successor_processes.get(&participant.node_id) == Some(&participant.boot_incarnation)
4117 })
4118 .map(|participant| NodeId(participant.node_id))
4119 .collect::<std::collections::BTreeSet<_>>();
4120 let mut counts = assignment.iter().copied().fold(
4121 std::collections::BTreeMap::<NodeId, usize>::new(),
4122 |mut counts, owner| {
4123 *counts.entry(owner).or_default() += 1;
4124 counts
4125 },
4126 );
4127
4128 for participant in retained.iter().copied() {
4129 if counts.get(&participant).copied().unwrap_or_default() != 0 {
4130 continue;
4131 }
4132 let mut donor = None;
4133 for (index, owner) in assignment.iter().copied().enumerate() {
4134 let owner_count = counts.get(&owner).copied().unwrap_or_default();
4135 if retained.contains(&owner) && owner_count <= 1 {
4136 continue;
4137 }
4138 if donor.is_none_or(|(_, best_count)| owner_count > best_count) {
4139 donor = Some((index, owner_count));
4140 }
4141 }
4142 let Some((index, _)) = donor else {
4143 return Err(format!(
4144 "recovery placement cannot retain healthy predecessor {participant}"
4145 ));
4146 };
4147 let displaced = assignment[index];
4148 assignment[index] = participant;
4149 *counts
4150 .get_mut(&displaced)
4151 .expect("assignment owner has a placement count") -= 1;
4152 *counts.entry(participant).or_default() += 1;
4153 }
4154 Ok(())
4155}
4156
4157#[cfg(test)]
4158pub(crate) async fn admit_cluster_checkpoint_artifacts_for_test(
4159 authority: &LeaderLeaseStore,
4160 proof: &laminar_core::checkpoint::LeaderProof,
4161 index: &laminar_core::checkpoint::CommittedCheckpointIndex,
4162) {
4163 use laminar_core::checkpoint::CheckpointAttempt;
4164 use laminar_core::checkpoint_decision::CheckpointArtifactInventory;
4165
4166 let inventory = CheckpointArtifactInventory {
4167 deployment_id: index.deployment_id.clone(),
4168 pipeline_identity: index.pipeline_identity.clone(),
4169 attempt: CheckpointAttempt::new(index.epoch, index.checkpoint_id),
4170 assignment_fence: index.assignment_fence.clone(),
4171 sink_artifact_intent_protocol: true,
4172 };
4173 assert_eq!(
4174 authority
4175 .begin_cluster_checkpoint_artifacts(proof, inventory.clone())
4176 .await
4177 .unwrap(),
4178 inventory
4179 );
4180}
4181
4182#[cfg(test)]
4183pub(crate) async fn record_assignment_checkpoint_for_test(
4184 authority: &LeaderLeaseStore,
4185 authority_store: &Arc<dyn object_store::ObjectStore>,
4186 fence: &CheckpointAssignmentFence,
4187 proof: &laminar_core::checkpoint::LeaderProof,
4188) -> CommittedCheckpointRef {
4189 use laminar_core::checkpoint::{
4190 CheckpointScope, CommittedCheckpointIndex, CommittedParticipantRef, PipelineIdentity,
4191 COMMITTED_CHECKPOINT_INDEX_VERSION,
4192 };
4193 use laminar_core::checkpoint_decision::{CheckpointDecisionStore, CheckpointVerdict};
4194
4195 let epoch = authority
4196 .highest_cluster_terminal_outcome()
4197 .await
4198 .unwrap()
4199 .map_or(1, |outcome| outcome.epoch.checked_add(1).unwrap());
4200 let predecessor = authority
4201 .highest_cluster_committed_outcome()
4202 .await
4203 .unwrap()
4204 .and_then(|outcome| outcome.committed_checkpoint);
4205 let deployment_id = CheckpointDecisionStore::new(Arc::clone(authority_store))
4206 .load_or_create_deployment_id()
4207 .await
4208 .unwrap();
4209 let index = CommittedCheckpointIndex {
4210 version: COMMITTED_CHECKPOINT_INDEX_VERSION,
4211 deployment_id,
4212 pipeline_identity: PipelineIdentity::empty(),
4213 epoch,
4214 checkpoint_id: epoch,
4215 scope: CheckpointScope::Cluster,
4216 vnode_count: u16::try_from(fence.vnode_count).unwrap(),
4217 assignment_fence: Some(fence.clone()),
4218 reassignment_portable: true,
4219 predecessor,
4220 participants: fence
4221 .participants
4222 .iter()
4223 .map(|participant| CommittedParticipantRef {
4224 participant_id: participant.node_id,
4225 manifest_len: 1,
4226 manifest_sha256: "0".repeat(64),
4227 node_data_len: 0,
4228 node_data_sha256: "1".repeat(64),
4229 })
4230 .collect(),
4231 source_names: Vec::new(),
4232 source_offsets: Default::default(),
4233 channel_progress: Vec::new(),
4234 source_watermarks: Default::default(),
4235 checkpoint_watermark: None,
4236 };
4237 admit_cluster_checkpoint_artifacts_for_test(authority, proof, &index).await;
4238 let reference = authority.create_committed_checkpoint(&index).await.unwrap();
4239 authority
4240 .record_cluster_outcome(
4241 proof,
4242 epoch,
4243 epoch,
4244 fence.clone(),
4245 CheckpointVerdict::Commit,
4246 Some(reference.clone()),
4247 )
4248 .await
4249 .unwrap();
4250 reference
4251}
4252
4253async fn pre_rotation_checkpoint(
4255 db: &Arc<LaminarDB>,
4256 controller: &ClusterController,
4257 transition: &AssignmentDrainTransition,
4258 config: RebalanceConfig,
4259) -> Result<CommittedCheckpointRef, String> {
4260 let deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
4261 let ckpt = db
4262 .checkpoint_until(deadline)
4263 .await
4264 .map_err(|e| e.to_string())?;
4265 if !ckpt.success {
4266 return Err(ckpt
4267 .error
4268 .unwrap_or_else(|| "pre-rotation checkpoint did not commit".into()));
4269 }
4270 if let Some(error) = ckpt.continuation_error() {
4271 return Err(format!(
4272 "pre-rotation checkpoint committed but cannot continue safely: {error}"
4273 ));
4274 }
4275
4276 let authority = controller
4277 .checkpoint_authority()
4278 .map_err(|error| error.to_string())?;
4279 let verification_deadline = tokio::time::Instant::now()
4284 .checked_add(config.checkpoint_timeout)
4285 .ok_or_else(|| "pre-rotation checkpoint authority-read deadline overflowed".to_string())?;
4286 let outcome =
4287 tokio::time::timeout_at(verification_deadline, authority.cluster_outcome(ckpt.epoch))
4288 .await
4289 .map_err(|_| "pre-rotation checkpoint authority read timed out".to_string())?
4290 .map_err(|error| error.to_string())?
4291 .ok_or_else(|| {
4292 "pre-rotation checkpoint disappeared from cluster authority".to_string()
4293 })?;
4294 if !outcome.is_commit()
4295 || outcome.epoch != ckpt.epoch
4296 || outcome.checkpoint_id != ckpt.checkpoint_id
4297 || outcome.assignment_fence.as_ref() != Some(&transition.predecessor)
4298 {
4299 return Err("pre-rotation checkpoint authority does not bind the predecessor cut".into());
4300 }
4301 outcome
4302 .committed_checkpoint
4303 .ok_or_else(|| "pre-rotation Commit has no checkpoint reference".to_string())
4304}
4305
4306async fn await_drain_quorum(
4310 controller: &Arc<ClusterController>,
4311 transition: &AssignmentDrainTransition,
4312 deadline: tokio::time::Instant,
4313) -> bool {
4314 tokio::time::timeout_at(deadline, async {
4315 loop {
4316 match controller.drain_ack_quorum_reached(transition).await {
4317 Ok(true) => return,
4318 Ok(false) => {}
4319 Err(error) => {
4320 debug!(%error, version = transition.target.assignment_version, "durable drain quorum observation failed; retrying");
4321 }
4322 }
4323 tokio::time::sleep(Duration::from_millis(100)).await;
4324 }
4325 })
4326 .await
4327 .is_ok()
4328}
4329
4330fn assignment_binds_local_process(
4331 snapshot: &AssignmentSnapshot,
4332 controller: &ClusterController,
4333) -> Result<bool, String> {
4334 let fence = snapshot
4335 .assignment_fence()
4336 .map_err(|error| error.to_string())?;
4337 let owners = snapshot
4338 .to_vnode_vec(fence.vnode_count)
4339 .map_err(|error| error.to_string())?;
4340 let local = controller.instance_id().0;
4341 let owns = owners.iter().any(|owner| owner.0 == local);
4342 let bound = fence.participant_incarnation(local);
4343 Ok(if owns {
4344 bound == Some(controller.recovery_incarnation())
4345 } else {
4346 bound.is_none()
4347 })
4348}
4349
4350#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4354enum DrainFinalizationMode {
4355 Live,
4356 StoppedRecovery,
4361}
4362
4363async fn finalize_drain_snapshot(
4364 db: &Arc<LaminarDB>,
4365 store: &Arc<AssignmentSnapshotStore>,
4366 controller: &ClusterController,
4367 draining: &AssignmentSnapshot,
4368 predecessor: &AssignmentSnapshot,
4369 requested_verdict: AssignmentDrainVerdict,
4370 handoff_checkpoint: Option<CommittedCheckpointRef>,
4371 config: RebalanceConfig,
4372 mode: DrainFinalizationMode,
4373) -> Result<Option<u64>, String> {
4374 if mode == DrainFinalizationMode::StoppedRecovery
4375 && (requested_verdict != AssignmentDrainVerdict::Abort || handoff_checkpoint.is_some())
4376 {
4377 return Err("stopped recovery may only finalize a source drain as Abort".into());
4378 }
4379 let finalization_deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
4380 let finalization_guard = if mode == DrainFinalizationMode::Live {
4381 let guard =
4382 tokio::time::timeout_at(finalization_deadline, db.assignment_adoption_lock.lock())
4383 .await
4384 .map_err(|_| {
4385 "source-drain finalization timed out waiting for assignment serialization"
4386 .to_string()
4387 })?;
4388 if controller.is_recovering() {
4389 return Err(
4390 "live source-drain finalization yielded to coordinated recovery authority".into(),
4391 );
4392 }
4393 Some(guard)
4394 } else {
4395 None
4396 };
4397 let transition = draining
4398 .drain_transition
4399 .as_ref()
4400 .ok_or_else(|| "draining assignment has no exact transition".to_string())?;
4401 if predecessor
4402 .assignment_fence()
4403 .map_err(|error| error.to_string())?
4404 != transition.predecessor
4405 {
4406 return Err("drain finalization predecessor does not match the transition".into());
4407 }
4408 let deciding_proof = controller
4409 .capture_leader_proof()
4410 .ok_or_else(|| "drain finalization requires a current leader proof".to_string())?;
4411 let requested = match (requested_verdict, handoff_checkpoint) {
4412 (AssignmentDrainVerdict::Commit, Some(reference)) => {
4413 AssignmentDrainDecision::commit(transition, deciding_proof.clone(), reference)
4414 }
4415 (AssignmentDrainVerdict::Abort, None) => {
4416 AssignmentDrainDecision::abort(transition, deciding_proof.clone())
4417 }
4418 _ => Err("drain finalization has an invalid handoff checkpoint".into()),
4419 }?;
4420 let authority = controller
4421 .checkpoint_authority()
4422 .map_err(|error| error.to_string())?;
4423 let decision = match authority
4424 .record_assignment_drain_decision(&deciding_proof, requested)
4425 .await
4426 .map_err(|error| error.to_string())?
4427 {
4428 RecordAssignmentDrainDecisionResult::Created(decision)
4429 | RecordAssignmentDrainDecisionResult::Unchanged(decision) => decision,
4430 RecordAssignmentDrainDecisionResult::Conflict { winner } => {
4431 warn!(
4432 requested = ?requested_verdict,
4433 winner = ?winner.verdict,
4434 version = draining.version,
4435 "another authority-sequenced drain decision already won"
4436 );
4437 winner
4438 }
4439 };
4440 if decision.transition != *transition {
4441 return Err(format!(
4442 "authority decision for assignment {} binds a different drain transition",
4443 draining.version
4444 ));
4445 }
4446 if mode == DrainFinalizationMode::StoppedRecovery
4447 && decision.verdict != AssignmentDrainVerdict::Abort
4448 {
4449 return Err(format!(
4450 "stopped recovery cannot materialize the {:?} authority winner for assignment {}; the round must yield",
4451 decision.verdict, draining.version
4452 ));
4453 }
4454 let proposal = match decision.verdict {
4455 AssignmentDrainVerdict::Commit => draining
4456 .committed_target()
4457 .map_err(|error| error.to_string())?,
4458 AssignmentDrainVerdict::Abort => draining
4459 .aborted_target(predecessor)
4460 .map_err(|error| error.to_string())?,
4461 };
4462
4463 let durable = match store
4464 .finalize_drain(draining, &proposal)
4465 .await
4466 .map_err(|error| error.to_string())?
4467 {
4468 RotateOutcome::Rotated => proposal,
4469 RotateOutcome::Conflict(winner) => *winner,
4470 };
4471 let source_outcome = finalized_drain_outcome(transition, &durable)?;
4472 let authority_outcome = match decision.verdict {
4473 AssignmentDrainVerdict::Commit => laminar_connectors::connector::SourceDrainOutcome::Commit,
4474 AssignmentDrainVerdict::Abort => laminar_connectors::connector::SourceDrainOutcome::Abort,
4475 };
4476 if source_outcome != authority_outcome {
4477 return Err(format!(
4478 "materialized assignment {} conflicts with the authority drain decision",
4479 durable.version
4480 ));
4481 }
4482 drop(finalization_guard);
4486
4487 let version = durable.version;
4488 if mode == DrainFinalizationMode::Live && assignment_binds_local_process(&durable, controller)?
4489 {
4490 let adoption_deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
4491 db.adopt_assignment_snapshot(durable, adoption_deadline)
4492 .await
4493 .map_err(|error| error.to_string())?;
4494 if local_drain_participant(controller, transition).is_some() {
4495 db.resolve_local_source_drain(
4496 transition.id(),
4497 source_outcome,
4498 tokio::time::Instant::now() + config.drain_ack_timeout,
4499 )
4500 .await
4501 .map_err(|error| error.to_string())?;
4502 }
4503 } else {
4504 warn!(
4505 version,
4506 "materialized drain terminal names a predecessor process; local adoption waits for a recovery successor"
4507 );
4508 }
4509 clear_settled_source_drain(controller, transition)?;
4510 let oldest_retained = version.saturating_sub(1);
4511 match store.prune_before(oldest_retained).await {
4512 Ok(()) => {
4513 if let Err(error) = authority
4514 .prune_assignment_drain_decisions_before(&deciding_proof, oldest_retained)
4515 .await
4516 {
4517 warn!(%error, "assignment drain authority prune failed after snapshot prune");
4518 }
4519 }
4520 Err(error) => {
4521 warn!(%error, "snapshot prune failed after drain finalization");
4522 }
4523 }
4524 Ok(Some(version))
4525}
4526
4527#[cfg(test)]
4530pub(crate) async fn settle_source_drain_before_recovery(
4531 db: &Arc<LaminarDB>,
4532 controller: &ClusterController,
4533 config: RebalanceConfig,
4534) -> Result<Option<u64>, String> {
4535 let published_transition = controller.checkpoint_drain_transition();
4536 let Some(store) = db.assignment_snapshot_store.lock().clone() else {
4537 return if published_transition.is_some() {
4538 Err("recovery source-drain settlement has no assignment store".into())
4539 } else {
4540 Ok(None)
4541 };
4542 };
4543 let Some(draining) = store.load().await.map_err(|error| error.to_string())? else {
4544 return Ok(None);
4545 };
4546 if !draining.draining {
4547 return if published_transition.is_some() {
4548 Err("local source-drain authority outlived its durable draining assignment".into())
4549 } else {
4550 Ok(None)
4551 };
4552 }
4553 let transition = draining
4554 .drain_transition
4555 .as_ref()
4556 .ok_or_else(|| "draining assignment has no exact source transition".to_string())?;
4557 if published_transition.as_ref() != Some(transition) {
4558 return Err("durable and locally published source-drain transitions differ".into());
4559 }
4560 let predecessor = store
4561 .load_version(transition.predecessor.assignment_version)
4562 .await
4563 .map_err(|error| error.to_string())?
4564 .ok_or_else(|| {
4565 format!(
4566 "source-drain predecessor assignment {} is unavailable during recovery",
4567 transition.predecessor.assignment_version
4568 )
4569 })?;
4570 finalize_drain_snapshot(
4571 db,
4572 &store,
4573 controller,
4574 &draining,
4575 &predecessor,
4576 AssignmentDrainVerdict::Abort,
4577 None,
4578 config,
4579 DrainFinalizationMode::Live,
4580 )
4581 .await
4582}
4583
4584pub(crate) async fn settle_stopped_source_drain_after_recovery_quorum(
4590 db: &Arc<LaminarDB>,
4591 controller: &ClusterController,
4592 expected_predecessor: &CheckpointAssignmentFence,
4593 config: RebalanceConfig,
4594) -> Result<Option<u64>, String> {
4595 if DbState::load(&db.state) != DbState::Created
4596 || !db.runtime_shutdown.read().is_cancelled()
4597 || !db.cluster_intake_fenced()
4598 || !db.coordinated_recovery_in_progress()
4599 || !controller.is_recovering()
4600 || !controller.process_lease_is_live()
4601 {
4602 return Err(
4603 "stopped source-drain settlement requires the exact fenced Created recovery generation"
4604 .into(),
4605 );
4606 }
4607 if !expected_predecessor.is_canonical() {
4608 return Err("stopped source-drain predecessor certificate is not canonical".into());
4609 }
4610 let store = db
4611 .assignment_snapshot_store
4612 .lock()
4613 .clone()
4614 .ok_or_else(|| "stopped source-drain settlement has no assignment store".to_string())?;
4615 let deadline = tokio::time::Instant::now() + config.checkpoint_timeout;
4616 let Some(draining) = tokio::time::timeout_at(deadline, store.load())
4617 .await
4618 .map_err(|_| "stopped source-drain head read timed out".to_string())?
4619 .map_err(|error| error.to_string())?
4620 else {
4621 return Err("stopped source-drain settlement has no durable assignment head".into());
4622 };
4623 let target_version = expected_predecessor
4624 .assignment_version
4625 .checked_add(1)
4626 .ok_or_else(|| "stopped source-drain assignment version overflow".to_string())?;
4627 if !draining.draining {
4628 if draining.version == expected_predecessor.assignment_version {
4629 if draining
4630 .assignment_fence()
4631 .map_err(|error| error.to_string())?
4632 != *expected_predecessor
4633 {
4634 return Err(
4635 "stable source-drain head does not match the stopped recovery predecessor"
4636 .into(),
4637 );
4638 }
4639 tokio::time::timeout_at(
4640 deadline,
4641 audit_assignment_snapshot_authority_outcome(&store, Some(controller), &draining),
4642 )
4643 .await
4644 .map_err(|_| {
4645 "stopped source-drain predecessor authority audit timed out".to_string()
4646 })??;
4647 let confirmed = tokio::time::timeout_at(deadline, store.load())
4648 .await
4649 .map_err(|_| "stopped source-drain predecessor recheck timed out".to_string())?
4650 .map_err(|error| error.to_string())?;
4651 if confirmed.as_ref() != Some(&draining) {
4652 return Err("stopped source-drain predecessor changed during audit".into());
4653 }
4654 return Ok(None);
4655 }
4656 if draining.version != target_version {
4657 return Err(format!(
4658 "stable assignment {} superseded stopped source-drain predecessor {}",
4659 draining.version, expected_predecessor.assignment_version
4660 ));
4661 }
4662 let audited = tokio::time::timeout_at(
4663 deadline,
4664 audit_assignment_snapshot_authority_outcome(&store, Some(controller), &draining),
4665 )
4666 .await
4667 .map_err(|_| "stopped source-drain terminal authority audit timed out".to_string())??;
4668 let terminal = audited.into_terminal().ok_or_else(|| {
4669 format!(
4670 "stable assignment {} is not an authority-sequenced drain terminal",
4671 draining.version
4672 )
4673 })?;
4674 if terminal.transition.predecessor != *expected_predecessor
4675 || terminal.transition.target.assignment_version != target_version
4676 {
4677 return Err(format!(
4678 "stable assignment {} is not an exact terminal for stopped predecessor {}",
4679 draining.version, expected_predecessor.assignment_version
4680 ));
4681 }
4682 let target_fence = draining
4683 .assignment_fence()
4684 .map_err(|error| error.to_string())?;
4685 let predecessor = tokio::time::timeout_at(
4686 deadline,
4687 store.load_version(expected_predecessor.assignment_version),
4688 )
4689 .await
4690 .map_err(|_| "stopped source-drain predecessor read timed out".to_string())?
4691 .map_err(|error| error.to_string())?
4692 .ok_or_else(|| {
4693 format!(
4694 "source-drain predecessor assignment {} is unavailable after stop quorum",
4695 expected_predecessor.assignment_version
4696 )
4697 })?;
4698 if predecessor.draining
4699 || predecessor
4700 .assignment_fence()
4701 .map_err(|error| error.to_string())?
4702 != *expected_predecessor
4703 {
4704 return Err("stable source-drain terminal lost its frozen predecessor".into());
4705 }
4706 let commit_pin = if terminal.is_abort() {
4707 if draining.vnodes != predecessor.vnodes
4708 || draining.participants != predecessor.participants
4709 || target_fence.assignment_version != target_version
4710 || target_fence.partitioning_abi_version
4711 != expected_predecessor.partitioning_abi_version
4712 || target_fence.vnode_count != expected_predecessor.vnode_count
4713 || target_fence.assignment_digest != expected_predecessor.assignment_digest
4714 || target_fence.participants != expected_predecessor.participants
4715 {
4716 return Err(
4717 "stable source-drain Abort does not retain the exact stopped predecessor topology"
4718 .into(),
4719 );
4720 }
4721 None
4722 } else if terminal.is_commit() {
4723 if target_fence != terminal.transition.target {
4724 return Err(
4725 "stable source-drain Commit does not materialize its exact target fence".into(),
4726 );
4727 }
4728 let handoff = terminal
4729 .committed_handoff_checkpoint()
4730 .cloned()
4731 .ok_or_else(|| {
4732 "stable source-drain Commit has no exact handoff checkpoint".to_string()
4733 })?;
4734 let authority = controller
4735 .checkpoint_authority()
4736 .map_err(|error| error.to_string())?;
4737 let pinned = tokio::time::timeout_at(
4738 deadline,
4739 authority.assignment_handoff_checkpoint(&target_fence),
4740 )
4741 .await
4742 .map_err(|_| "stopped source-drain Commit handoff-pin read timed out".to_string())?
4743 .map_err(|error| error.to_string())?;
4744 if pinned.as_ref() != Some(&handoff) {
4745 return Err(
4746 "stable source-drain Commit no longer retains its exact active handoff checkpoint pin"
4747 .into(),
4748 );
4749 }
4750 Some((authority, handoff))
4751 } else {
4752 return Err("stable source-drain terminal has no supported verdict".into());
4753 };
4754 let confirmed = tokio::time::timeout_at(deadline, store.load())
4755 .await
4756 .map_err(|_| "stopped source-drain terminal recheck timed out".to_string())?
4757 .map_err(|error| error.to_string())?;
4758 if confirmed.as_ref() != Some(&draining) {
4759 return Err("stopped source-drain terminal changed during reconciliation".into());
4760 }
4761 let confirmed_authority = tokio::time::timeout_at(
4762 deadline,
4763 audit_assignment_snapshot_authority_outcome(&store, Some(controller), &draining),
4764 )
4765 .await
4766 .map_err(|_| "stopped source-drain terminal authority recheck timed out".to_string())??;
4767 if confirmed_authority.into_terminal().as_ref() != Some(&terminal) {
4768 return Err(
4769 "stopped source-drain terminal authority changed during reconciliation".into(),
4770 );
4771 }
4772 if let Some((authority, handoff)) = commit_pin {
4773 let confirmed_pin = tokio::time::timeout_at(
4774 deadline,
4775 authority.assignment_handoff_checkpoint(&target_fence),
4776 )
4777 .await
4778 .map_err(|_| "stopped source-drain Commit handoff-pin recheck timed out".to_string())?
4779 .map_err(|error| error.to_string())?;
4780 if confirmed_pin.as_ref() != Some(&handoff) {
4781 return Err(
4782 "stable source-drain Commit handoff pin changed during reconciliation".into(),
4783 );
4784 }
4785 }
4786 clear_settled_source_drain(controller, &terminal.transition)?;
4787 return Ok(Some(target_version));
4788 }
4789 let transition = draining
4790 .drain_transition
4791 .as_ref()
4792 .ok_or_else(|| "draining assignment has no exact source transition".to_string())?;
4793 if transition.predecessor != *expected_predecessor
4794 || transition.target.assignment_version != target_version
4795 || draining.version != transition.target.assignment_version
4796 {
4797 return Err(format!(
4798 "draining assignment {} does not bind the stopped recovery predecessor {}",
4799 draining.version, expected_predecessor.assignment_version
4800 ));
4801 }
4802 if controller
4803 .checkpoint_drain_transition()
4804 .as_ref()
4805 .is_some_and(|published| published != transition)
4806 {
4807 return Err("durable and locally published source-drain transitions differ".into());
4808 }
4809 let predecessor = tokio::time::timeout_at(
4810 deadline,
4811 store.load_version(expected_predecessor.assignment_version),
4812 )
4813 .await
4814 .map_err(|_| "stopped source-drain predecessor read timed out".to_string())?
4815 .map_err(|error| error.to_string())?
4816 .ok_or_else(|| {
4817 format!(
4818 "source-drain predecessor assignment {} is unavailable after stop quorum",
4819 expected_predecessor.assignment_version
4820 )
4821 })?;
4822 if predecessor.draining
4823 || predecessor
4824 .assignment_fence()
4825 .map_err(|error| error.to_string())?
4826 != *expected_predecessor
4827 {
4828 return Err("stopped source-drain predecessor no longer matches the frozen round".into());
4829 }
4830 let confirmed = tokio::time::timeout_at(deadline, store.load())
4831 .await
4832 .map_err(|_| "stopped source-drain head recheck timed out".to_string())?
4833 .map_err(|error| error.to_string())?;
4834 if confirmed.as_ref() != Some(&draining) {
4835 return Err("stopped source-drain head changed during settlement".into());
4836 }
4837 finalize_drain_snapshot(
4838 db,
4839 &store,
4840 controller,
4841 &draining,
4842 &predecessor,
4843 AssignmentDrainVerdict::Abort,
4844 None,
4845 config,
4846 DrainFinalizationMode::StoppedRecovery,
4847 )
4848 .await
4849}
4850
4851async fn adopt_any(
4853 db: &Arc<LaminarDB>,
4854 store: &AssignmentSnapshotStore,
4855 controller: &ClusterController,
4856 registry: &VnodeRegistry,
4857 snap: AssignmentSnapshot,
4858 operation_timeout: Duration,
4859) -> Result<(), String> {
4860 let deadline = tokio::time::Instant::now() + operation_timeout;
4861 let audited = tokio::time::timeout_at(
4862 deadline,
4863 audit_assignment_snapshot_authority_outcome(store, Some(controller), &snap),
4864 )
4865 .await
4866 .map_err(|_| format!("assignment {} authority audit timed out", snap.version))??;
4867 if snap.draining {
4868 db.validate_source_drain_snapshot(&snap)
4869 .map_err(|error| error.to_string())?;
4870 } else if audited.is_recovery() {
4871 prepare_recovery_assignment_adoption(db, store, controller, registry, &snap, deadline)
4872 .await?;
4873 db.adopt_recovery_assignment_snapshot(snap, operation_timeout)
4874 .await
4875 .map_err(|error| error.to_string())?;
4876 } else {
4877 db.adopt_assignment_snapshot(snap, deadline)
4878 .await
4879 .map_err(|e| e.to_string())?;
4880 }
4881 Ok(())
4882}
4883
4884#[cfg(test)]
4885mod tests;