laminar_core/cluster/control/leader_lease/
assignment_drain.rs

1//! Durable intent and materialization of an exact graceful assignment drain.
2
3use super::topology_admission::MAX_ADMISSION_ATTEMPTS;
4use super::*;
5use crate::cluster::control::snapshot::AssignmentSnapshot;
6
7const DRAIN_PROPOSAL_PREFIX: &str = "control/assignment-drain-proposals/v1/";
8
9fn drain_path(reference: &AssignmentSnapshotRef) -> OsPath {
10    OsPath::from(format!(
11        "{DRAIN_PROPOSAL_PREFIX}v{:020}/{}.json",
12        reference.version, reference.sha256
13    ))
14}
15
16impl LeaderLeaseStore {
17    async fn load_drain_proposal(
18        &self,
19        reference: &AssignmentSnapshotRef,
20    ) -> Result<AssignmentSnapshot, LeaseError> {
21        reference
22            .validate()
23            .map_err(|e| LeaseError::Invalid(e.to_string()))?;
24        let bytes = self
25            .read_admission_blob(&drain_path(reference), reference.encoded_len)
26            .await?;
27        let proposal: AssignmentSnapshot =
28            serde_json::from_slice(&bytes).map_err(|e| LeaseError::Invalid(e.to_string()))?;
29        let (canonical, actual) = proposal
30            .encode_drain_proposal()
31            .map_err(|e| LeaseError::Invalid(e.to_string()))?;
32        if actual != *reference || canonical.as_slice() != bytes.as_ref() {
33            return Err(LeaseError::Invalid(
34                "drain proposal does not match its immutable reference".into(),
35            ));
36        }
37        Ok(proposal)
38    }
39
40    /// Reserve and publish an exact graceful assignment drain through shared authority.
41    ///
42    /// A reservation survives cancellation before snapshot publication. Checkpoint/migration
43    /// admission sees that reservation; only an exact drain/recovery decision can release it.
44    /// The supplied store must be the namespace-verified assignment store used by the controller.
45    ///
46    /// # Errors
47    /// Rejects stale authority, another reserved operation, a changed predecessor, or storage I/O.
48    pub async fn publish_assignment_drain(
49        &self,
50        proof: &LeaderProof,
51        assignments: &AssignmentSnapshotStore,
52        proposal: &AssignmentSnapshot,
53    ) -> Result<RotateOutcome, ClusterCheckpointAuthorityError> {
54        let (encoded, reference) = proposal
55            .encode_drain_proposal()
56            .map_err(snapshot_authority_error)?;
57        let transition = proposal
58            .drain_transition
59            .as_ref()
60            .ok_or_else(|| LeaseError::Invalid("missing drain transition".into()))?;
61        if &transition.leader != proof || !proof.is_canonical() {
62            return Err(ClusterCheckpointAuthorityError::Fenced);
63        }
64        for _ in 0..MAX_ADMISSION_ATTEMPTS {
65            let published = self
66                .load_published_authority_head()
67                .await?
68                .ok_or(ClusterCheckpointAuthorityError::Fenced)?;
69            let current = &published.record;
70            if !current.lease.matches_proof(proof) {
71                return Err(ClusterCheckpointAuthorityError::Fenced);
72            }
73            current.reject_topology_preparation()?;
74            self.validate_topology_assignment_proposal_from(current, &transition.target)
75                .await?;
76            if let Some(reservation) = &current.assignment_drain_reservation {
77                if reservation.proposal != reference || &reservation.transition != transition {
78                    return Err(DecisionError::Conflict(
79                        "another assignment drain is reserved".into(),
80                    )
81                    .into());
82                }
83                self.load_drain_proposal(&reference).await?;
84                return assignments
85                    .save_if_version(proposal, transition.predecessor.assignment_version)
86                    .await
87                    .map_err(snapshot_authority_error);
88            }
89            if current
90                .recovery_fault_slots
91                .iter()
92                .any(|slot| slot.active || slot.disposition == RecoveryFaultDisposition::Terminal)
93            {
94                return Err(DecisionError::Conflict(
95                    "assignment drain cannot overtake recovery".into(),
96                )
97                .into());
98            }
99            let predecessor = assignments
100                .load()
101                .await
102                .map_err(snapshot_authority_error)?
103                .ok_or_else(|| LeaseError::Invalid("assignment predecessor is missing".into()))?;
104            if predecessor.version == proposal.version {
105                // Compatibility with a drain published before this binary's coordinated upgrade.
106                // Definitive settlement still uses the shared authority as before.
107                return Ok(RotateOutcome::Conflict(Box::new(predecessor)));
108            }
109            if predecessor.draining
110                || predecessor
111                    .assignment_fence()
112                    .map_err(snapshot_authority_error)?
113                    != transition.predecessor
114            {
115                return Err(
116                    DecisionError::Conflict("assignment drain predecessor changed".into()).into(),
117                );
118            }
119            self.reject_consumed_checkpoint_assignment(current, &transition.predecessor)
120                .await?;
121            self.stage_admission_blob(&drain_path(&reference), &encoded)
122                .await?;
123            let mut lease = current.lease.clone();
124            lease.seq = lease
125                .seq
126                .checked_add(1)
127                .ok_or_else(|| LeaseError::Invalid("authority sequence exhausted".into()))?;
128            let sequence = lease.seq;
129            let mut next = current.preserve_with_lease(lease);
130            next.version = next.version.max(TOPOLOGY_ADMISSION_RECORD_VERSION);
131            next.assignment_drain_reservation = Some(AssignmentDrainReservation {
132                proposal: reference.clone(),
133                transition: transition.clone(),
134                authority_sequence: sequence,
135            });
136            match self
137                .create_authority_record(Some(&published), &next)
138                .await?
139            {
140                AuthorityCreateOutcome::Created | AuthorityCreateOutcome::ExistingIdentical => {
141                    return assignments
142                        .save_if_version(proposal, transition.predecessor.assignment_version)
143                        .await
144                        .map_err(snapshot_authority_error);
145                }
146                AuthorityCreateOutcome::Contended(_) => tokio::task::yield_now().await,
147            }
148        }
149        Err(LeaseError::Io(
150            "assignment drain admission contention exhausted its fixed retry bound".into(),
151        )
152        .into())
153    }
154
155    /// Materialize an already-admitted drain after cancellation or leader replacement.
156    ///
157    /// This publishes the original fenced intent; it never transfers vnode ownership or reopens
158    /// intake. A replacement leader settles it through the existing drain/recovery decision path.
159    ///
160    /// # Errors
161    /// Rejects corrupt/missing proposals or conflicting assignment history.
162    pub async fn materialize_reserved_assignment_drain(
163        &self,
164        assignments: &AssignmentSnapshotStore,
165    ) -> Result<Option<RotateOutcome>, ClusterCheckpointAuthorityError> {
166        let Some(head) = self.load_record().await? else {
167            return Ok(None);
168        };
169        let Some(reservation) = &head.assignment_drain_reservation else {
170            return Ok(None);
171        };
172        self.materialize_drain_reservation(reservation, assignments)
173            .await
174            .map(Some)
175    }
176
177    pub(super) async fn materialize_drain_reservation(
178        &self,
179        reservation: &AssignmentDrainReservation,
180        assignments: &AssignmentSnapshotStore,
181    ) -> Result<RotateOutcome, ClusterCheckpointAuthorityError> {
182        let anchor = read_authority_record(self.store.as_ref(), reservation.authority_sequence)
183            .await?
184            .ok_or_else(|| {
185                LeaseError::Invalid("assignment reservation authority anchor is missing".into())
186            })?;
187        if anchor.assignment_drain_reservation.as_ref() != Some(reservation) {
188            return Err(LeaseError::Invalid(
189                "assignment reservation differs from its authority anchor".into(),
190            )
191            .into());
192        }
193        let proposal = self.load_drain_proposal(&reservation.proposal).await?;
194        if proposal.drain_transition.as_ref() != Some(&reservation.transition) {
195            return Err(LeaseError::Invalid(
196                "reserved drain transition differs from its proposal".into(),
197            )
198            .into());
199        }
200        let outcome = assignments
201            .save_if_version(
202                &proposal,
203                reservation.transition.predecessor.assignment_version,
204            )
205            .await
206            .map_err(snapshot_authority_error)?;
207        if matches!(&outcome, RotateOutcome::Conflict(winner) if winner.draining && winner.as_ref() != &proposal)
208        {
209            return Err(LeaseError::Invalid(
210                "reserved drain lost to a different drain proposal".into(),
211            )
212            .into());
213        }
214        Ok(outcome)
215    }
216
217    pub(super) async fn prune_drain_proposals(&self, before: u64) -> Result<(), LeaseError> {
218        if before == 0 {
219            return Ok(());
220        }
221        let prefix = OsPath::from(DRAIN_PROPOSAL_PREFIX);
222        let mut objects = self.store.list(Some(&prefix));
223        let mut candidates = Vec::new();
224        let mut scanned = 0;
225        while let Some(object) = objects.next().await {
226            let object = object.map_err(|e| LeaseError::Io(e.to_string()))?;
227            let suffix = object
228                .location
229                .as_ref()
230                .strip_prefix(DRAIN_PROPOSAL_PREFIX)
231                .ok_or_else(|| {
232                    LeaseError::Invalid("drain proposal escaped its namespace".into())
233                })?;
234            let (version, digest) = suffix
235                .split_once('/')
236                .ok_or_else(|| LeaseError::Invalid("invalid drain proposal path".into()))?;
237            let digits = version
238                .strip_prefix('v')
239                .ok_or_else(|| LeaseError::Invalid("invalid drain proposal version".into()))?;
240            let hash = digest
241                .strip_suffix(".json")
242                .ok_or_else(|| LeaseError::Invalid("invalid drain proposal suffix".into()))?;
243            if digits.len() != 20
244                || !digits.bytes().all(|byte| byte.is_ascii_digit())
245                || hash.len() != 64
246                || !hash
247                    .bytes()
248                    .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
249            {
250                return Err(LeaseError::Invalid(
251                    "noncanonical drain proposal path".into(),
252                ));
253            }
254            let version = digits
255                .parse::<u64>()
256                .map_err(|e| LeaseError::Invalid(e.to_string()))?;
257            if version < before {
258                candidates.push(object.location);
259            }
260            scanned += 1;
261            if candidates.len() == 64 || scanned == 256 {
262                break;
263            }
264        }
265        drop(objects);
266        // A later scheduled retention pass resumes this bounded batch. The admitted floor has
267        // already excluded every candidate from checkpoint/assignment reuse.
268        for candidate in candidates {
269            match self.store.delete(&candidate).await {
270                Ok(()) | Err(object_store::Error::NotFound { .. }) => {}
271                Err(error) => return Err(LeaseError::Io(error.to_string())),
272            }
273        }
274        Ok(())
275    }
276}
277
278fn snapshot_authority_error(error: SnapshotError) -> ClusterCheckpointAuthorityError {
279    match error {
280        SnapshotError::Io(reason) => LeaseError::Io(reason).into(),
281        error => LeaseError::Invalid(error.to_string()).into(),
282    }
283}
laminar_core/cluster/control/leader_lease/assignment_drain.rs

laminar_core/cluster/control/leader_lease/
assignment_drain.rs
