1use uuid::Uuid;
4
5use super::LeaderProof;
6use crate::state::{KeyGroupCount, PARTITIONING_ABI_VERSION};
7
8pub const MAX_CHECKPOINT_PARTICIPANTS: usize = 128 + 1;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
17pub struct CheckpointParticipant {
18 pub node_id: u64,
20 pub boot_incarnation: Uuid,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
26#[serde(try_from = "UncheckedCheckpointAssignmentAdoption")]
27pub struct CheckpointAssignmentAdoption {
28 pub participant: CheckpointParticipant,
30 pub assignment_version: u64,
32 pub partitioning_abi_version: u16,
34 pub vnode_count: u32,
36 pub assignment_digest: [u8; 32],
38}
39
40#[derive(serde::Deserialize)]
41struct UncheckedCheckpointAssignmentAdoption {
42 participant: CheckpointParticipant,
43 assignment_version: u64,
44 partitioning_abi_version: u16,
45 vnode_count: u32,
46 assignment_digest: [u8; 32],
47}
48
49impl TryFrom<UncheckedCheckpointAssignmentAdoption> for CheckpointAssignmentAdoption {
50 type Error = &'static str;
51
52 fn try_from(unchecked: UncheckedCheckpointAssignmentAdoption) -> Result<Self, Self::Error> {
53 let adoption = Self {
54 participant: unchecked.participant,
55 assignment_version: unchecked.assignment_version,
56 partitioning_abi_version: unchecked.partitioning_abi_version,
57 vnode_count: unchecked.vnode_count,
58 assignment_digest: unchecked.assignment_digest,
59 };
60 adoption
61 .is_canonical()
62 .then_some(adoption)
63 .ok_or("checkpoint assignment adoption is not canonical")
64 }
65}
66
67impl CheckpointAssignmentAdoption {
68 #[must_use]
70 pub fn is_canonical(&self) -> bool {
71 self.participant.node_id != 0
72 && !self.participant.boot_incarnation.is_nil()
73 && self.assignment_version != 0
74 && self.partitioning_abi_version == PARTITIONING_ABI_VERSION
75 && KeyGroupCount::try_from(self.vnode_count).is_ok()
76 && self.assignment_digest != [0; 32]
77 }
78
79 #[must_use]
81 pub fn matches_fence(&self, fence: &CheckpointAssignmentFence) -> bool {
82 self.assignment_version == fence.assignment_version
83 && self.partitioning_abi_version == fence.partitioning_abi_version
84 && self.vnode_count == fence.vnode_count
85 && self.assignment_digest == fence.assignment_digest
86 && fence.participant_incarnation(self.participant.node_id)
87 == Some(self.participant.boot_incarnation)
88 }
89}
90
91#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
93#[serde(try_from = "UncheckedCheckpointAssignmentFence")]
94pub struct CheckpointAssignmentFence {
95 pub assignment_version: u64,
97 pub partitioning_abi_version: u16,
99 pub vnode_count: u32,
101 pub assignment_digest: [u8; 32],
103 pub participants: Vec<CheckpointParticipant>,
105}
106
107#[derive(serde::Deserialize)]
108struct UncheckedCheckpointAssignmentFence {
109 assignment_version: u64,
110 partitioning_abi_version: u16,
111 vnode_count: u32,
112 assignment_digest: [u8; 32],
113 participants: Vec<CheckpointParticipant>,
114}
115
116impl TryFrom<UncheckedCheckpointAssignmentFence> for CheckpointAssignmentFence {
117 type Error = &'static str;
118
119 fn try_from(unchecked: UncheckedCheckpointAssignmentFence) -> Result<Self, Self::Error> {
120 let fence = Self {
121 assignment_version: unchecked.assignment_version,
122 partitioning_abi_version: unchecked.partitioning_abi_version,
123 vnode_count: unchecked.vnode_count,
124 assignment_digest: unchecked.assignment_digest,
125 participants: unchecked.participants,
126 };
127 fence
128 .is_canonical()
129 .then_some(fence)
130 .ok_or("checkpoint assignment certificate is not canonical")
131 }
132}
133
134impl CheckpointAssignmentFence {
135 pub fn from_owner_map(
140 assignment_version: u64,
141 owners: &[u64],
142 participants: Vec<CheckpointParticipant>,
143 ) -> Result<Self, String> {
144 if participants.len() > MAX_CHECKPOINT_PARTICIPANTS {
145 return Err(format!(
146 "checkpoint assignment has {} participants; maximum is {MAX_CHECKPOINT_PARTICIPANTS}",
147 participants.len()
148 ));
149 }
150 let vnode_count = u32::try_from(owners.len())
151 .map_err(|_| "checkpoint assignment has more than u32::MAX vnodes".to_string())?;
152 KeyGroupCount::try_from(vnode_count).map_err(|_| {
153 format!(
154 "checkpoint assignment key-group count must be between 1 and {}, got {vnode_count}",
155 crate::state::MAX_KEY_GROUP_COUNT
156 )
157 })?;
158 let fence = Self {
159 assignment_version,
160 partitioning_abi_version: PARTITIONING_ABI_VERSION,
161 vnode_count,
162 assignment_digest: Self::owner_map_digest(vnode_count, owners),
163 participants,
164 };
165 if !fence.is_canonical() {
166 return Err("checkpoint assignment certificate is not canonical".into());
167 }
168 let owner_ids: std::collections::BTreeSet<u64> = owners.iter().copied().collect();
169 if owner_ids.len() != fence.participants.len()
170 || fence
171 .participants
172 .iter()
173 .any(|participant| !owner_ids.contains(&participant.node_id))
174 {
175 return Err("checkpoint participants must be the exact vnode-owner set".into());
176 }
177 Ok(fence)
178 }
179
180 #[must_use]
182 pub fn owner_map_digest(vnode_count: u32, owners: &[u64]) -> [u8; 32] {
183 Self::owner_map_digest_for_abi(PARTITIONING_ABI_VERSION, vnode_count, owners)
184 }
185
186 fn owner_map_digest_for_abi(
187 partitioning_abi_version: u16,
188 vnode_count: u32,
189 owners: &[u64],
190 ) -> [u8; 32] {
191 use sha2::{Digest, Sha256};
192
193 let mut hash = Sha256::new();
194 hash.update(b"laminardb-vnode-owner-map-v2\0");
195 hash.update(partitioning_abi_version.to_le_bytes());
196 hash.update(vnode_count.to_le_bytes());
197 hash.update(
198 u64::try_from(owners.len())
199 .unwrap_or(u64::MAX)
200 .to_le_bytes(),
201 );
202 for owner in owners {
203 hash.update(owner.to_le_bytes());
204 }
205 hash.finalize().into()
206 }
207
208 #[must_use]
210 pub fn is_canonical(&self) -> bool {
211 self.assignment_version != 0
212 && self.partitioning_abi_version == PARTITIONING_ABI_VERSION
213 && KeyGroupCount::try_from(self.vnode_count).is_ok()
214 && self.assignment_digest != [0; 32]
215 && !self.participants.is_empty()
216 && self.participants.len() <= MAX_CHECKPOINT_PARTICIPANTS
217 && self
218 .participants
219 .windows(2)
220 .all(|pair| pair[0].node_id < pair[1].node_id)
221 && self.participants.iter().all(|participant| {
222 participant.node_id != 0 && !participant.boot_incarnation.is_nil()
223 })
224 }
225
226 #[must_use]
228 pub fn contains(&self, node_id: u64) -> bool {
229 self.participants
230 .binary_search_by_key(&node_id, |participant| participant.node_id)
231 .is_ok()
232 }
233
234 #[must_use]
236 pub fn participant_incarnation(&self, node_id: u64) -> Option<Uuid> {
237 self.participants
238 .binary_search_by_key(&node_id, |participant| participant.node_id)
239 .ok()
240 .map(|index| self.participants[index].boot_incarnation)
241 }
242
243 #[must_use]
245 pub fn participant_ids(&self) -> Vec<u64> {
246 self.participants
247 .iter()
248 .map(|participant| participant.node_id)
249 .collect()
250 }
251
252 #[must_use]
254 pub fn matches_owner_map(&self, owners: &[u64]) -> bool {
255 self.partitioning_abi_version == PARTITIONING_ABI_VERSION
256 && usize::try_from(self.vnode_count).ok() == Some(owners.len())
257 && self.assignment_digest == Self::owner_map_digest(self.vnode_count, owners)
258 && {
259 let owner_ids: std::collections::BTreeSet<u64> = owners.iter().copied().collect();
260 owner_ids.len() == self.participants.len()
261 && self
262 .participants
263 .iter()
264 .all(|participant| owner_ids.contains(&participant.node_id))
265 }
266 }
267
268 #[must_use]
270 pub fn digest(&self) -> [u8; 32] {
271 use sha2::{Digest, Sha256};
272
273 let mut hash = Sha256::new();
274 hash.update(b"laminardb-checkpoint-assignment-v3\0");
275 hash.update(self.assignment_version.to_le_bytes());
276 hash.update(self.partitioning_abi_version.to_le_bytes());
277 hash.update(self.vnode_count.to_le_bytes());
278 hash.update(self.assignment_digest);
279 hash.update(
280 u64::try_from(self.participants.len())
281 .unwrap_or(u64::MAX)
282 .to_le_bytes(),
283 );
284 for participant in &self.participants {
285 hash.update(participant.node_id.to_le_bytes());
286 hash.update(participant.boot_incarnation.as_bytes());
287 }
288 hash.finalize().into()
289 }
290}
291
292#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
294pub struct AssignmentDrainId {
295 pub predecessor_version: u64,
297 pub target_version: u64,
299 pub digest: [u8; 32],
301}
302
303impl AssignmentDrainId {
304 #[must_use]
306 pub fn is_canonical(self) -> bool {
307 self.predecessor_version != 0
308 && self.predecessor_version.checked_add(1) == Some(self.target_version)
309 && self.digest != [0; 32]
310 }
311}
312
313#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
318#[serde(try_from = "UncheckedAssignmentDrainTransition")]
319pub struct AssignmentDrainTransition {
320 pub predecessor: CheckpointAssignmentFence,
322 pub target: CheckpointAssignmentFence,
324 pub leader: LeaderProof,
326}
327
328#[derive(serde::Deserialize)]
329struct UncheckedAssignmentDrainTransition {
330 predecessor: CheckpointAssignmentFence,
331 target: CheckpointAssignmentFence,
332 leader: LeaderProof,
333}
334
335impl TryFrom<UncheckedAssignmentDrainTransition> for AssignmentDrainTransition {
336 type Error = &'static str;
337
338 fn try_from(unchecked: UncheckedAssignmentDrainTransition) -> Result<Self, Self::Error> {
339 Self::new(unchecked.predecessor, unchecked.target, unchecked.leader)
340 .map_err(|_| "assignment drain transition is not canonical")
341 }
342}
343
344impl AssignmentDrainTransition {
345 pub fn new(
351 predecessor: CheckpointAssignmentFence,
352 target: CheckpointAssignmentFence,
353 leader: LeaderProof,
354 ) -> Result<Self, String> {
355 let transition = Self {
356 predecessor,
357 target,
358 leader,
359 };
360 if !transition.is_canonical() {
361 return Err("assignment drain transition is not canonical".into());
362 }
363 Ok(transition)
364 }
365
366 #[must_use]
368 pub fn is_canonical(&self) -> bool {
369 self.predecessor.is_canonical()
370 && self.target.is_canonical()
371 && self.leader.is_canonical()
372 && self.predecessor.vnode_count == self.target.vnode_count
373 && self.predecessor.assignment_version.checked_add(1)
374 == Some(self.target.assignment_version)
375 && (self
376 .predecessor
377 .participant_incarnation(self.leader.owner.node_id)
378 == Some(self.leader.owner.boot_id)
379 || self
380 .target
381 .participant_incarnation(self.leader.owner.node_id)
382 == Some(self.leader.owner.boot_id))
383 }
384
385 #[must_use]
387 pub fn id(&self) -> AssignmentDrainId {
388 AssignmentDrainId {
389 predecessor_version: self.predecessor.assignment_version,
390 target_version: self.target.assignment_version,
391 digest: self.digest(),
392 }
393 }
394
395 #[must_use]
397 pub fn required_participants(&self) -> &[CheckpointParticipant] {
398 &self.predecessor.participants
399 }
400
401 #[must_use]
403 pub fn digest(&self) -> [u8; 32] {
404 use sha2::{Digest, Sha256};
405
406 let mut hash = Sha256::new();
407 hash.update(b"laminardb-assignment-drain-transition-v1\0");
408 hash.update(self.predecessor.digest());
409 hash.update(self.target.digest());
410 hash.update(self.leader.owner.node_id.to_le_bytes());
411 hash.update(self.leader.owner.boot_id.as_bytes());
412 hash.update(self.leader.owner.process_term.to_le_bytes());
413 hash.update(self.leader.fencing_token.to_le_bytes());
414 hash.finalize().into()
415 }
416}
417
418#[cfg(test)]
419mod tests {
420 use super::{
421 AssignmentDrainTransition, CheckpointAssignmentAdoption, CheckpointAssignmentFence,
422 CheckpointParticipant, MAX_CHECKPOINT_PARTICIPANTS,
423 };
424 use crate::checkpoint::{LeaderProof, LeaderProofOwner};
425 use crate::state::{MAX_KEY_GROUP_COUNT, PARTITIONING_ABI_VERSION};
426 use uuid::Uuid;
427
428 fn participant(node_id: u64, boot: u128) -> CheckpointParticipant {
429 CheckpointParticipant {
430 node_id,
431 boot_incarnation: Uuid::from_u128(boot),
432 }
433 }
434
435 fn leader(node_id: u64, boot: u128) -> LeaderProof {
436 LeaderProof {
437 owner: LeaderProofOwner {
438 node_id,
439 boot_id: Uuid::from_u128(boot),
440 process_term: 4,
441 },
442 fencing_token: 9,
443 }
444 }
445
446 #[test]
447 fn certificate_binds_map_version_and_process_roster() {
448 let fence = CheckpointAssignmentFence::from_owner_map(
449 7,
450 &[1, 2, 1, 9],
451 vec![participant(1, 11), participant(2, 22), participant(9, 99)],
452 )
453 .unwrap();
454 assert!(fence.is_canonical());
455 assert!(fence.matches_owner_map(&[1, 2, 1, 9]));
456 assert_eq!(fence.participant_ids(), [1, 2, 9]);
457
458 let different_map =
459 CheckpointAssignmentFence::from_owner_map(7, &[2, 1, 1, 9], fence.participants.clone())
460 .unwrap();
461 assert_ne!(fence.digest(), different_map.digest());
462
463 let restarted = CheckpointAssignmentFence::from_owner_map(
464 7,
465 &[1, 2, 1, 9],
466 vec![participant(1, 111), participant(2, 22), participant(9, 99)],
467 )
468 .unwrap();
469 assert_ne!(fence.digest(), restarted.digest());
470 }
471
472 #[test]
473 fn certificate_digest_binds_partitioning_abi() {
474 let fence = CheckpointAssignmentFence::from_owner_map(
475 7,
476 &[1, 2],
477 vec![participant(1, 11), participant(2, 22)],
478 )
479 .unwrap();
480 let mut wrong_abi = fence.clone();
481 wrong_abi.partitioning_abi_version = PARTITIONING_ABI_VERSION + 1;
482
483 assert_ne!(fence.digest(), wrong_abi.digest());
484 assert_ne!(
485 CheckpointAssignmentFence::owner_map_digest(2, &[1, 2]),
486 CheckpointAssignmentFence::owner_map_digest_for_abi(
487 PARTITIONING_ABI_VERSION + 1,
488 2,
489 &[1, 2]
490 )
491 );
492 assert!(!wrong_abi.is_canonical());
493 assert!(!wrong_abi.matches_owner_map(&[1, 2]));
494 }
495
496 #[test]
497 fn certificate_and_adoption_require_current_partitioning_abi() {
498 let fence =
499 CheckpointAssignmentFence::from_owner_map(7, &[1], vec![participant(1, 11)]).unwrap();
500 let mut missing_fence = serde_json::to_value(&fence).unwrap();
501 missing_fence
502 .as_object_mut()
503 .unwrap()
504 .remove("partitioning_abi_version");
505 assert!(serde_json::from_value::<CheckpointAssignmentFence>(missing_fence).is_err());
506
507 let mut wrong_fence = fence.clone();
508 wrong_fence.partitioning_abi_version = PARTITIONING_ABI_VERSION + 1;
509 assert!(serde_json::from_value::<CheckpointAssignmentFence>(
510 serde_json::to_value(wrong_fence).unwrap()
511 )
512 .is_err());
513
514 let adoption = CheckpointAssignmentAdoption {
515 participant: participant(1, 11),
516 assignment_version: fence.assignment_version,
517 partitioning_abi_version: PARTITIONING_ABI_VERSION,
518 vnode_count: fence.vnode_count,
519 assignment_digest: fence.assignment_digest,
520 };
521 assert!(adoption.is_canonical());
522 assert!(adoption.matches_fence(&fence));
523
524 let mut missing_adoption = serde_json::to_value(&adoption).unwrap();
525 missing_adoption
526 .as_object_mut()
527 .unwrap()
528 .remove("partitioning_abi_version");
529 assert!(serde_json::from_value::<CheckpointAssignmentAdoption>(missing_adoption).is_err());
530
531 let mut wrong_adoption = adoption;
532 wrong_adoption.partitioning_abi_version = PARTITIONING_ABI_VERSION + 1;
533 assert!(serde_json::from_value::<CheckpointAssignmentAdoption>(
534 serde_json::to_value(wrong_adoption).unwrap()
535 )
536 .is_err());
537 }
538
539 #[test]
540 fn malformed_or_incomplete_certificates_fail_closed() {
541 assert!(
542 CheckpointAssignmentFence::from_owner_map(7, &[1, 2], vec![participant(1, 11)])
543 .is_err()
544 );
545 assert!(CheckpointAssignmentFence::from_owner_map(
546 7,
547 &[1],
548 vec![participant(1, 11), participant(1, 12)]
549 )
550 .is_err());
551 assert!(
552 CheckpointAssignmentFence::from_owner_map(0, &[1], vec![participant(1, 11)]).is_err()
553 );
554 assert!(matches!(
555 CheckpointAssignmentFence::from_owner_map(
556 7,
557 &[1],
558 vec![participant(1, 11), participant(2, 22)]
559 ),
560 Err(message) if message.contains("exact vnode-owner set")
561 ));
562 }
563
564 #[test]
565 fn certificate_rejects_more_than_the_partitioning_abi_limit() {
566 let owner_count = usize::try_from(MAX_KEY_GROUP_COUNT).unwrap() + 1;
567 let owners = vec![1; owner_count];
568
569 assert!(matches!(
570 CheckpointAssignmentFence::from_owner_map(7, &owners, vec![participant(1, 11)]),
571 Err(message) if message.contains("key-group count")
572 ));
573 }
574
575 #[test]
576 fn certificate_participant_limit_accepts_129_and_rejects_130() {
577 let maximum = u64::try_from(MAX_CHECKPOINT_PARTICIPANTS).unwrap();
578 let participants = (1..=maximum)
579 .map(|node_id| participant(node_id, u128::from(node_id)))
580 .collect();
581 let owners = (1..=maximum).collect::<Vec<_>>();
582 let fence = CheckpointAssignmentFence::from_owner_map(7, &owners, participants).unwrap();
583 assert!(fence.is_canonical());
584 assert_eq!(fence.participants.len(), MAX_CHECKPOINT_PARTICIPANTS);
585
586 let oversized = (1..=maximum + 1)
587 .map(|node_id| participant(node_id, u128::from(node_id)))
588 .collect();
589 let oversized_owners = (1..=maximum + 1).collect::<Vec<_>>();
590 assert!(matches!(
591 CheckpointAssignmentFence::from_owner_map(8, &oversized_owners, oversized),
592 Err(message) if message.contains("maximum is 129")
593 ));
594
595 let mut forged = fence;
596 forged
597 .participants
598 .push(participant(maximum + 1, u128::from(maximum + 1)));
599 assert!(!forged.is_canonical());
600
601 let encoded = serde_json::to_vec(&forged).unwrap();
602 assert!(serde_json::from_slice::<CheckpointAssignmentFence>(&encoded).is_err());
603 }
604
605 #[test]
606 fn drain_transition_acks_predecessor_roster_and_accepts_target_only_leader() {
607 let predecessor = CheckpointAssignmentFence::from_owner_map(
608 7,
609 &[1, 2],
610 vec![participant(1, 11), participant(2, 22)],
611 )
612 .unwrap();
613 let target = CheckpointAssignmentFence::from_owner_map(
614 8,
615 &[2, 3],
616 vec![participant(2, 22), participant(3, 33)],
617 )
618 .unwrap();
619 let transition =
620 AssignmentDrainTransition::new(predecessor, target, leader(3, 33)).unwrap();
621
622 assert_eq!(
623 transition
624 .required_participants()
625 .iter()
626 .map(|participant| participant.node_id)
627 .collect::<Vec<_>>(),
628 [1, 2]
629 );
630 assert!(transition.id().is_canonical());
631 assert_ne!(transition.digest(), [0; 32]);
632 }
633
634 #[test]
635 fn drain_transition_accepts_predecessor_only_leader() {
636 let predecessor = CheckpointAssignmentFence::from_owner_map(
637 7,
638 &[1, 2],
639 vec![participant(1, 11), participant(2, 22)],
640 )
641 .unwrap();
642 let target = CheckpointAssignmentFence::from_owner_map(
643 8,
644 &[2, 3],
645 vec![participant(2, 22), participant(3, 33)],
646 )
647 .unwrap();
648 assert!(AssignmentDrainTransition::new(predecessor, target, leader(1, 11)).is_ok());
649 }
650
651 #[test]
652 fn drain_identity_rejects_version_overflow() {
653 let identity = super::AssignmentDrainId {
654 predecessor_version: u64::MAX,
655 target_version: u64::MAX,
656 digest: [1; 32],
657 };
658 assert!(!identity.is_canonical());
659 }
660}