laminar_core/cluster/control/snapshot/
model.rs

1//! Canonical assignment snapshots, recovery references, and drain finalizations.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use serde::{Deserialize, Serialize};
6use sha2::{Digest, Sha256};
7
8use super::{
9    current_time_millis, SnapshotError, DRAIN_FINALIZATION_VERSION, MAX_RECOVERY_PROPOSAL_BYTES,
10    RECOVERY_MATERIALIZATION_VERSION,
11};
12use crate::checkpoint::{
13    AssignmentDrainTransition, CheckpointAssignmentFence, CheckpointParticipant, LeaderProof,
14    MAX_CHECKPOINT_PARTICIPANTS,
15};
16use crate::cluster::discovery::NodeId;
17use crate::state::{KeyGroupCount, PARTITIONING_ABI_VERSION};
18
19/// Durable vnode-to-instance assignment snapshot.
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(deny_unknown_fields)]
22pub struct AssignmentSnapshot {
23    /// Monotonic version. Writers bump on each update.
24    pub version: u64,
25    /// Key encoding, hashing, and key-group mapping contract used by this owner map.
26    pub partitioning_abi_version: u16,
27    /// Vnode id → owning instance. `BTreeMap` (not `Vec`) so snapshots
28    /// with different `vnode_count` are still deserializable — sparse
29    /// indices surface as missing keys the caller can diagnose.
30    pub vnodes: BTreeMap<u32, NodeId>,
31    /// Exact process roster certified for this assignment generation. A stable node restart
32    /// changes this roster and therefore requires a new version even when vnode owners do not.
33    pub participants: Vec<CheckpointParticipant>,
34    /// Wall-clock timestamp of the last update, millis since epoch.
35    pub updated_at_ms: i64,
36    /// Pre-rotation drain phase: when set, this snapshot carries the *intended*
37    /// next assignment but ownership has NOT changed yet. Nodes mark the vnodes
38    /// they are about to lose as draining (pausing those source partitions) so the
39    /// pre-rotation checkpoint is a clean cut; the leader then publishes the same
40    /// assignment with `draining = false` to commit the rotation.
41    pub draining: bool,
42    /// Exact predecessor, successor, and durable leader term for a draining generation.
43    /// Present if and only if `draining` is true.
44    #[serde(deserialize_with = "deserialize_drain_transition")]
45    pub drain_transition: Option<AssignmentDrainTransition>,
46}
47
48fn deserialize_drain_transition<'de, D>(
49    deserializer: D,
50) -> Result<Option<AssignmentDrainTransition>, D::Error>
51where
52    D: serde::Deserializer<'de>,
53{
54    Option::deserialize(deserializer)
55}
56
57/// Exact immutable reference to a staged failure-recovery successor.
58#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
59#[serde(deny_unknown_fields)]
60pub struct AssignmentSnapshotRef {
61    /// Assignment generation carried by the referenced snapshot.
62    pub version: u64,
63    /// Lowercase hexadecimal SHA-256 of its canonical JSON body.
64    pub sha256: String,
65    /// Exact canonical body length.
66    pub encoded_len: u64,
67}
68
69/// Create-only logical assignment head for an authority-selected recovery proposal.
70///
71/// Recovery uses a separate namespace from graceful-drain intents. A delayed drain writer may
72/// still win the raw snapshot key after losing leadership, but it cannot replace this immutable
73/// materialization winner for the same assignment version.
74#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
75#[serde(deny_unknown_fields)]
76pub(super) struct RecoveryMaterialization {
77    protocol_version: u16,
78    pub(super) proposal: AssignmentSnapshotRef,
79    pub(super) snapshot: AssignmentSnapshot,
80}
81
82impl RecoveryMaterialization {
83    pub(super) fn new(
84        proposal: AssignmentSnapshotRef,
85        snapshot: AssignmentSnapshot,
86    ) -> Result<Self, SnapshotError> {
87        let materialization = Self {
88            protocol_version: RECOVERY_MATERIALIZATION_VERSION,
89            proposal,
90            snapshot,
91        };
92        materialization.validate()?;
93        Ok(materialization)
94    }
95
96    pub(super) fn validate(&self) -> Result<(), SnapshotError> {
97        self.proposal.validate()?;
98        if self.protocol_version != RECOVERY_MATERIALIZATION_VERSION {
99            return Err(SnapshotError::Invalid(format!(
100                "unsupported recovery materialization version {}",
101                self.protocol_version
102            )));
103        }
104        let (_, actual_reference) = self.snapshot.encode_recovery_proposal()?;
105        if actual_reference != self.proposal {
106            return Err(SnapshotError::Invalid(
107                "recovery materialization body does not match its proposal reference".into(),
108            ));
109        }
110        Ok(())
111    }
112}
113
114impl AssignmentSnapshotRef {
115    /// Validate the reference independently of the staged object.
116    ///
117    /// # Errors
118    /// Rejects a non-successor version, malformed digest, or unsafe encoded length.
119    pub fn validate(&self) -> Result<(), SnapshotError> {
120        if self.version < 2 {
121            return Err(SnapshotError::Invalid(
122                "recovery proposal must be a successor generation".into(),
123            ));
124        }
125        if self.sha256.len() != 64
126            || !self
127                .sha256
128                .bytes()
129                .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
130        {
131            return Err(SnapshotError::Invalid(
132                "recovery proposal SHA-256 must be 64 lowercase hexadecimal characters".into(),
133            ));
134        }
135        let encoded_len = usize::try_from(self.encoded_len).map_err(|_| {
136            SnapshotError::Invalid("recovery proposal encoded length exceeds usize".into())
137        })?;
138        if encoded_len == 0 || encoded_len > MAX_RECOVERY_PROPOSAL_BYTES {
139            return Err(SnapshotError::Invalid(format!(
140                "recovery proposal encoded length {} is outside 1..={MAX_RECOVERY_PROPOSAL_BYTES}",
141                self.encoded_len
142            )));
143        }
144        Ok(())
145    }
146}
147
148/// Immutable winner that settles one draining snapshot without changing its certified target
149/// version. The original transition remains append-only at the snapshot path.
150#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
151#[serde(deny_unknown_fields)]
152pub(super) struct DrainFinalization {
153    protocol_version: u16,
154    transition_digest: [u8; 32],
155    pub(super) proposal: AssignmentSnapshot,
156}
157
158impl DrainFinalization {
159    pub(super) fn new(
160        draining: &AssignmentSnapshot,
161        proposal: AssignmentSnapshot,
162    ) -> Result<Self, SnapshotError> {
163        let transition = draining.drain_transition.as_ref().ok_or_else(|| {
164            SnapshotError::Invalid("drain finalization requires a draining transition".into())
165        })?;
166        let finalization = Self {
167            protocol_version: DRAIN_FINALIZATION_VERSION,
168            transition_digest: transition.digest(),
169            proposal,
170        };
171        finalization.validate_against(draining)?;
172        Ok(finalization)
173    }
174
175    pub(super) fn validate_against(
176        &self,
177        draining: &AssignmentSnapshot,
178    ) -> Result<(), SnapshotError> {
179        draining.validate()?;
180        self.proposal.validate()?;
181        let transition = draining.drain_transition.as_ref().ok_or_else(|| {
182            SnapshotError::Invalid("drain finalization requires a draining transition".into())
183        })?;
184        if self.protocol_version != DRAIN_FINALIZATION_VERSION
185            || self.transition_digest != transition.digest()
186            || !draining.draining
187            || self.proposal.draining
188            || self.proposal.version != draining.version
189            || self.proposal.drain_transition.is_some()
190        {
191            return Err(SnapshotError::Invalid(
192                "drain finalization does not preserve the exact transition identity".into(),
193            ));
194        }
195        let proposed_fence = self.proposal.assignment_fence()?;
196        let predecessor = &transition.predecessor;
197        let commits_target = proposed_fence == transition.target;
198        let aborts_to_predecessor = proposed_fence.assignment_version
199            == transition.target.assignment_version
200            && proposed_fence.vnode_count == predecessor.vnode_count
201            && proposed_fence.assignment_digest == predecessor.assignment_digest
202            && proposed_fence.participants == predecessor.participants;
203        if !commits_target && !aborts_to_predecessor {
204            return Err(SnapshotError::Invalid(
205                "drain finalization is neither the certified target nor exact predecessor rollback"
206                    .into(),
207            ));
208        }
209        Ok(())
210    }
211}
212
213impl AssignmentSnapshot {
214    /// Empty snapshot at version 0.
215    #[must_use]
216    pub fn empty() -> Self {
217        Self {
218            version: 0,
219            partitioning_abi_version: PARTITIONING_ABI_VERSION,
220            vnodes: BTreeMap::new(),
221            participants: Vec::new(),
222            updated_at_ms: 0,
223            draining: false,
224            drain_transition: None,
225        }
226    }
227
228    /// Next snapshot with bumped version and current wall-clock time.
229    ///
230    /// # Errors
231    ///
232    /// Returns an error if the version overflows or the successor is noncanonical.
233    pub fn next(&self, vnodes: BTreeMap<u32, NodeId>) -> Result<Self, SnapshotError> {
234        self.next_for_participants(vnodes, self.participants.clone())
235    }
236
237    /// Next snapshot bound to the supplied canonical process roster.
238    ///
239    /// # Errors
240    ///
241    /// Returns an error if the version overflows or the successor is noncanonical.
242    pub fn next_for_participants(
243        &self,
244        vnodes: BTreeMap<u32, NodeId>,
245        participants: Vec<CheckpointParticipant>,
246    ) -> Result<Self, SnapshotError> {
247        let version = self
248            .version
249            .checked_add(1)
250            .ok_or_else(|| SnapshotError::Invalid("assignment snapshot version overflow".into()))?;
251        let next = Self {
252            version,
253            partitioning_abi_version: self.partitioning_abi_version,
254            vnodes,
255            participants,
256            updated_at_ms: current_time_millis(),
257            draining: false,
258            drain_transition: None,
259        };
260        next.validate()?;
261        Ok(next)
262    }
263
264    fn validate_assignment(&self) -> Result<(), SnapshotError> {
265        if self.partitioning_abi_version != PARTITIONING_ABI_VERSION {
266            return Err(SnapshotError::Invalid(format!(
267                "assignment snapshot partitioning ABI {} does not match runtime ABI {PARTITIONING_ABI_VERSION}",
268                self.partitioning_abi_version
269            )));
270        }
271        if self.participants.len() > MAX_CHECKPOINT_PARTICIPANTS {
272            return Err(SnapshotError::Invalid(format!(
273                "assignment snapshot has {} participants; maximum is {MAX_CHECKPOINT_PARTICIPANTS}",
274                self.participants.len()
275            )));
276        }
277        let vnode_count = u32::try_from(self.vnodes.len()).map_err(|_| {
278            SnapshotError::Invalid("assignment snapshot has more than u32::MAX key groups".into())
279        })?;
280        KeyGroupCount::try_from(vnode_count).map_err(|_| {
281            SnapshotError::Invalid(format!(
282                "assignment snapshot key-group count must be between 1 and {}, got {vnode_count}",
283                crate::state::MAX_KEY_GROUP_COUNT
284            ))
285        })?;
286        let dense = !self.vnodes.is_empty()
287            && self
288                .vnodes
289                .keys()
290                .copied()
291                .zip(0_u32..)
292                .all(|(actual, expected)| actual == expected);
293        let canonical_participants = !self.participants.is_empty()
294            && self
295                .participants
296                .windows(2)
297                .all(|pair| pair[0].node_id < pair[1].node_id)
298            && self.participants.iter().all(|participant| {
299                participant.node_id != 0 && !participant.boot_incarnation.is_nil()
300            })
301            && {
302                let owners: BTreeSet<u64> = self.vnodes.values().map(|owner| owner.0).collect();
303                owners.len() == self.participants.len()
304                    && self
305                        .participants
306                        .iter()
307                        .all(|participant| owners.contains(&participant.node_id))
308            };
309        if self.version == 0 || !dense || !canonical_participants {
310            return Err(SnapshotError::Invalid(
311                "assignment snapshot is not canonical".into(),
312            ));
313        }
314        Ok(())
315    }
316
317    /// Validate the durable owner map, process roster, and optional drain transition.
318    ///
319    /// # Errors
320    ///
321    /// Returns an error when the snapshot or drain transition is noncanonical.
322    pub fn validate(&self) -> Result<(), SnapshotError> {
323        self.validate_assignment()?;
324        match (self.draining, self.drain_transition.as_ref()) {
325            (false, None) => Ok(()),
326            (true, Some(transition)) => {
327                let target = self.assignment_fence_unchecked()?;
328                if !transition.is_canonical() || transition.target != target {
329                    return Err(SnapshotError::Invalid(
330                        "draining snapshot does not match its exact target transition".into(),
331                    ));
332                }
333                Ok(())
334            }
335            _ => Err(SnapshotError::Invalid(
336                "assignment drain flag and transition disagree".into(),
337            )),
338        }
339    }
340
341    pub(super) fn encode_recovery_proposal(
342        &self,
343    ) -> Result<(Vec<u8>, AssignmentSnapshotRef), SnapshotError> {
344        self.validate()?;
345        if self.draining || self.drain_transition.is_some() || self.version < 2 {
346            return Err(SnapshotError::Invalid(
347                "recovery proposal must be a committed successor generation".into(),
348            ));
349        }
350        let encoded = serde_json::to_vec(self)?;
351        if encoded.len() > MAX_RECOVERY_PROPOSAL_BYTES {
352            return Err(SnapshotError::Invalid(format!(
353                "encoded recovery proposal is {} bytes; maximum is {MAX_RECOVERY_PROPOSAL_BYTES}",
354                encoded.len()
355            )));
356        }
357        let reference = AssignmentSnapshotRef {
358            version: self.version,
359            sha256: format!("{:x}", Sha256::digest(&encoded)),
360            encoded_len: u64::try_from(encoded.len()).map_err(|_| {
361                SnapshotError::Invalid("recovery proposal encoded length overflow".into())
362            })?,
363        };
364        reference.validate()?;
365        Ok((encoded, reference))
366    }
367
368    fn assignment_fence_unchecked(&self) -> Result<CheckpointAssignmentFence, SnapshotError> {
369        let owners: Vec<u64> = self.vnodes.values().map(|owner| owner.0).collect();
370        CheckpointAssignmentFence::from_owner_map(self.version, &owners, self.participants.clone())
371            .map_err(SnapshotError::Invalid)
372    }
373
374    /// Exact checkpoint certificate represented by this snapshot.
375    ///
376    /// # Errors
377    ///
378    /// Returns an error when the assignment map or process roster is noncanonical.
379    pub fn assignment_fence(&self) -> Result<CheckpointAssignmentFence, SnapshotError> {
380        self.validate_assignment()?;
381        self.assignment_fence_unchecked()
382    }
383
384    /// Create the exact successor as a leader-fenced draining generation.
385    ///
386    /// # Errors
387    /// Rejects a non-committed predecessor, invalid successor, or leader proof outside both
388    /// certified rosters.
389    pub fn next_draining(
390        &self,
391        vnodes: BTreeMap<u32, NodeId>,
392        participants: Vec<CheckpointParticipant>,
393        leader: LeaderProof,
394    ) -> Result<Self, SnapshotError> {
395        self.validate()?;
396        if self.draining {
397            return Err(SnapshotError::Invalid(
398                "cannot start a drain from a draining assignment".into(),
399            ));
400        }
401        let predecessor = self.assignment_fence()?;
402        let mut target = self.next_for_participants(vnodes, participants)?;
403        let target_fence = target.assignment_fence()?;
404        target.drain_transition = Some(
405            AssignmentDrainTransition::new(predecessor, target_fence, leader)
406                .map_err(SnapshotError::Invalid)?,
407        );
408        target.draining = true;
409        target.validate()?;
410        Ok(target)
411    }
412
413    /// Convert a draining generation into its committed target without changing identity.
414    ///
415    /// # Errors
416    ///
417    /// Returns an error unless this is a valid draining snapshot.
418    pub fn committed_target(&self) -> Result<Self, SnapshotError> {
419        self.validate()?;
420        if !self.draining {
421            return Err(SnapshotError::Invalid(
422                "only a draining assignment has a target to commit".into(),
423            ));
424        }
425        let mut committed = self.clone();
426        committed.draining = false;
427        committed.drain_transition = None;
428        committed.updated_at_ms = current_time_millis();
429        committed.validate()?;
430        Ok(committed)
431    }
432
433    /// Convert a draining generation into a committed rollback of its predecessor map.
434    ///
435    /// # Errors
436    ///
437    /// Returns an error when either snapshot is invalid or the predecessor does not match.
438    pub fn aborted_target(&self, predecessor: &Self) -> Result<Self, SnapshotError> {
439        self.validate()?;
440        predecessor.validate()?;
441        let transition = self.drain_transition.as_ref().ok_or_else(|| {
442            SnapshotError::Invalid("draining assignment has no transition".into())
443        })?;
444        if predecessor.draining
445            || predecessor.assignment_fence()? != transition.predecessor
446            || self.version != predecessor.version.saturating_add(1)
447        {
448            return Err(SnapshotError::Invalid(
449                "drain rollback does not match the exact predecessor".into(),
450            ));
451        }
452        let mut aborted = predecessor.clone();
453        aborted.version = self.version;
454        aborted.updated_at_ms = current_time_millis();
455        aborted.draining = false;
456        aborted.drain_transition = None;
457        aborted.validate()?;
458        Ok(aborted)
459    }
460
461    /// Whether the durable process roster is canonical and covers every vnode owner.
462    #[must_use]
463    pub fn has_canonical_participants(&self) -> bool {
464        self.validate().is_ok()
465    }
466
467    /// Convert a `Vec<NodeId>` (one entry per vnode id, dense) into the
468    /// `BTreeMap` shape this snapshot uses. Mirrors the layout returned
469    /// by `rendezvous_assignment`.
470    #[must_use]
471    pub fn vnodes_from_vec(assignment: &[NodeId]) -> BTreeMap<u32, NodeId> {
472        #[allow(clippy::cast_possible_truncation)]
473        assignment
474            .iter()
475            .enumerate()
476            .map(|(i, n)| (i as u32, *n))
477            .collect()
478    }
479
480    /// Convert the canonical owner map to a dense vector of exactly `vnode_count` entries.
481    ///
482    /// # Errors
483    ///
484    /// Returns an error when the snapshot is invalid or its vnode map is not exactly dense.
485    pub fn to_vnode_vec(&self, vnode_count: u32) -> Result<Vec<NodeId>, SnapshotError> {
486        self.validate()?;
487        if usize::try_from(vnode_count).ok() != Some(self.vnodes.len()) {
488            return Err(SnapshotError::Invalid(format!(
489                "assignment {} vnode cardinality {} does not match runtime cardinality {vnode_count}",
490                self.version,
491                self.vnodes.len()
492            )));
493        }
494        (0..vnode_count)
495            .map(|v| {
496                self.vnodes.get(&v).copied().ok_or_else(|| {
497                    SnapshotError::Invalid(format!(
498                        "assignment {} is missing vnode {v}",
499                        self.version
500                    ))
501                })
502            })
503            .collect()
504    }
505}
laminar_core/cluster/control/snapshot/model.rs

laminar_core/cluster/control/snapshot/
model.rs
