Skip to main content

laminar_db/rebalance/
mod.rs

1//! Dynamic vnode rebalance control plane.
2
3#![cfg(feature = "cluster")]
4#![allow(clippy::disallowed_types)] // cold path
5
6use 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/// Tunables for the rebalance control plane.
46#[derive(Debug, Clone, Copy)]
47pub struct RebalanceConfig {
48    /// Interval between snapshot-store polls.
49    pub watcher_poll: Duration,
50    /// Quiet period before a membership change triggers rotation.
51    pub rebalance_debounce: Duration,
52    /// Upper bound used for the pre-rotation checkpoint and each serialized recovery-adoption
53    /// phase.
54    pub checkpoint_timeout: Duration,
55    /// Delay before retrying a failed rotation.
56    pub retry_delay: Duration,
57    /// Bound on waiting for the frozen predecessor roster to ack the draining snapshot before
58    /// the pre-rotation checkpoint; on timeout the rotation aborts. Must exceed `watcher_poll`.
59    pub drain_ack_timeout: Duration,
60    /// Locality tier the placement metrics group by (0 = coarsest).
61    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            // A healthy pre-rotation drain commits in well under a second; the longer budget
70            // absorbs slow external source cuts without weakening the frozen-roster quorum.
71            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    /// Fast timings for tests — 500ms debounce thrashes in production.
81    #[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    // Cancel first: a compute-cycle read guard may itself be blocked in shuffle admission.
106    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    // A watcher that already owned the adoption lock could have republished and reopened after
111    // the first cancellation. Reassert the full closure and retain serialization through process
112    // fencing, fault publication, and the execution drain.
113    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    // RECOVERY: Prepare cancels graph work, so publish its fault before waiting for that work to
129    // release the execution fence. Process fencing above prevents a live predecessor from causing
130    // a spurious fault, and the serialized local authority is already closed.
131    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
145/// Fail closed for a transient durable snapshot read without forcing a new assignment version.
146/// The exact retained certificate can resume after the same durable head is audited again.
147async 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    // Preserve the certificate and its sequence domain, but cancel its active lifetime before
158    // waiting for compute cycles that may be blocked on that lifetime.
159    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
166/// Suspend authority after the caller has serialized assignment adoption.
167async fn suspend_local_assignment_authority_locked(
168    db: &Arc<LaminarDB>,
169    controller: Option<&ClusterController>,
170    deadline: tokio::time::Instant,
171) -> Result<(), String> {
172    // A watcher that owned the adoption lock before the first cancellation could have republished
173    // and resumed. Reassert the full suspension while serialized, before draining execution.
174    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
189/// Recovery heads must not suspend predecessor assignment authority while the graph still owns a
190/// pending vnode transition. Holding the adoption lock makes the preflight stable against both
191/// assignment adoption and startup staging; the graph may only finish (not create) that work.
192async 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
237/// Spawn the per-node snapshot watcher. Exits on `shutdown`.
238///
239/// # Panics
240///
241/// The spawned watcher panics if an already-validated draining snapshot lacks its transition.
242struct 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    /// Install an observed durable drain at the same local boundary used by checkpoint
293    /// assignment admission. A checkpoint that already owns the mutex may finish publishing its
294    /// predecessor Prepare first; otherwise the re-audit below makes the drain visible before any
295    /// later admission can read assignment flags.
296    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        // The durable-head audit and authority publication are independently bounded phases.
392        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            // The durable head used above predates an authority closure by another adoption.
398            // Keep that closure in force and re-read the head on the next tick.
399            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        // Serialize the report with assignment adoption and boot staging. A transition publishes
414        // its false report under this same lock before exposing pending state, so a
415        // delayed watcher can never overwrite that withdrawal with stale readiness.
416        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                    // A stopped owner removed from the target, or an evidence-only reporter, has
474                    // no target vnode state to restore. Its false report is still the durable
475                    // topology-observed acknowledgement that permits a gap-free Prepare handoff.
476                    false
477                };
478                // Do not suppress same-value writes here. Startup, assignment adoption, and
479                // recovery staging also publish this durable slot directly; a watcher-local cache
480                // cannot observe those writes and could otherwise strand a newer false report.
481                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        // Publish a version-bound, owner-complete fence off the hot path on every node.
517        // The adoption lock and authority revision reject a proof computed from a head that
518        // another task fenced while this watcher was awaiting durable I/O.
519        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                            // Cancel the cached lifetime before waiting for compute cycles that
585                            // may be blocked in send. Suspension is safe if another authority
586                            // path already installed this exact successor; a genuinely higher
587                            // certificate resets the delivery domain during installation.
588                            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                    // An incomplete adoption or process-roster read is not proof that the
626                    // installed certificate was superseded. Close admission while retaining
627                    // its delivery sequence; installing a later concrete successor resets the
628                    // delivery domain from that certified boundary.
629                    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        // Recovery may reconcile a drain terminal only after its exact predecessor owner roster is
674        // stopped. Those source tasks are intentionally gone, so neither Commit nor Abort can be
675        // resolved on this generation. Let topology-only adoption publish the terminal; the
676        // replacement source generation consumes it in pre-Release reconciliation.
677        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            // The durable namespace is authoritative, so every tick audits it even when the local
790            // version has not changed. This closes the crash window after a successful CAS but
791            // before local adoption.
792            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                // Drain phase: hold every predecessor source at one global input frontier;
855                // ownership is unchanged so the registry version stays put.
856                Ok(Ok(Some(snap))) if snap.draining && snap.version > local => {
857                    // A draining head describes the successor map; checkpoint and shuffle still
858                    // belong to the exact committed predecessor until the pre-rotation cut lands.
859                    // Re-audit that predecessor on every tick, including after this process starts
860                    // while a drain is already in progress. Retaining a previously cached fence
861                    // without this read would trust an object that may no longer be the durable
862                    // predecessor of the observed head.
863                    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                                                // A target-only joining process has no predecessor
918                                                // input authority and is intentionally absent from
919                                                // the receipt quorum.
920                                                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                                                // Receipt production may consume the entire head
944                                                // budget. Re-read and recertify from a fresh
945                                                // durable head on the next tick.
946                                                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                                    // The transition may have been staged before its predecessor
1042                                    // transport certificate became active. Repair that exact audited
1043                                    // authority before waiting for the newer durable head.
1044                                    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                            // Suspension succeeded with no pending live transition. Do not let the
1058                            // watcher recertify the stale local generation while the requested fault
1059                            // is still moving the graph to the cold-recovery boundary.
1060                            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                                // The local graph still needs the installed predecessor authority
1140                                // to consume its staged vnode transition. Keep the prior audited
1141                                // snapshot/certificate cached and retry this successor next tick.
1142                                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                                // Fall through to predecessor authority publication below. A
1160                                // staged transition cannot complete if its retained certificate
1161                                // was suspended or its first activation previously timed out.
1162                            }
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
1214/// Spawn the per-node snapshot watcher. Exits on shutdown.
1215///
1216/// # Panics
1217///
1218/// The spawned watcher panics if an already-validated draining snapshot lacks its transition.
1219pub 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}
1229/// Per-node checkpoint fence. A reported version is insufficient: the exact current
1230/// assignment must be owner-complete over the same canonical participant set.
1231async 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                // Transport activation requires exact assignment adoption, including while vnode
1284                // state is still installing. Graceful rotation applies the stronger semantic
1285                // readiness gate separately immediately before its durable CAS.
1286                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/// Durable authority proof for one materialized assignment target.
1384///
1385/// This capability binds only the target's immediate predecessor. Older authority history may
1386/// already have been pruned and must not be reopened while installing or recertifying the target.
1387#[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    /// The unresolved cut is still pinned to this exact recovery target.
1400    Active {
1401        checkpoint: CommittedCheckpointRef,
1402        origin: CheckpointAssignmentFence,
1403    },
1404    /// The unresolved cut has been carried forward to the exact next recovery target, whose
1405    /// immutable authority decision is durable but whose proposal is not materialized yet.
1406    Propagated {
1407        checkpoint: CommittedCheckpointRef,
1408        origin: CheckpointAssignmentFence,
1409        successor: CheckpointAssignmentFence,
1410    },
1411    /// This target has committed its own successor cut and consumed the older handoff pin.
1412    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
1574/// A committed drain transition returned only after durable authority audit.
1575///
1576/// The private field makes this a capability: sibling modules may consume it, but cannot mint one
1577/// from an arbitrary canonical transition.
1578pub(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    /// The exact transition whose target durably committed, if this was not an abort.
1622    #[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
1629/// Audit a materialized assignment-drain outcome against the shared authority sequence.
1630/// Ordinary assignments and still-draining transitions require no terminal decision.
1631///
1632/// # Errors
1633/// Returns an error for missing/malformed authority history or a materialization that disagrees
1634/// with the immutable terminal decision.
1635pub 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
1645/// Select the last committed assignment that a cluster process may adopt during startup.
1646/// A draining head has not transferred ownership, so startup retains its audited predecessor.
1647///
1648/// # Errors
1649///
1650/// Returns an error when the head or retained predecessor lacks valid durable authority.
1651pub 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            // A Commit under the recovery target atomically consumes its handoff pin. The older
1970            // decision checkpoint may then be garbage-collected; validate and recover from the
1971            // exact target-fenced head instead of reopening that retired artifact.
1972            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            // Recording the exact next recovery decision moves the unresolved pin from this
2007            // materialized target to its successor before that successor proposal is installed.
2008            // Retain authority for this transient head only when all four immutable edges agree:
2009            // D(n), D(n+1), the successor-keyed pin, and the unchanged committed cut.
2010            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
2101/// Resolve the old assignment's in-flight checkpoint once an immutable recovery decision has
2102/// made that attempt abort-only. This publishes only terminal authority: artifact sealing and
2103/// deletion remain behind coordinated recovery's stopped quorum because participant writers may
2104/// still be unwinding when they observe the Abort.
2105fn 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        // Followers only consume the shared outcome. A stale local proof also fails atomically in the
2113        // authority CAS and is retried after the next leader observation.
2114        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
2143/// Settle an already-audited predecessor checkpoint in a separate boxed state machine. Keeping the
2144/// authority audit and settlement state distinct bounds debug-build stack usage on the recovery
2145/// path without weakening any of the post-audit revalidation below.
2146fn 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        // Once a target-generation inventory exists, single-active-inventory admission proves the
2197        // predecessor attempt was already terminal and cleaned. Recovery-origin assignments stay
2198        // identifiable after their handoff pin is legitimately cleared by a target Commit.
2199        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                    &current_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/// Apply a durable drain outcome before adopting any later assignment. The exact terminal
2448/// generation must be installed before connector resolution; skipping it would make the source
2449/// cut unverifiable.
2450#[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
2590/// Reapply a materialized drain terminal to a source generation created by coordinated recovery.
2591/// The caller keeps recovery and intake fenced until this returns.
2592pub(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        &registry,
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    // Receipt production can block on a connector FIFO. Re-read both the transition and its
2815    // leader lease before making that receipt visible to the frozen predecessor quorum.
2816    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
2827/// Publish per-domain owner counts. Resets the gauge so disappeared domains don't leave stale series.
2828fn 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
2857/// Spawn the leader-gated rebalance controller. Runs on every node;
2858/// leadership is re-checked after the debounce.
2859pub 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(())) => {}       // another change; keep waiting
2896                                Ok(Err(_)) => return,  // sender dropped
2897                                Err(_) => break,        // quiet period elapsed
2898                            }
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                // Assignable = Active, non-draining — Draining/Suspected nodes never receive vnodes.
2910                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(&registry),
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
2943/// Poll the durable assignment snapshot until `me` owns no vnodes (its
2944/// state has been reassigned elsewhere) or `deadline` elapses. Returns
2945/// true if fully drained. A version materialized from a drain is accepted only with its matching
2946/// shared-authority decision; `None` therefore supports ordinary snapshots but fails closed for a
2947/// drain finalization. Used by a draining node to know when it is safe to exit.
2948pub 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            // Assignment adoption and startup recovery create pending vnode work under this lock.
3363            // Keep it through the durable CAS so it cannot appear after preflight but
3364            // before V+1 is published. Once the recovery gate below admits publication, the
3365            // conditional write must resolve without an outer cancellation deadline: releasing
3366            // this lock while an admitted remote write can still commit would let recovery select
3367            // the predecessor before V+1 appears. The scope must end before conflict adoption or
3368            // drain finalization, both of which acquire the same lock.
3369            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(&registry, &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            // This mutable report scan is a best-current preflight, not a cross-node lease. It
3422            // prevents assignment chaining before every participant's initial install; a later
3423            // withdrawal is contained by source drain and forced-checkpoint abort/recovery.
3424            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                // The assignment-adoption claim is also checkpoint admission's linearization
3474                // boundary. Publish the local transition before releasing it, so either an
3475                // ordinary checkpoint durably prepares against the predecessor first or every
3476                // later admission observes this exact drain.
3477                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                // Wait for every process in the snapshot's frozen boot roster to durably ack the
3498                // exact source cut. Target-only joiners intentionally do not acknowledge inputs
3499                // they never owned under the predecessor.
3500                let local_receipt = match local_drain_participant(&controller, &transition) {
3501                    Some(participant) => {
3502                        prepare_and_announce_local_drain(
3503                            &db,
3504                            &store,
3505                            &registry,
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                        &registry,
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                    // RECOVERY: a missing receipt means at least one provider FIFO cut is
3539                    // unknown. Live Abort adoption fences predecessor execution before that
3540                    // source can finish draining its queued pre-cut payload. Keep the draining
3541                    // head authoritative and let stopped recovery materialize Abort only after
3542                    // the complete predecessor roster is fenced.
3543                    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                    &registry,
3554                    &drain,
3555                    &controller,
3556                    tokio::time::Instant::now() + config.checkpoint_timeout,
3557                )
3558                .await?;
3559                // Abort the drain on failure OR timeout (not just Ok(false)) — a bare
3560                // `?` here would leave nodes stuck draining.
3561                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                                &registry,
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                                &current,
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                    &registry,
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                    &current,
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                    &registry,
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        // Materialization and predecessor settlement are separate durable operations. A leader
3661        // may fail between them, so every later leader audit of a recovery-origin head retries the
3662        // decision-only Abort even when this process has already installed that assignment.
3663        if current_authority.is_recovery() {
3664            abort_predecessor_checkpoint_for_recovery(&store, &controller, &current, 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            &current,
3674            &current_owners,
3675            local_assignment.version(),
3676            local_assignment.owners(),
3677        )?;
3678
3679        // A propagated pin means D(current + 1) is already durable while `current` remains the
3680        // materialized assignment head. A faulted graph may cold-bootstrap directly from the
3681        // pinned origin, and a process that already installed `current` may advance the decision.
3682        // A healthy lagging process must install `current` first: materializing the successor here
3683        // would skip its only exact live-transition edge and strand it behind the newer head.
3684        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(&registry),
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            // The retained predecessor may name an owner process that died during the drain. Close
3754            // intake before publishing the rollback, but preserve its exact delivery domain: live
3755            // finalization may yield to coordinated recovery, which must be able to recertify this
3756            // same predecessor version before it can stop the owner-complete roster. A successful
3757            // terminal adoption resets transport under the higher assignment version.
3758            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                &current,
3765                &prior,
3766                AssignmentDrainVerdict::Abort,
3767                None,
3768                config,
3769                DrainFinalizationMode::Live,
3770            )
3771            .await;
3772        }
3773
3774        // Settle an existing drain before considering the current placement input. An empty live set
3775        // must not strand a durable transition without its authority-sequenced outcome.
3776        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(&current_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(&current, &controller)?
3812        {
3813            return adopt_materialized_recovery_head(
3814                &db,
3815                &store,
3816                &controller,
3817                &registry,
3818                &current,
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            // A writer can fail after its durable CAS succeeds but before local adoption. Adopt
3827            // that exact live-process generation before planning another rotation. A replacement
3828            // process must not adopt a retained certificate naming its predecessor incarnation;
3829            // it proceeds below through the authority-sequenced recovery-successor path.
3830            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        // Reconcile a staged successor only after this process has installed the durable current
3843        // generation. Otherwise a lagging v1 process could materialize v3 before adopting v2 and
3844        // lose the exact-successor recovery path.
3845        if let Some(version) = reconcile_pending_recovery_decision_owned(
3846            Arc::clone(&db),
3847            Arc::clone(&store),
3848            Arc::clone(&controller),
3849            Arc::clone(&registry),
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        // The successor checkpoint quorum is the exact successor owner set. Deriving it from the
3861        // current membership/checkpoint roster retains a gracefully departing process after all of
3862        // its vnodes move away; the first checkpoint after that process exits can then never reach
3863        // quorum. The predecessor roster remains authoritative for the drain cut and is carried by
3864        // `current`.
3865        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, &current, &current_owners, &successors)
3874                .await;
3875        if let Some(reason) = recovery_reason {
3876            retain_recovery_predecessors(
3877                &mut new_assignment,
3878                &current.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(&registry),
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            &registry,
3988            &current,
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            &current,
4011            proposal,
4012            operation_timeout,
4013            &reason,
4014        )
4015        .await?;
4016        materialize_recovery_decision(
4017            &db,
4018            &store,
4019            &controller,
4020            &registry,
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
4253/// Seal and resolve the exact predecessor-fenced checkpoint used for state handoff.
4254async 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    // The attempt is already terminal when `checkpoint_until` returns. Its private Abort/rollback
4280    // or post-Commit continuation budget may legitimately finish after the admission deadline, so
4281    // use a fresh bounded read window solely to verify the immutable terminal record. This cannot
4282    // admit or mutate checkpoint work.
4283    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
4306/// Wait until every exact process in the draining assignment certificate has durably proved it
4307/// reached the global input cut. Durable read errors are retried within the existing deadline;
4308/// they never become quorum.
4309async 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/// Settle a draining generation without changing the target version certified by source receipts.
4351/// The terminal verdict first enters the shared leader/checkpoint authority sequence; the snapshot
4352/// store then publishes that immutable verdict as the assignment materialization.
4353#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4354enum DrainFinalizationMode {
4355    Live,
4356    /// Coordinated recovery has stopped this exact predecessor generation. Publish an Abort
4357    /// rollback topology, but leave source resolution to the replacement generation's Release
4358    /// reconciliation and topology adoption to the snapshot watcher. The caller holds the
4359    /// assignment-adoption lock across this publication and its subsequent recovery decision.
4360    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    // Durable terminalization is the only portion mutually exclusive with recovery target
4483    // selection. Ordinary live adoption takes this same lock itself and therefore begins only
4484    // after the authority decision and snapshot materialization are indivisible to recovery.
4485    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/// Settle a durable drain before coordinated recovery freezes a new process-local source cut.
4528/// Recovery restarts source tasks, so a pre-restart receipt can never authorize a later Release.
4529#[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
4584/// Settle an exact draining predecessor only after coordinated recovery has stopped its complete
4585/// owner roster and reconciled every checkpoint artifact. A still-draining head is aborted; an
4586/// already materialized exact terminal is reconciled idempotently. Its topology is published on
4587/// this stopped generation, while source-task resolution is deliberately deferred until the fresh
4588/// recovery round has rebuilt the source generation and reaches Release reconciliation.
4589pub(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
4851/// Adopt a snapshot whether it is a draining or a committed one.
4852async 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;