1use std::collections::BTreeMap;
75use std::fmt;
76use std::path::PathBuf;
77use std::sync::{Arc, Mutex};
78
79use mongreldb_types::errors::ErrorCategory;
80use mongreldb_types::hlc::HlcTimestamp;
81use mongreldb_types::ids::{MetadataVersion, RaftGroupId, TableId, TabletId};
82use serde::{Deserialize, Serialize};
83use sha2::{Digest, Sha256};
84
85use crate::meta::MetaRejectionReason;
86use crate::node::ClusterError;
87use crate::tablet::{
88 Bound, Key, PartitionBounds, ReplicaDescriptor, ReplicaRole, RoutingError, TabletDescriptor,
89 TabletError, TabletLayout, TabletState,
90};
91
92#[derive(Debug, thiserror::Error)]
98pub enum SplitError {
99 #[error(transparent)]
101 Tablet(#[from] TabletError),
102 #[error(transparent)]
104 Fault(#[from] mongreldb_fault::Fault),
105 #[error(transparent)]
107 MetaPlane(#[from] MetaRejectionReason),
108 #[error(transparent)]
110 TabletData(#[from] TabletDataError),
111 #[error("source tablet {tablet} is in state {state}, expected Active")]
113 SourceNotActive {
114 tablet: TabletId,
116 state: TabletState,
118 },
119 #[error(
121 "cannot derive a deterministic midpoint split key from unbounded bounds; \
122 supply an explicit split key"
123 )]
124 UnboundedMidpoint,
125 #[error("no key lies strictly between the bounds' endpoints; the tablet cannot be split")]
127 UnsplittableBounds,
128 #[error("split key {key} does not partition the source bounds into two non-empty halves")]
130 InvalidSplitKey {
131 key: Key,
133 },
134 #[error("invalid split plan: {0}")]
136 InvalidPlan(String),
137 #[error("applied key {0} lies outside the source partition")]
140 KeyOutsideSource(Key),
141 #[error("source tablet {tablet} is retained by {pins} old-generation pin(s)")]
143 SourceRetained {
144 tablet: TabletId,
146 pins: usize,
148 },
149 #[error("cannot abort the split of tablet {tablet}: it already reached phase {phase}")]
153 CannotAbort {
154 tablet: TabletId,
156 phase: SplitPhase,
158 },
159}
160
161#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
166pub enum TabletDataError {
167 #[error("keyspace operation failed: {0}")]
169 Keyspace(String),
170 #[error("no staged build in progress")]
172 NoStagedBuild,
173 #[error("child state sink failed: {0}")]
175 Sink(String),
176}
177
178#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
185pub enum SplitPhase {
186 Started,
188 MarkedSplitting,
190 ChildrenCreated,
193 SnapshotPinned,
195 ChildrenBuilt,
197 CaughtUp,
200 Published,
202 SourceRetired,
205}
206
207impl SplitPhase {
208 pub fn hook_name(self) -> Option<&'static str> {
212 Some(match self {
213 Self::Started => return None,
214 Self::MarkedSplitting => "tablet.split.phase.1",
215 Self::ChildrenCreated => "tablet.split.phase.2",
216 Self::SnapshotPinned => "tablet.split.phase.3",
217 Self::ChildrenBuilt => "tablet.split.phase.4",
218 Self::CaughtUp => "tablet.split.phase.5",
219 Self::Published => "tablet.split.phase.6",
220 Self::SourceRetired => "tablet.split.phase.7",
221 })
222 }
223}
224
225impl fmt::Display for SplitPhase {
226 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227 let name = match self {
228 Self::Started => "Started",
229 Self::MarkedSplitting => "MarkedSplitting",
230 Self::ChildrenCreated => "ChildrenCreated",
231 Self::SnapshotPinned => "SnapshotPinned",
232 Self::ChildrenBuilt => "ChildrenBuilt",
233 Self::CaughtUp => "CaughtUp",
234 Self::Published => "Published",
235 Self::SourceRetired => "SourceRetired",
236 };
237 f.write_str(name)
238 }
239}
240
241#[derive(Clone, Debug, PartialEq, Eq)]
247pub enum SplitKeySelection {
248 Midpoint,
251 Explicit(Key),
254}
255
256fn bound_key(bound: &Bound<Key>) -> Option<&Key> {
258 match bound {
259 Bound::Unbounded => None,
260 Bound::Included(key) | Bound::Excluded(key) => Some(key),
261 }
262}
263
264pub fn midpoint_key(low: &Key, high: &Key) -> Option<Key> {
274 let (low, high) = (low.as_bytes(), high.as_bytes());
275 debug_assert!(low < high, "midpoint requires ordered endpoints");
276 midpoint_bytes(low, high).map(Key::from_bytes)
277}
278
279fn midpoint_bytes(low: &[u8], high: &[u8]) -> Option<Vec<u8>> {
280 let mut index = 0;
281 while index < low.len() && index < high.len() && low[index] == high[index] {
282 index += 1;
283 }
284 if index == low.len() {
285 if high.len() == index + 1 && high[index] == 0 {
290 return None;
291 }
292 let mut mid = low.to_vec();
293 mid.push(high[index] / 2);
294 return Some(mid);
295 }
296 let (low_byte, high_byte) = (low[index], high[index]);
298 if high_byte >= low_byte + 2 {
299 let mut mid = low[..index].to_vec();
301 mid.push(low_byte + (high_byte - low_byte) / 2);
302 Some(mid)
303 } else {
304 let mut mid = low.to_vec();
307 mid.push(0x80);
308 Some(mid)
309 }
310}
311
312fn choose_split_key(
314 bounds: &PartitionBounds,
315 selection: &SplitKeySelection,
316) -> Result<Key, SplitError> {
317 match selection {
318 SplitKeySelection::Explicit(key) => {
319 if bounds.split_at(key).is_none() {
320 return Err(SplitError::InvalidSplitKey { key: key.clone() });
321 }
322 Ok(key.clone())
323 }
324 SplitKeySelection::Midpoint => {
325 let (Some(low), Some(high)) = (bound_key(&bounds.low), bound_key(&bounds.high)) else {
326 return Err(SplitError::UnboundedMidpoint);
327 };
328 midpoint_key(low, high).ok_or(SplitError::UnsplittableBounds)
329 }
330 }
331}
332
333#[derive(Clone, Debug, PartialEq, Eq)]
342pub struct ChildAllocation {
343 pub tablet_id: TabletId,
345 pub raft_group_id: RaftGroupId,
347 pub replicas: Vec<ReplicaDescriptor>,
349}
350
351#[derive(Clone, Debug)]
354pub struct ChildPlan {
355 pub bounds: PartitionBounds,
357 pub layout: TabletLayout,
359 pub replicas: Vec<ReplicaDescriptor>,
361}
362
363impl ChildPlan {
364 pub fn descriptor(
367 &self,
368 table_id: TableId,
369 generation: u64,
370 state: TabletState,
371 role: ReplicaRole,
372 ) -> TabletDescriptor {
373 TabletDescriptor {
374 tablet_id: self.layout.tablet_id(),
375 database_id: mongreldb_types::ids::DatabaseId::ZERO,
376 table_id,
377 raft_group_id: self.layout.raft_group_id(),
378 partition: self.bounds.clone(),
379 replicas: self
380 .replicas
381 .iter()
382 .map(|replica| ReplicaDescriptor { role, ..*replica })
383 .collect(),
384 leader_hint: None,
385 generation,
386 state,
387 }
388 }
389
390 pub fn progress(&self) -> ChildProgress {
392 ChildProgress {
393 node_data: self.layout.node_data().to_path_buf(),
394 tablet_id: self.layout.tablet_id(),
395 raft_group_id: self.layout.raft_group_id(),
396 bounds: self.bounds.clone(),
397 replicas: self.replicas.clone(),
398 }
399 }
400}
401
402#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
405#[serde(deny_unknown_fields)]
406pub struct ChildProgress {
407 pub node_data: PathBuf,
409 pub tablet_id: TabletId,
411 pub raft_group_id: RaftGroupId,
413 pub bounds: PartitionBounds,
415 pub replicas: Vec<ReplicaDescriptor>,
417}
418
419impl ChildProgress {
420 pub fn plan(&self) -> ChildPlan {
422 ChildPlan {
423 bounds: self.bounds.clone(),
424 layout: TabletLayout::new(self.node_data.clone(), self.tablet_id, self.raft_group_id),
425 replicas: self.replicas.clone(),
426 }
427 }
428}
429
430#[derive(Clone, Debug)]
434pub struct SplitPlan {
435 pub source: TabletDescriptor,
437 pub children: [ChildPlan; 2],
439 pub split_key: Key,
441 pub split_ts: HlcTimestamp,
443}
444
445impl SplitPlan {
446 pub fn validate(&self) -> Result<(), SplitError> {
450 self.source.validate()?;
451 let (lower, upper) = self
452 .source
453 .partition
454 .split_at(&self.split_key)
455 .ok_or_else(|| SplitError::InvalidSplitKey {
456 key: self.split_key.clone(),
457 })?;
458 if self.children[0].bounds != lower || self.children[1].bounds != upper {
459 return Err(SplitError::InvalidPlan(
460 "child bounds are not the source bounds split at the split key".to_owned(),
461 ));
462 }
463 let child_ids = [
464 (
465 self.children[0].layout.tablet_id(),
466 self.children[0].layout.raft_group_id(),
467 ),
468 (
469 self.children[1].layout.tablet_id(),
470 self.children[1].layout.raft_group_id(),
471 ),
472 ];
473 if child_ids[0] == child_ids[1] || child_ids[0].0 == child_ids[1].0 {
474 return Err(SplitError::InvalidPlan(
475 "child tablet and raft group ids must be distinct".to_owned(),
476 ));
477 }
478 if child_ids.iter().any(|ids| ids.0 == self.source.tablet_id) {
479 return Err(SplitError::InvalidPlan(
480 "child tablet ids must differ from the source's".to_owned(),
481 ));
482 }
483 if child_ids
484 .iter()
485 .any(|ids| ids.1 == self.source.raft_group_id)
486 {
487 return Err(SplitError::InvalidPlan(
488 "child raft group ids must differ from the source's".to_owned(),
489 ));
490 }
491 for descriptor in self.child_descriptors() {
492 descriptor.validate()?;
493 }
494 Ok(())
495 }
496
497 pub fn child_descriptors(&self) -> [TabletDescriptor; 2] {
501 let generation = self
502 .source
503 .generation
504 .checked_add(1)
505 .expect("descriptor generation overflows u64");
506 self.children.clone().map(|child| {
507 child.descriptor(
508 self.source.table_id,
509 generation,
510 TabletState::Creating,
511 ReplicaRole::Learner,
512 )
513 })
514 }
515}
516
517#[derive(Clone, Debug)]
519pub struct TabletSplitPlanner {
520 node_data: PathBuf,
521}
522
523impl TabletSplitPlanner {
524 pub fn new(node_data: impl Into<PathBuf>) -> Self {
526 Self {
527 node_data: node_data.into(),
528 }
529 }
530
531 pub fn plan(
535 &self,
536 source: &TabletDescriptor,
537 selection: SplitKeySelection,
538 split_ts: HlcTimestamp,
539 allocations: [ChildAllocation; 2],
540 ) -> Result<SplitPlan, SplitError> {
541 if source.state != TabletState::Active {
542 return Err(SplitError::SourceNotActive {
543 tablet: source.tablet_id,
544 state: source.state,
545 });
546 }
547 source.validate()?;
548 let split_key = choose_split_key(&source.partition, &selection)?;
549 let (lower, upper) =
550 source
551 .partition
552 .split_at(&split_key)
553 .ok_or_else(|| SplitError::InvalidSplitKey {
554 key: split_key.clone(),
555 })?;
556 let [lower_alloc, upper_alloc] = allocations;
557 let child = |bounds: PartitionBounds, alloc: ChildAllocation| ChildPlan {
558 bounds,
559 layout: TabletLayout::new(self.node_data.clone(), alloc.tablet_id, alloc.raft_group_id),
560 replicas: alloc.replicas,
561 };
562 let plan = SplitPlan {
563 source: source.clone(),
564 children: [child(lower, lower_alloc), child(upper, upper_alloc)],
565 split_key,
566 split_ts,
567 };
568 plan.validate()?;
569 Ok(plan)
570 }
571}
572
573#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
586pub struct SplitPublishCommand {
587 pub source: TabletDescriptor,
589 pub children: [TabletDescriptor; 2],
594 pub split_key: Key,
596 pub split_ts: HlcTimestamp,
598}
599
600impl SplitPublishCommand {
601 pub fn from_plan(plan: &SplitPlan) -> Result<Self, SplitError> {
606 let marked = plan.source.published_transition(TabletState::Splitting)?;
607 let source = marked.published_transition(TabletState::Retiring)?;
608 let mut children = plan.child_descriptors();
609 for child in &mut children {
610 *child = child.published_transition(TabletState::Active)?;
611 for replica in &mut child.replicas {
612 replica.role = ReplicaRole::Voter;
613 }
614 }
615 let command = Self {
616 source,
617 children,
618 split_key: plan.split_key.clone(),
619 split_ts: plan.split_ts,
620 };
621 command.validate()?;
622 Ok(command)
623 }
624
625 pub fn publish_generation(&self) -> u64 {
627 self.source.generation
628 }
629
630 pub fn validate(&self) -> Result<(), SplitError> {
634 if self.source.state != TabletState::Retiring {
635 return Err(SplitError::InvalidPlan(format!(
636 "published source must be Retiring, is {}",
637 self.source.state
638 )));
639 }
640 let generation = self.source.generation;
641 for child in &self.children {
642 if child.state != TabletState::Active {
643 return Err(SplitError::InvalidPlan(format!(
644 "published child {} must be Active, is {}",
645 child.tablet_id, child.state
646 )));
647 }
648 if child.generation != generation {
649 return Err(SplitError::InvalidPlan(
650 "publication assigns one generation to all descriptors".to_owned(),
651 ));
652 }
653 if child.table_id != self.source.table_id {
654 return Err(SplitError::InvalidPlan(
655 "children and source name different tables".to_owned(),
656 ));
657 }
658 child.validate()?;
659 }
660 self.source.validate()?;
661 let (lower, upper) = self
662 .source
663 .partition
664 .split_at(&self.split_key)
665 .ok_or_else(|| SplitError::InvalidSplitKey {
666 key: self.split_key.clone(),
667 })?;
668 if self.children[0].partition != lower || self.children[1].partition != upper {
669 return Err(SplitError::InvalidPlan(
670 "published child bounds are not the source bounds split at the split key"
671 .to_owned(),
672 ));
673 }
674 if self.children[0].tablet_id == self.children[1].tablet_id
675 || self
676 .children
677 .iter()
678 .any(|c| c.tablet_id == self.source.tablet_id)
679 {
680 return Err(SplitError::InvalidPlan(
681 "publication tablet ids must be distinct".to_owned(),
682 ));
683 }
684 Ok(())
685 }
686}
687
688pub const SPLIT_PROGRESS_FILENAME: &str = "split.json";
694pub const SPLIT_PROGRESS_FORMAT_VERSION: u32 = 1;
696pub const MIN_SUPPORTED_SPLIT_PROGRESS_FORMAT_VERSION: u32 = 1;
698
699#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
704#[serde(deny_unknown_fields)]
705pub struct SplitProgress {
706 pub source: TabletDescriptor,
708 pub split_key: Key,
710 pub split_ts: HlcTimestamp,
712 pub children: [ChildProgress; 2],
714 pub phase: SplitPhase,
716}
717
718impl SplitProgress {
719 pub fn from_plan(plan: &SplitPlan, phase: SplitPhase) -> Self {
721 Self {
722 source: plan.source.clone(),
723 split_key: plan.split_key.clone(),
724 split_ts: plan.split_ts,
725 children: plan.children.clone().map(|child| child.progress()),
726 phase,
727 }
728 }
729
730 pub fn plan(&self) -> SplitPlan {
732 SplitPlan {
733 source: self.source.clone(),
734 children: self.children.clone().map(|child| child.plan()),
735 split_key: self.split_key.clone(),
736 split_ts: self.split_ts,
737 }
738 }
739}
740
741#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
744#[serde(deny_unknown_fields)]
745struct SplitProgressFile {
746 format_version: u32,
748 checksum: String,
750 progress: SplitProgress,
752}
753
754impl SplitProgressFile {
755 fn envelope(progress: &SplitProgress) -> Result<Self, SplitError> {
756 Ok(Self {
757 format_version: SPLIT_PROGRESS_FORMAT_VERSION,
758 checksum: progress_checksum(progress).map_err(meta_io)?,
759 progress: progress.clone(),
760 })
761 }
762}
763
764fn progress_checksum(progress: &SplitProgress) -> Result<String, ClusterError> {
766 let bytes = serde_json::to_vec(progress).map_err(|error| ClusterError::CorruptMetadata {
767 file: SPLIT_PROGRESS_FILENAME,
768 detail: format!("encode: {error}"),
769 })?;
770 Ok(hex_encode(&Sha256::digest(&bytes)))
771}
772
773fn meta_io(error: ClusterError) -> SplitError {
775 SplitError::Tablet(TabletError::Metadata(error))
776}
777
778fn hex_encode(bytes: &[u8]) -> String {
779 const HEX: &[u8; 16] = b"0123456789abcdef";
780 let mut out = String::with_capacity(bytes.len() * 2);
781 for byte in bytes {
782 out.push(HEX[(byte >> 4) as usize] as char);
783 out.push(HEX[(byte & 0x0f) as usize] as char);
784 }
785 out
786}
787
788pub trait TabletMetaPlane {
797 fn set_tablet(&mut self, descriptor: &TabletDescriptor) -> Result<(), MetaRejectionReason>;
802
803 fn tablet(&self, tablet_id: TabletId) -> Option<TabletDescriptor>;
805
806 fn remove_tablet(
809 &mut self,
810 tablet_id: TabletId,
811 generation: u64,
812 ) -> Result<(), MetaRejectionReason>;
813
814 fn publish_split(&mut self, command: &SplitPublishCommand) -> Result<(), MetaRejectionReason> {
820 command
821 .validate()
822 .map_err(|error| MetaRejectionReason::Invalid {
823 reason: error.to_string(),
824 })?;
825 for child in &command.children {
826 self.set_tablet(child)?;
827 }
828 self.set_tablet(&command.source)
829 }
830}
831
832#[derive(Clone, Default)]
839pub struct InMemoryMetaPlane {
840 tablets: Arc<Mutex<BTreeMap<TabletId, TabletDescriptor>>>,
841}
842
843impl InMemoryMetaPlane {
844 pub fn new() -> Self {
846 Self::default()
847 }
848
849 pub fn descriptors(&self) -> Vec<TabletDescriptor> {
851 self.tablets
852 .lock()
853 .expect("meta plane lock poisoned")
854 .values()
855 .cloned()
856 .collect()
857 }
858}
859
860impl TabletMetaPlane for InMemoryMetaPlane {
861 fn set_tablet(&mut self, descriptor: &TabletDescriptor) -> Result<(), MetaRejectionReason> {
862 descriptor
863 .validate()
864 .map_err(|error| MetaRejectionReason::Invalid {
865 reason: error.to_string(),
866 })?;
867 let mut tablets = self.tablets.lock().expect("meta plane lock poisoned");
868 match tablets.get(&descriptor.tablet_id) {
869 Some(existing) => {
870 if descriptor.generation > existing.generation {
871 tablets.insert(descriptor.tablet_id, descriptor.clone());
872 Ok(())
873 } else if descriptor.generation == existing.generation {
874 if existing == descriptor {
875 Ok(())
876 } else {
877 Err(MetaRejectionReason::Conflict {
878 resource: format!("tablet {}", descriptor.tablet_id),
879 reason: "generation already used for different content".to_owned(),
880 })
881 }
882 } else {
883 Err(MetaRejectionReason::StaleWrite {
884 resource: format!("tablet {}", descriptor.tablet_id),
885 current: MetadataVersion(existing.generation),
886 attempted: MetadataVersion(descriptor.generation),
887 })
888 }
889 }
890 None => {
891 tablets.insert(descriptor.tablet_id, descriptor.clone());
892 Ok(())
893 }
894 }
895 }
896
897 fn tablet(&self, tablet_id: TabletId) -> Option<TabletDescriptor> {
898 self.tablets
899 .lock()
900 .expect("meta plane lock poisoned")
901 .get(&tablet_id)
902 .cloned()
903 }
904
905 fn remove_tablet(
906 &mut self,
907 tablet_id: TabletId,
908 generation: u64,
909 ) -> Result<(), MetaRejectionReason> {
910 let mut tablets = self.tablets.lock().expect("meta plane lock poisoned");
911 match tablets.get(&tablet_id) {
912 None => Ok(()),
913 Some(existing) => {
914 if generation >= existing.generation {
915 tablets.remove(&tablet_id);
916 Ok(())
917 } else {
918 Err(MetaRejectionReason::StaleWrite {
919 resource: format!("tablet {tablet_id}"),
920 current: MetadataVersion(existing.generation),
921 attempted: MetadataVersion(generation),
922 })
923 }
924 }
925 }
926 }
927}
928
929pub trait SnapshotPin: Send {
938 fn pinned_at(&self) -> HlcTimestamp;
940}
941
942pub type RecordStream<'a> = Box<dyn Iterator<Item = (Key, Vec<u8>)> + 'a>;
945
946#[derive(Clone, Debug, PartialEq, Eq)]
948pub enum TabletMutation {
949 Upsert(Key, Vec<u8>),
951 Delete(Key),
953}
954
955pub type MutationStream<'a> = Box<dyn Iterator<Item = TabletMutation> + 'a>;
957
958pub trait TabletKeyspace {
963 fn pin_snapshot(&mut self, ts: HlcTimestamp) -> Result<Box<dyn SnapshotPin>, TabletDataError>;
966
967 fn snapshot_at(&self, ts: HlcTimestamp) -> Result<RecordStream<'_>, TabletDataError>;
970
971 fn deltas_after(&self, ts: HlcTimestamp) -> Result<MutationStream<'_>, TabletDataError>;
974}
975
976pub trait ChildStateSink {
983 fn begin_build(&mut self) -> Result<(), TabletDataError>;
985
986 fn stage(&mut self, key: &Key, value: &[u8]) -> Result<(), TabletDataError>;
988
989 fn install_staged(&mut self) -> Result<(), TabletDataError>;
991
992 fn apply_delta(&mut self, mutation: &TabletMutation) -> Result<(), TabletDataError>;
994}
995
996#[derive(Clone, Default)]
1001pub struct MapKeyspace {
1002 state: Arc<Mutex<MapKeyspaceState>>,
1003}
1004
1005type VersionChain = Vec<(HlcTimestamp, Option<Vec<u8>>)>;
1006
1007#[derive(Default)]
1008struct MapKeyspaceState {
1009 rows: BTreeMap<Key, VersionChain>,
1011 pins: usize,
1013}
1014
1015impl MapKeyspace {
1016 pub fn new() -> Self {
1018 Self::default()
1019 }
1020
1021 pub fn insert(&self, key: Key, ts: HlcTimestamp, value: Vec<u8>) {
1023 let mut state = self.state.lock().expect("keyspace lock poisoned");
1024 let chain = state.rows.entry(key).or_default();
1025 chain.push((ts, Some(value)));
1026 chain.sort_by_key(|(version, _)| *version);
1027 }
1028
1029 pub fn delete(&self, key: Key, ts: HlcTimestamp) {
1031 let mut state = self.state.lock().expect("keyspace lock poisoned");
1032 let chain = state.rows.entry(key).or_default();
1033 chain.push((ts, None));
1034 chain.sort_by_key(|(version, _)| *version);
1035 }
1036
1037 pub fn pin_count(&self) -> usize {
1039 self.state.lock().expect("keyspace lock poisoned").pins
1040 }
1041
1042 pub fn rows_at(&self, ts: HlcTimestamp) -> BTreeMap<Key, Vec<u8>> {
1044 let state = self.state.lock().expect("keyspace lock poisoned");
1045 state
1046 .rows
1047 .iter()
1048 .filter_map(|(key, chain)| {
1049 let visible = chain.iter().rfind(|(version, _)| *version <= ts)?;
1050 visible.1.as_ref().map(|value| (key.clone(), value.clone()))
1051 })
1052 .collect()
1053 }
1054}
1055
1056impl TabletKeyspace for MapKeyspace {
1057 fn pin_snapshot(&mut self, ts: HlcTimestamp) -> Result<Box<dyn SnapshotPin>, TabletDataError> {
1058 self.state.lock().expect("keyspace lock poisoned").pins += 1;
1059 Ok(Box::new(MapSnapshotPin {
1060 ts,
1061 state: self.state.clone(),
1062 }))
1063 }
1064
1065 fn snapshot_at(&self, ts: HlcTimestamp) -> Result<RecordStream<'_>, TabletDataError> {
1066 Ok(Box::new(self.rows_at(ts).into_iter()))
1067 }
1068
1069 fn deltas_after(&self, ts: HlcTimestamp) -> Result<MutationStream<'_>, TabletDataError> {
1070 let before = self.rows_at(ts);
1071 let current = self.rows_at(HlcTimestamp {
1072 physical_micros: u64::MAX,
1073 logical: u32::MAX,
1074 node_tiebreaker: u32::MAX,
1075 });
1076 let mut deltas = Vec::new();
1077 for (key, value) in ¤t {
1078 if before.get(key) != Some(value) {
1079 deltas.push(TabletMutation::Upsert(key.clone(), value.clone()));
1080 }
1081 }
1082 for key in before.keys() {
1083 if !current.contains_key(key) {
1084 deltas.push(TabletMutation::Delete(key.clone()));
1085 }
1086 }
1087 Ok(Box::new(deltas.into_iter()))
1088 }
1089}
1090
1091struct MapSnapshotPin {
1093 ts: HlcTimestamp,
1094 state: Arc<Mutex<MapKeyspaceState>>,
1095}
1096
1097impl SnapshotPin for MapSnapshotPin {
1098 fn pinned_at(&self) -> HlcTimestamp {
1099 self.ts
1100 }
1101}
1102
1103impl Drop for MapSnapshotPin {
1104 fn drop(&mut self) {
1105 self.state.lock().expect("keyspace lock poisoned").pins -= 1;
1106 }
1107}
1108
1109#[derive(Clone, Default)]
1113pub struct MapChildSink {
1114 state: Arc<Mutex<MapChildSinkState>>,
1115}
1116
1117#[derive(Default)]
1118struct MapChildSinkState {
1119 staged: Option<BTreeMap<Key, Vec<u8>>>,
1120 installed: BTreeMap<Key, Vec<u8>>,
1121}
1122
1123impl MapChildSink {
1124 pub fn new() -> Self {
1126 Self::default()
1127 }
1128
1129 pub fn rows(&self) -> BTreeMap<Key, Vec<u8>> {
1131 self.state
1132 .lock()
1133 .expect("child sink lock poisoned")
1134 .installed
1135 .clone()
1136 }
1137}
1138
1139impl ChildStateSink for MapChildSink {
1140 fn begin_build(&mut self) -> Result<(), TabletDataError> {
1141 self.state.lock().expect("child sink lock poisoned").staged = Some(BTreeMap::new());
1142 Ok(())
1143 }
1144
1145 fn stage(&mut self, key: &Key, value: &[u8]) -> Result<(), TabletDataError> {
1146 let mut state = self.state.lock().expect("child sink lock poisoned");
1147 let staged = state
1148 .staged
1149 .as_mut()
1150 .ok_or(TabletDataError::NoStagedBuild)?;
1151 staged.insert(key.clone(), value.to_vec());
1152 Ok(())
1153 }
1154
1155 fn install_staged(&mut self) -> Result<(), TabletDataError> {
1156 let mut state = self.state.lock().expect("child sink lock poisoned");
1157 let staged = state.staged.take().ok_or(TabletDataError::NoStagedBuild)?;
1158 state.installed = staged;
1159 Ok(())
1160 }
1161
1162 fn apply_delta(&mut self, mutation: &TabletMutation) -> Result<(), TabletDataError> {
1163 let mut state = self.state.lock().expect("child sink lock poisoned");
1164 match mutation {
1165 TabletMutation::Upsert(key, value) => {
1166 state.installed.insert(key.clone(), value.clone());
1167 }
1168 TabletMutation::Delete(key) => {
1169 state.installed.remove(key);
1170 }
1171 }
1172 Ok(())
1173 }
1174}
1175
1176#[derive(Debug)]
1189pub struct SourceRetentionGuard {
1190 source: TabletId,
1191 retired_generation: u64,
1192 inner: Mutex<RetentionInner>,
1193}
1194
1195#[derive(Debug, Default)]
1196struct RetentionInner {
1197 next_pin: u64,
1198 pins: BTreeMap<u64, u64>,
1200}
1201
1202impl SourceRetentionGuard {
1203 pub fn new(source: TabletId, retired_generation: u64) -> Self {
1205 Self {
1206 source,
1207 retired_generation,
1208 inner: Mutex::new(RetentionInner::default()),
1209 }
1210 }
1211
1212 pub fn source(&self) -> TabletId {
1214 self.source
1215 }
1216
1217 pub fn retired_generation(&self) -> u64 {
1219 self.retired_generation
1220 }
1221
1222 pub fn pin(&self, used_generation: u64) -> u64 {
1225 let mut inner = self.inner.lock().expect("retention lock poisoned");
1226 let pin = inner.next_pin;
1227 inner.next_pin += 1;
1228 inner.pins.insert(pin, used_generation);
1229 pin
1230 }
1231
1232 pub fn unpin(&self, pin: u64) -> bool {
1234 self.inner
1235 .lock()
1236 .expect("retention lock poisoned")
1237 .pins
1238 .remove(&pin)
1239 .is_some()
1240 }
1241
1242 pub fn pin_count(&self) -> usize {
1244 self.inner
1245 .lock()
1246 .expect("retention lock poisoned")
1247 .pins
1248 .len()
1249 }
1250
1251 pub fn old_generation_pins(&self) -> usize {
1254 self.inner
1255 .lock()
1256 .expect("retention lock poisoned")
1257 .pins
1258 .values()
1259 .filter(|generation| **generation < self.retired_generation)
1260 .count()
1261 }
1262
1263 pub fn ready_for_removal(&self) -> bool {
1266 self.old_generation_pins() == 0
1267 }
1268}
1269
1270#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1279pub enum RetryGuidance {
1280 AwaitSplitPublish {
1283 tablet_id: TabletId,
1285 },
1286 RefreshAndReroute {
1290 tablet_id: TabletId,
1292 },
1293 RefreshAndRetry {
1295 tablet_id: TabletId,
1297 },
1298}
1299
1300impl RetryGuidance {
1301 pub fn category(&self) -> ErrorCategory {
1303 match self {
1304 Self::AwaitSplitPublish { .. } => ErrorCategory::TabletSplitting,
1305 Self::RefreshAndReroute { .. } => ErrorCategory::TabletMoved,
1306 Self::RefreshAndRetry { .. } => ErrorCategory::StaleMetadata,
1307 }
1308 }
1309
1310 pub fn failure(&self) -> crate::routing::Failure {
1314 crate::routing::Failure::new(self.category())
1315 }
1316}
1317
1318pub fn retry_guidance(error: &RoutingError) -> RetryGuidance {
1323 match *error {
1324 RoutingError::TabletSplit { tablet_id, .. } => {
1325 RetryGuidance::AwaitSplitPublish { tablet_id }
1326 }
1327 RoutingError::TabletMoved { tablet_id, .. } => {
1328 RetryGuidance::RefreshAndReroute { tablet_id }
1329 }
1330 RoutingError::StaleMetadata { tablet_id, .. } => {
1331 RetryGuidance::RefreshAndRetry { tablet_id }
1332 }
1333 }
1334}
1335
1336pub struct SplitExecutor<M, K, S> {
1349 progress: SplitProgress,
1350 source_layout: TabletLayout,
1351 meta: M,
1352 keyspace: K,
1353 sinks: [S; 2],
1354 snapshot_pin: Option<Box<dyn SnapshotPin>>,
1355 retention: Option<SourceRetentionGuard>,
1356}
1357
1358impl<M: TabletMetaPlane, K: TabletKeyspace, S: ChildStateSink> SplitExecutor<M, K, S> {
1359 pub fn begin(
1363 plan: SplitPlan,
1364 source_layout: TabletLayout,
1365 meta: M,
1366 keyspace: K,
1367 sinks: [S; 2],
1368 ) -> Result<Self, SplitError> {
1369 plan.validate()?;
1370 if plan.source.state != TabletState::Active {
1371 return Err(SplitError::SourceNotActive {
1372 tablet: plan.source.tablet_id,
1373 state: plan.source.state,
1374 });
1375 }
1376 if source_layout.tablet_id() != plan.source.tablet_id
1377 || source_layout.raft_group_id() != plan.source.raft_group_id
1378 {
1379 return Err(TabletError::TabletMismatch {
1380 path: source_layout.tablet_dir(),
1381 expected: source_layout.tablet_id(),
1382 found: plan.source.tablet_id,
1383 expected_group: source_layout.raft_group_id(),
1384 found_group: plan.source.raft_group_id,
1385 }
1386 .into());
1387 }
1388 source_layout.validate()?;
1390 let executor = Self {
1391 progress: SplitProgress::from_plan(&plan, SplitPhase::Started),
1392 source_layout,
1393 meta,
1394 keyspace,
1395 sinks,
1396 snapshot_pin: None,
1397 retention: None,
1398 };
1399 executor.persist_progress()?;
1400 Ok(executor)
1401 }
1402
1403 pub fn resume(
1407 source_layout: TabletLayout,
1408 meta: M,
1409 keyspace: K,
1410 sinks: [S; 2],
1411 ) -> Result<Option<Self>, SplitError> {
1412 let Some(progress) = load_progress(&source_layout)? else {
1413 return Ok(None);
1414 };
1415 progress.plan().validate()?;
1416 Ok(Some(Self {
1417 progress,
1418 source_layout,
1419 meta,
1420 keyspace,
1421 sinks,
1422 snapshot_pin: None,
1423 retention: None,
1424 }))
1425 }
1426
1427 pub fn phase(&self) -> SplitPhase {
1429 self.progress.phase
1430 }
1431
1432 pub fn progress(&self) -> &SplitProgress {
1434 &self.progress
1435 }
1436
1437 pub fn plan(&self) -> SplitPlan {
1439 self.progress.plan()
1440 }
1441
1442 pub fn retention(&self) -> Option<&SourceRetentionGuard> {
1444 self.retention.as_ref()
1445 }
1446
1447 pub fn meta(&self) -> &M {
1449 &self.meta
1450 }
1451
1452 pub fn keyspace(&self) -> &K {
1454 &self.keyspace
1455 }
1456
1457 pub fn sinks(&self) -> &[S; 2] {
1459 &self.sinks
1460 }
1461
1462 pub fn step(&mut self) -> Result<SplitPhase, SplitError> {
1466 use SplitPhase::{
1467 CaughtUp, ChildrenBuilt, ChildrenCreated, MarkedSplitting, Published, SnapshotPinned,
1468 SourceRetired, Started,
1469 };
1470 let next = match self.progress.phase {
1471 Started => {
1472 self.mark_source_splitting()?;
1473 MarkedSplitting
1474 }
1475 MarkedSplitting => {
1476 self.create_children()?;
1477 ChildrenCreated
1478 }
1479 ChildrenCreated => {
1480 self.pin_source_snapshot()?;
1481 SnapshotPinned
1482 }
1483 SnapshotPinned => {
1484 self.build_children()?;
1485 ChildrenBuilt
1486 }
1487 ChildrenBuilt => {
1488 self.catch_up_children()?;
1489 CaughtUp
1490 }
1491 CaughtUp => {
1492 self.publish_children()?;
1493 Published
1494 }
1495 Published => {
1496 self.remove_source()?;
1497 SourceRetired
1498 }
1499 SourceRetired => SourceRetired,
1500 };
1501 self.progress.phase = next;
1504 if next != SourceRetired {
1505 self.persist_progress()?;
1506 }
1507 if let Some(hook) = next.hook_name() {
1508 mongreldb_fault::inject(hook)?;
1509 }
1510 Ok(next)
1511 }
1512
1513 pub fn run(&mut self) -> Result<(), SplitError> {
1517 while self.progress.phase != SplitPhase::SourceRetired {
1518 self.step()?;
1519 }
1520 Ok(())
1521 }
1522
1523 pub fn run_until(&mut self, phase: SplitPhase) -> Result<(), SplitError> {
1525 while self.progress.phase != phase {
1526 self.step()?;
1527 }
1528 Ok(())
1529 }
1530
1531 fn mark_source_splitting(&mut self) -> Result<(), SplitError> {
1534 let marked = self
1535 .progress
1536 .source
1537 .published_transition(TabletState::Splitting)?;
1538 self.meta.set_tablet(&marked)?;
1539 self.source_layout.store_metadata(&marked)?;
1540 Ok(())
1541 }
1542
1543 fn create_children(&mut self) -> Result<(), SplitError> {
1546 let plan = self.plan();
1547 for (descriptor, child) in plan.child_descriptors().iter().zip(plan.children.iter()) {
1548 child.layout.create(descriptor)?;
1549 self.meta.set_tablet(descriptor)?;
1550 }
1551 Ok(())
1552 }
1553
1554 fn pin_source_snapshot(&mut self) -> Result<(), SplitError> {
1557 self.ensure_snapshot_pin()
1558 }
1559
1560 fn ensure_snapshot_pin(&mut self) -> Result<(), SplitError> {
1563 if self.snapshot_pin.is_none() {
1564 self.snapshot_pin = Some(self.keyspace.pin_snapshot(self.progress.split_ts)?);
1565 }
1566 Ok(())
1567 }
1568
1569 fn build_children(&mut self) -> Result<(), SplitError> {
1572 self.ensure_snapshot_pin()?;
1573 for sink in &mut self.sinks {
1574 sink.begin_build()?;
1575 }
1576 let plan = self.plan();
1577 let snapshot = self.keyspace.snapshot_at(self.progress.split_ts)?;
1578 for (key, value) in snapshot {
1579 let index = route_child(&plan, &key)?;
1580 self.sinks[index].stage(&key, &value)?;
1581 }
1582 for sink in &mut self.sinks {
1583 sink.install_staged()?;
1584 }
1585 Ok(())
1586 }
1587
1588 fn catch_up_children(&mut self) -> Result<(), SplitError> {
1591 self.ensure_snapshot_pin()?;
1592 let plan = self.plan();
1593 let deltas = self.keyspace.deltas_after(self.progress.split_ts)?;
1594 for mutation in deltas {
1595 let key = match &mutation {
1596 TabletMutation::Upsert(key, _) | TabletMutation::Delete(key) => key,
1597 };
1598 let index = route_child(&plan, key)?;
1599 self.sinks[index].apply_delta(&mutation)?;
1600 }
1601 Ok(())
1602 }
1603
1604 fn publish_children(&mut self) -> Result<(), SplitError> {
1609 let plan = self.plan();
1610 let command = SplitPublishCommand::from_plan(&plan)?;
1611 mongreldb_fault::inject("tablet.split.before")?;
1612 self.meta.publish_split(&command)?;
1613 mongreldb_fault::inject("tablet.split.after")?;
1614 for (descriptor, child) in command.children.iter().zip(plan.children.iter()) {
1616 child.layout.store_metadata(descriptor)?;
1617 }
1618 self.source_layout.store_metadata(&command.source)?;
1619 self.retention = Some(SourceRetentionGuard::new(
1620 plan.source.tablet_id,
1621 command.publish_generation(),
1622 ));
1623 self.snapshot_pin = None;
1625 Ok(())
1626 }
1627
1628 fn remove_source(&mut self) -> Result<(), SplitError> {
1632 let source_id = self.progress.source.tablet_id;
1633 if let Some(guard) = &self.retention {
1634 if !guard.ready_for_removal() {
1635 return Err(SplitError::SourceRetained {
1636 tablet: source_id,
1637 pins: guard.old_generation_pins(),
1638 });
1639 }
1640 }
1641 if let Some(current) = self.meta.tablet(source_id) {
1642 let retired = if current.state == TabletState::Retired {
1643 current
1644 } else {
1645 let retired = current.published_transition(TabletState::Retired)?;
1646 self.meta.set_tablet(&retired)?;
1647 retired
1648 };
1649 self.meta.remove_tablet(source_id, retired.generation)?;
1650 }
1651 self.source_layout.teardown()?;
1655 Ok(())
1656 }
1657
1658 fn persist_progress(&self) -> Result<(), SplitError> {
1661 let file = SplitProgressFile::envelope(&self.progress)?;
1662 let bytes = crate::node::encode_json(SPLIT_PROGRESS_FILENAME, &file).map_err(meta_io)?;
1663 crate::node::write_meta_atomic(
1664 &self.source_layout.tablet_dir(),
1665 SPLIT_PROGRESS_FILENAME,
1666 &bytes,
1667 )
1668 .map_err(ClusterError::Io)
1669 .map_err(meta_io)?;
1670 Ok(())
1671 }
1672}
1673
1674fn route_child(plan: &SplitPlan, key: &Key) -> Result<usize, SplitError> {
1677 if plan.children[0].bounds.contains(key) {
1678 return Ok(0);
1679 }
1680 if plan.children[1].bounds.contains(key) {
1681 return Ok(1);
1682 }
1683 Err(SplitError::KeyOutsideSource(key.clone()))
1684}
1685
1686#[derive(Clone, Debug, PartialEq, Eq)]
1692pub struct SplitAbortReport {
1693 pub source: TabletId,
1695 pub phase: Option<SplitPhase>,
1698 pub children_removed: Vec<TabletId>,
1700 pub source_after: Option<TabletDescriptor>,
1704}
1705
1706pub fn abort_split<M: TabletMetaPlane>(
1723 source_layout: &TabletLayout,
1724 meta: &mut M,
1725) -> Result<SplitAbortReport, SplitError> {
1726 let Some(progress) = load_progress(source_layout)? else {
1727 return Ok(SplitAbortReport {
1728 source: source_layout.tablet_id(),
1729 phase: None,
1730 children_removed: Vec::new(),
1731 source_after: None,
1732 });
1733 };
1734 if progress.phase >= SplitPhase::Published {
1735 return Err(SplitError::CannotAbort {
1736 tablet: progress.source.tablet_id,
1737 phase: progress.phase,
1738 });
1739 }
1740 let mut children_removed = Vec::new();
1745 for child in &progress.children {
1746 if let Some(current) = meta.tablet(child.tablet_id) {
1747 meta.remove_tablet(child.tablet_id, current.generation)?;
1748 children_removed.push(child.tablet_id);
1749 }
1750 }
1751 let current = meta.tablet(progress.source.tablet_id).ok_or_else(|| {
1757 SplitError::InvalidPlan(format!(
1758 "source tablet {} is missing from the meta plane mid-abort",
1759 progress.source.tablet_id
1760 ))
1761 })?;
1762 let restored = if current.state == TabletState::Active {
1763 current
1764 } else {
1765 let restored = current.published_transition(TabletState::Active)?;
1766 meta.set_tablet(&restored)?;
1767 restored
1768 };
1769 source_layout.store_metadata(&restored)?;
1771 for child in &progress.children {
1774 child.plan().layout.teardown()?;
1775 }
1776 let record = source_layout.tablet_dir().join(SPLIT_PROGRESS_FILENAME);
1779 match std::fs::remove_file(&record) {
1780 Ok(()) => {}
1781 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
1782 Err(error) => {
1783 return Err(meta_io(ClusterError::Io(error)));
1784 }
1785 }
1786 Ok(SplitAbortReport {
1787 source: progress.source.tablet_id,
1788 phase: Some(progress.phase),
1789 children_removed,
1790 source_after: Some(restored),
1791 })
1792}
1793
1794fn load_progress(source_layout: &TabletLayout) -> Result<Option<SplitProgress>, SplitError> {
1797 let path = source_layout.tablet_dir().join(SPLIT_PROGRESS_FILENAME);
1798 let Some(bytes) = crate::node::read_meta_file(&path).map_err(meta_io)? else {
1799 return Ok(None);
1800 };
1801 let file: SplitProgressFile =
1802 crate::node::decode_json(SPLIT_PROGRESS_FILENAME, &bytes).map_err(meta_io)?;
1803 if file.format_version < MIN_SUPPORTED_SPLIT_PROGRESS_FORMAT_VERSION
1804 || file.format_version > SPLIT_PROGRESS_FORMAT_VERSION
1805 {
1806 return Err(meta_io(ClusterError::UnsupportedFormatVersion {
1807 file: SPLIT_PROGRESS_FILENAME,
1808 found: file.format_version,
1809 min: MIN_SUPPORTED_SPLIT_PROGRESS_FORMAT_VERSION,
1810 max: SPLIT_PROGRESS_FORMAT_VERSION,
1811 }));
1812 }
1813 if file.checksum != progress_checksum(&file.progress).map_err(meta_io)? {
1814 return Err(meta_io(ClusterError::CorruptMetadata {
1815 file: SPLIT_PROGRESS_FILENAME,
1816 detail: "checksum mismatch".to_owned(),
1817 }));
1818 }
1819 let progress = file.progress;
1820 if progress.source.tablet_id != source_layout.tablet_id()
1821 || progress.source.raft_group_id != source_layout.raft_group_id()
1822 {
1823 return Err(TabletError::TabletMismatch {
1824 path: source_layout.tablet_dir(),
1825 expected: source_layout.tablet_id(),
1826 found: progress.source.tablet_id,
1827 expected_group: source_layout.raft_group_id(),
1828 found_group: progress.source.raft_group_id,
1829 }
1830 .into());
1831 }
1832 Ok(Some(progress))
1833}
1834
1835pub fn split_progress(source_layout: &TabletLayout) -> Result<Option<SplitProgress>, SplitError> {
1839 load_progress(source_layout)
1840}
1841
1842#[cfg(test)]
1847pub(crate) static EXECUTOR_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
1848
1849#[cfg(test)]
1850mod tests {
1851 use std::time::Duration;
1852
1853 use mongreldb_types::ids::NodeId;
1854
1855 use super::*;
1856 use crate::routing::{
1857 GroupKey, OperationDescriptor, RetryAction, RetryPolicy, RetryState, RoutingCache,
1858 };
1859 use crate::tablet::{
1860 check_generation, find_tablet_for_key, tablets_overlapping, KeyValue, RowKeyEncoder,
1861 };
1862
1863 fn node(byte: u8) -> NodeId {
1864 NodeId::from_bytes([byte; 16])
1865 }
1866
1867 fn tablet_id(byte: u8) -> TabletId {
1868 TabletId::from_bytes([byte; 16])
1869 }
1870
1871 fn group_id(byte: u8) -> RaftGroupId {
1872 RaftGroupId::from_bytes([byte; 16])
1873 }
1874
1875 fn key(bytes: &[u8]) -> Key {
1876 Key::from_bytes(bytes.to_vec())
1877 }
1878
1879 fn text_key(text: &str) -> Key {
1880 RowKeyEncoder::encode_key(&[KeyValue::Text(text.to_owned())])
1881 }
1882
1883 fn ts(micros: u64) -> HlcTimestamp {
1884 HlcTimestamp {
1885 physical_micros: micros,
1886 logical: 0,
1887 node_tiebreaker: 0,
1888 }
1889 }
1890
1891 fn source_descriptor() -> TabletDescriptor {
1892 TabletDescriptor {
1893 tablet_id: tablet_id(1),
1894 table_id: TableId::new(3),
1895 database_id: mongreldb_types::ids::DatabaseId::ZERO,
1896 raft_group_id: group_id(1),
1897 partition: PartitionBounds::new(
1898 Bound::Included(text_key("a")),
1899 Bound::Excluded(text_key("z")),
1900 )
1901 .unwrap(),
1902 replicas: vec![
1903 ReplicaDescriptor {
1904 node_id: node(1),
1905 role: ReplicaRole::Voter,
1906 raft_node_id: 11,
1907 },
1908 ReplicaDescriptor {
1909 node_id: node(2),
1910 role: ReplicaRole::Voter,
1911 raft_node_id: 12,
1912 },
1913 ],
1914 leader_hint: Some(node(1)),
1915 generation: 5,
1916 state: TabletState::Active,
1917 }
1918 }
1919
1920 fn allocation(tablet: u8, group: u8, raft_base: u64) -> ChildAllocation {
1921 ChildAllocation {
1922 tablet_id: tablet_id(tablet),
1923 raft_group_id: group_id(group),
1924 replicas: vec![
1925 ReplicaDescriptor {
1926 node_id: node(3),
1927 role: ReplicaRole::Voter,
1928 raft_node_id: raft_base,
1929 },
1930 ReplicaDescriptor {
1931 node_id: node(4),
1932 role: ReplicaRole::Voter,
1933 raft_node_id: raft_base + 1,
1934 },
1935 ],
1936 }
1937 }
1938
1939 const LOWER_KEYS: [&str; 6] = ["b", "d", "f", "h", "j", "l"];
1941 const UPPER_KEYS: [&str; 7] = ["m", "o", "q", "s", "u", "w", "y"];
1942
1943 fn seed_keyspace(keyspace: &MapKeyspace) {
1944 for name in LOWER_KEYS.into_iter().chain(UPPER_KEYS) {
1945 keyspace.insert(
1946 text_key(name),
1947 ts(100),
1948 format!("v-{name}@100").into_bytes(),
1949 );
1950 }
1951 keyspace.insert(text_key("b"), ts(200), b"v-b@200".to_vec());
1953 keyspace.insert(text_key("y"), ts(200), b"v-y@200".to_vec());
1954 keyspace.insert(text_key("n"), ts(200), b"v-n@200".to_vec());
1955 }
1956
1957 struct SplitFixture {
1958 _dir: tempfile::TempDir,
1959 source: TabletDescriptor,
1960 source_layout: TabletLayout,
1961 meta: InMemoryMetaPlane,
1962 keyspace: MapKeyspace,
1963 sinks: [MapChildSink; 2],
1964 plan: SplitPlan,
1965 }
1966
1967 fn split_fixture() -> SplitFixture {
1968 let dir = tempfile::tempdir().unwrap();
1969 let source = source_descriptor();
1970 let source_layout = TabletLayout::new(dir.path(), source.tablet_id, source.raft_group_id);
1971 source_layout.create(&source).unwrap();
1972 let mut meta = InMemoryMetaPlane::new();
1973 meta.set_tablet(&source).unwrap();
1974 let keyspace = MapKeyspace::new();
1975 seed_keyspace(&keyspace);
1976 let planner = TabletSplitPlanner::new(dir.path());
1977 let plan = planner
1978 .plan(
1979 &source,
1980 SplitKeySelection::Explicit(text_key("m")),
1981 ts(150),
1982 [allocation(2, 2, 21), allocation(3, 3, 31)],
1983 )
1984 .unwrap();
1985 SplitFixture {
1986 _dir: dir,
1987 source,
1988 source_layout,
1989 meta,
1990 keyspace,
1991 sinks: [MapChildSink::new(), MapChildSink::new()],
1992 plan,
1993 }
1994 }
1995
1996 type TestExecutor = SplitExecutor<InMemoryMetaPlane, MapKeyspace, MapChildSink>;
1997
1998 fn begin_executor(fixture: &SplitFixture) -> TestExecutor {
1999 SplitExecutor::begin(
2000 fixture.plan.clone(),
2001 fixture.source_layout.clone(),
2002 fixture.meta.clone(),
2003 fixture.keyspace.clone(),
2004 fixture.sinks.clone(),
2005 )
2006 .unwrap()
2007 }
2008
2009 fn assert_split_completed(fixture: &SplitFixture) {
2010 assert!(fixture.meta.tablet(fixture.source.tablet_id).is_none());
2013 for (index, id) in [tablet_id(2), tablet_id(3)].into_iter().enumerate() {
2014 let child = fixture.meta.tablet(id).unwrap();
2015 assert_eq!(child.state, TabletState::Active);
2016 assert_eq!(child.generation, 7);
2017 assert!(child
2018 .replicas
2019 .iter()
2020 .all(|replica| replica.role == ReplicaRole::Voter));
2021 assert_eq!(child.partition, fixture.plan.children[index].bounds);
2022 assert_eq!(
2024 fixture.plan.children[index].layout.load_metadata().unwrap(),
2025 child
2026 );
2027 }
2028 assert!(!fixture.source_layout.tablet_dir().exists());
2030 assert!(!fixture.source_layout.group_dir().exists());
2031 assert!(!fixture
2032 .source_layout
2033 .tablet_dir()
2034 .join(SPLIT_PROGRESS_FILENAME)
2035 .exists());
2036 let lower_rows = fixture.sinks[0].rows();
2038 let upper_rows = fixture.sinks[1].rows();
2039 assert!(lower_rows.keys().all(|key| *key < text_key("m")));
2040 assert!(upper_rows.keys().all(|key| *key >= text_key("m")));
2041 let mut union = lower_rows.clone();
2042 for (key, value) in &upper_rows {
2043 assert!(
2044 union.insert(key.clone(), value.clone()).is_none(),
2045 "duplicate key {key} across the split boundary"
2046 );
2047 }
2048 assert_eq!(union, fixture.keyspace.rows_at(ts(u64::MAX)));
2049 }
2050
2051 #[test]
2054 fn midpoint_key_is_deterministic_and_strictly_between() {
2055 let cases: &[(&[u8], &[u8])] = &[
2056 (b"a", b"z"),
2057 (b"aa", b"ab"),
2058 (b"a", b"a\x01"),
2059 (b"m", b"n"),
2060 (b"\x00", b"\xff"),
2061 (b"abc", b"abd"),
2062 (b"a", b"aa"),
2063 (b"", b"\x01"),
2064 ];
2065 for (low, high) in cases {
2066 let mid = midpoint_key(&key(low), &key(high)).unwrap();
2067 assert!(key(low) < mid, "midpoint {mid} not above {low:?}");
2068 assert!(mid < key(high), "midpoint {mid} not below {high:?}");
2069 assert_eq!(mid, midpoint_key(&key(low), &key(high)).unwrap());
2071 }
2072 assert_eq!(midpoint_key(&key(b"a"), &key(b"a\x00")), None);
2074 assert_eq!(midpoint_key(&key(b"a"), &key(b"z")).unwrap(), key(b"m"));
2076 assert_eq!(
2077 midpoint_key(&key(b"\x10"), &key(b"\x20")).unwrap(),
2078 key(b"\x18")
2079 );
2080 }
2081
2082 #[test]
2083 fn planner_chooses_and_validates_split_keys() {
2084 let dir = tempfile::tempdir().unwrap();
2085 let source = source_descriptor();
2086 let planner = TabletSplitPlanner::new(dir.path());
2087
2088 let plan = planner
2090 .plan(
2091 &source,
2092 SplitKeySelection::Midpoint,
2093 ts(150),
2094 [allocation(2, 2, 21), allocation(3, 3, 31)],
2095 )
2096 .unwrap();
2097 let (expected_lower, expected_upper) = source.partition.split_at(&plan.split_key).unwrap();
2098 assert_eq!(plan.children[0].bounds, expected_lower);
2099 assert_eq!(plan.children[1].bounds, expected_upper);
2100 assert!(plan.children[0]
2101 .bounds
2102 .meets_start_of(&plan.children[1].bounds));
2103
2104 let mut unbounded = source.clone();
2106 unbounded.partition = PartitionBounds::unbounded();
2107 assert!(matches!(
2108 planner.plan(
2109 &unbounded,
2110 SplitKeySelection::Midpoint,
2111 ts(150),
2112 [allocation(2, 2, 21), allocation(3, 3, 31)],
2113 ),
2114 Err(SplitError::UnboundedMidpoint)
2115 ));
2116
2117 for bad in ["a", "z", "0"] {
2119 assert!(matches!(
2120 planner.plan(
2121 &source,
2122 SplitKeySelection::Explicit(text_key(bad)),
2123 ts(150),
2124 [allocation(2, 2, 21), allocation(3, 3, 31)],
2125 ),
2126 Err(SplitError::InvalidSplitKey { .. })
2127 ));
2128 }
2129
2130 let mut splitting = source.clone();
2132 splitting.state = TabletState::Splitting;
2133 assert!(matches!(
2134 planner.plan(
2135 &splitting,
2136 SplitKeySelection::Explicit(text_key("m")),
2137 ts(150),
2138 [allocation(2, 2, 21), allocation(3, 3, 31)],
2139 ),
2140 Err(SplitError::SourceNotActive {
2141 state: TabletState::Splitting,
2142 ..
2143 })
2144 ));
2145
2146 let mut colliding = source.clone();
2148 let error = planner
2149 .plan(
2150 &colliding,
2151 SplitKeySelection::Explicit(text_key("m")),
2152 ts(150),
2153 [allocation(2, 2, 21), allocation(2, 3, 31)],
2154 )
2155 .unwrap_err();
2156 assert!(matches!(error, SplitError::InvalidPlan(_)));
2157 colliding = source.clone();
2158 let error = planner
2159 .plan(
2160 &colliding,
2161 SplitKeySelection::Explicit(text_key("m")),
2162 ts(150),
2163 [allocation(1, 2, 21), allocation(3, 3, 31)],
2164 )
2165 .unwrap_err();
2166 assert!(matches!(error, SplitError::InvalidPlan(_)));
2167 }
2168
2169 #[test]
2170 fn child_descriptors_are_creating_learners_at_the_inception_generation() {
2171 let dir = tempfile::tempdir().unwrap();
2172 let planner = TabletSplitPlanner::new(dir.path());
2173 let plan = planner
2174 .plan(
2175 &source_descriptor(),
2176 SplitKeySelection::Explicit(text_key("m")),
2177 ts(150),
2178 [allocation(2, 2, 21), allocation(3, 3, 31)],
2179 )
2180 .unwrap();
2181 let children = plan.child_descriptors();
2182 assert_eq!(children[0].state, TabletState::Creating);
2183 assert_eq!(children[0].generation, 6); assert!(children
2185 .iter()
2186 .flat_map(|child| child.replicas.iter())
2187 .all(|replica| replica.role == ReplicaRole::Learner));
2188 for child in &children {
2189 child.validate().unwrap();
2190 }
2191 assert!(!children.iter().any(|child| child.state.is_routable()));
2193 }
2194
2195 #[test]
2198 fn publish_command_flips_three_descriptors_at_one_generation() {
2199 let dir = tempfile::tempdir().unwrap();
2200 let planner = TabletSplitPlanner::new(dir.path());
2201 let plan = planner
2202 .plan(
2203 &source_descriptor(),
2204 SplitKeySelection::Explicit(text_key("m")),
2205 ts(150),
2206 [allocation(2, 2, 21), allocation(3, 3, 31)],
2207 )
2208 .unwrap();
2209 let command = SplitPublishCommand::from_plan(&plan).unwrap();
2210 assert_eq!(command.publish_generation(), 7); assert_eq!(command.source.state, TabletState::Retiring);
2212 assert_eq!(command.source.generation, 7);
2213 for child in &command.children {
2214 assert_eq!(child.state, TabletState::Active);
2215 assert_eq!(child.generation, 7);
2216 assert!(child
2217 .replicas
2218 .iter()
2219 .all(|replica| replica.role == ReplicaRole::Voter));
2220 }
2221 let bytes = serde_json::to_vec(&command).unwrap();
2223 let back: SplitPublishCommand = serde_json::from_slice(&bytes).unwrap();
2224 assert_eq!(back, command);
2225
2226 let mut wrong_state = command.clone();
2228 wrong_state.children[0].state = TabletState::Creating;
2229 assert!(wrong_state.validate().is_err());
2230 let mut skewed = command.clone();
2232 skewed.children[0].generation = 8;
2233 assert!(skewed.validate().is_err());
2234 let mut wrong_bounds = command.clone();
2236 wrong_bounds.children[0].partition = PartitionBounds::unbounded();
2237 assert!(wrong_bounds.validate().is_err());
2238 let mut duplicate = command.clone();
2240 duplicate.children[1].tablet_id = duplicate.children[0].tablet_id;
2241 assert!(duplicate.validate().is_err());
2242 }
2243
2244 #[test]
2247 fn full_split_partitions_the_keyspace_and_flips_routing_atomically() {
2248 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2249 let fixture = split_fixture();
2250 let table = fixture.source.table_id;
2251 let mut executor = begin_executor(&fixture);
2252 assert_eq!(executor.phase(), SplitPhase::Started);
2253
2254 assert_eq!(executor.step().unwrap(), SplitPhase::MarkedSplitting);
2256 let marked = fixture.meta.tablet(fixture.source.tablet_id).unwrap();
2257 assert_eq!(marked.state, TabletState::Splitting);
2258 assert_eq!(marked.generation, 6);
2259 let error = check_generation(&marked, 5).unwrap_err();
2260 assert!(matches!(error, RoutingError::TabletSplit { .. }));
2261 assert!(matches!(
2262 retry_guidance(&error),
2263 RetryGuidance::AwaitSplitPublish { tablet_id } if tablet_id == fixture.source.tablet_id
2264 ));
2265 let tablets = fixture.meta.descriptors();
2266 for name in ["b", "l", "m", "y"] {
2267 assert_eq!(
2268 find_tablet_for_key(&tablets, table, &text_key(name))
2269 .unwrap()
2270 .tablet_id,
2271 fixture.source.tablet_id,
2272 "key {name} left the source during the split"
2273 );
2274 }
2275
2276 assert_eq!(executor.step().unwrap(), SplitPhase::ChildrenCreated);
2278 for id in [tablet_id(2), tablet_id(3)] {
2279 let child = fixture.meta.tablet(id).unwrap();
2280 assert_eq!(child.state, TabletState::Creating);
2281 assert_eq!(child.generation, 6);
2282 assert!(child
2283 .replicas
2284 .iter()
2285 .all(|replica| replica.role == ReplicaRole::Learner));
2286 }
2287 let tablets = fixture.meta.descriptors();
2288 for name in ["b", "y"] {
2289 assert_eq!(
2290 find_tablet_for_key(&tablets, table, &text_key(name))
2291 .unwrap()
2292 .tablet_id,
2293 fixture.source.tablet_id,
2294 "Creating child exposed key {name} before catch-up"
2295 );
2296 }
2297 for child in &fixture.plan.children {
2298 assert_eq!(
2299 child.layout.load_metadata().unwrap().state,
2300 TabletState::Creating
2301 );
2302 }
2303
2304 assert_eq!(executor.step().unwrap(), SplitPhase::SnapshotPinned);
2306 assert_eq!(fixture.keyspace.pin_count(), 1);
2307 assert_eq!(executor.phase(), SplitPhase::SnapshotPinned);
2308
2309 assert_eq!(executor.step().unwrap(), SplitPhase::ChildrenBuilt);
2312 assert_eq!(fixture.sinks[0].rows().len(), LOWER_KEYS.len());
2313 assert_eq!(fixture.sinks[1].rows().len(), UPPER_KEYS.len());
2314 assert_eq!(
2315 fixture.sinks[1].rows().get(&text_key("y")),
2316 Some(&b"v-y@100".to_vec()),
2317 "post-split write leaked into the pinned snapshot"
2318 );
2319
2320 assert_eq!(executor.step().unwrap(), SplitPhase::CaughtUp);
2322 assert_eq!(
2323 fixture.sinks[0].rows().get(&text_key("b")),
2324 Some(&b"v-b@200".to_vec())
2325 );
2326 assert_eq!(
2327 fixture.sinks[1].rows().get(&text_key("y")),
2328 Some(&b"v-y@200".to_vec())
2329 );
2330 assert_eq!(
2331 fixture.sinks[1].rows().get(&text_key("n")),
2332 Some(&b"v-n@200".to_vec())
2333 );
2334
2335 assert_eq!(executor.step().unwrap(), SplitPhase::Published);
2337 assert_eq!(fixture.keyspace.pin_count(), 0);
2338 let retiring = fixture.meta.tablet(fixture.source.tablet_id).unwrap();
2339 assert_eq!(retiring.state, TabletState::Retiring);
2340 assert_eq!(retiring.generation, 7);
2341 let tablets = fixture.meta.descriptors();
2342 assert_eq!(
2343 find_tablet_for_key(&tablets, table, &text_key("b"))
2344 .unwrap()
2345 .tablet_id,
2346 tablet_id(2)
2347 );
2348 assert_eq!(
2349 find_tablet_for_key(&tablets, table, &text_key("y"))
2350 .unwrap()
2351 .tablet_id,
2352 tablet_id(3)
2353 );
2354 assert_eq!(
2356 find_tablet_for_key(&tablets, table, &text_key("m"))
2357 .unwrap()
2358 .tablet_id,
2359 tablet_id(3)
2360 );
2361 let overlapping = tablets_overlapping(&tablets, table, &PartitionBounds::unbounded());
2363 assert_eq!(
2364 overlapping
2365 .iter()
2366 .map(|tablet| tablet.tablet_id)
2367 .collect::<Vec<_>>(),
2368 vec![tablet_id(2), tablet_id(3)]
2369 );
2370 let error = check_generation(&retiring, 5).unwrap_err();
2372 assert!(matches!(error, RoutingError::TabletMoved { .. }));
2373 assert!(matches!(
2374 retry_guidance(&error),
2375 RetryGuidance::RefreshAndReroute { .. }
2376 ));
2377 for id in [tablet_id(2), tablet_id(3)] {
2379 assert!(check_generation(&fixture.meta.tablet(id).unwrap(), 7).is_ok());
2380 }
2381
2382 let pin = executor.retention().unwrap().pin(5);
2384 assert!(matches!(
2385 executor.step(),
2386 Err(SplitError::SourceRetained { pins: 1, .. })
2387 ));
2388 assert_eq!(executor.phase(), SplitPhase::Published);
2389 let fresh = executor.retention().unwrap().pin(7);
2391 assert!(executor.retention().unwrap().unpin(pin));
2392 assert_eq!(executor.step().unwrap(), SplitPhase::SourceRetired);
2393 assert!(executor.retention().unwrap().unpin(fresh));
2394 assert_split_completed(&fixture);
2395 executor.run().unwrap();
2397 assert!(SplitExecutor::resume(
2398 fixture.source_layout.clone(),
2399 fixture.meta.clone(),
2400 fixture.keyspace.clone(),
2401 fixture.sinks.clone(),
2402 )
2403 .unwrap()
2404 .is_none());
2405 }
2406
2407 #[test]
2410 fn split_resumes_after_a_crash_at_every_durable_boundary() {
2411 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2412 let hooks = [
2413 "tablet.split.phase.1",
2414 "tablet.split.phase.2",
2415 "tablet.split.phase.3",
2416 "tablet.split.phase.4",
2417 "tablet.split.phase.5",
2418 "tablet.split.phase.6",
2419 "tablet.split.phase.7",
2420 "tablet.split.before",
2421 "tablet.split.after",
2422 ];
2423 for hook in hooks {
2424 let fixture = split_fixture();
2425 let mut executor = begin_executor(&fixture);
2426 {
2427 let _guard =
2428 mongreldb_fault::ScopedGuard::limited(hook, mongreldb_fault::Action::Fail, 1);
2429 assert!(
2430 matches!(executor.run(), Err(SplitError::Fault(_))),
2431 "hook {hook} did not fire"
2432 );
2433 }
2434 drop(executor);
2436 let resumed = SplitExecutor::resume(
2437 fixture.source_layout.clone(),
2438 fixture.meta.clone(),
2439 fixture.keyspace.clone(),
2440 fixture.sinks.clone(),
2441 )
2442 .unwrap();
2443 if hook == "tablet.split.phase.7" {
2444 assert!(resumed.is_none(), "hook {hook}");
2447 } else {
2448 resumed
2449 .unwrap()
2450 .run()
2451 .unwrap_or_else(|error| panic!("resume after {hook} failed: {error}"));
2452 }
2453 assert_split_completed(&fixture);
2454 }
2455 }
2456
2457 #[test]
2458 fn resume_replays_an_interrupted_step_idempotently() {
2459 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2460 let fixture = split_fixture();
2461 let mut executor = begin_executor(&fixture);
2462 executor.run_until(SplitPhase::ChildrenCreated).unwrap();
2463 drop(executor);
2465 let mut resumed = SplitExecutor::resume(
2466 fixture.source_layout.clone(),
2467 fixture.meta.clone(),
2468 fixture.keyspace.clone(),
2469 fixture.sinks.clone(),
2470 )
2471 .unwrap()
2472 .unwrap();
2473 assert_eq!(resumed.phase(), SplitPhase::ChildrenCreated);
2474 resumed.step().unwrap();
2475 assert_eq!(resumed.phase(), SplitPhase::SnapshotPinned);
2476 assert_eq!(fixture.keyspace.pin_count(), 1);
2477 resumed.run().unwrap();
2478 assert_split_completed(&fixture);
2479 }
2480
2481 #[test]
2484 fn retention_guard_tracks_old_generation_pins() {
2485 let guard = SourceRetentionGuard::new(tablet_id(1), 7);
2486 assert!(guard.ready_for_removal());
2487 assert_eq!(guard.source(), tablet_id(1));
2488 assert_eq!(guard.retired_generation(), 7);
2489
2490 let stale = guard.pin(5);
2491 let boundary = guard.pin(7); assert_eq!(guard.pin_count(), 2);
2493 assert_eq!(guard.old_generation_pins(), 1);
2494 assert!(!guard.ready_for_removal());
2495
2496 assert!(guard.unpin(stale));
2497 assert!(guard.ready_for_removal());
2498 assert!(guard.unpin(boundary));
2499 assert!(!guard.unpin(boundary), "double release must not succeed");
2500 assert_eq!(guard.pin_count(), 0);
2501 }
2502
2503 #[test]
2506 fn stale_requests_classify_and_guidance_feeds_the_retry_policy() {
2507 let mut descriptor = source_descriptor();
2508 assert!(check_generation(&descriptor, 5).is_ok());
2510
2511 descriptor.state = TabletState::Splitting;
2513 descriptor.generation = 6;
2514 let error = check_generation(&descriptor, 5).unwrap_err();
2515 assert_eq!(
2516 error,
2517 RoutingError::TabletSplit {
2518 tablet_id: tablet_id(1),
2519 used_generation: 5,
2520 current_generation: 6,
2521 }
2522 );
2523 let guidance = retry_guidance(&error);
2524 assert_eq!(guidance.category(), ErrorCategory::TabletSplitting);
2525
2526 descriptor.state = TabletState::Retiring;
2528 descriptor.generation = 7;
2529 let error = check_generation(&descriptor, 5).unwrap_err();
2530 assert!(matches!(error, RoutingError::TabletMoved { .. }));
2531 let guidance = retry_guidance(&error);
2532 assert_eq!(guidance.category(), ErrorCategory::TabletMoved);
2533
2534 descriptor.state = TabletState::Active;
2537 let error = check_generation(&descriptor, 5).unwrap_err();
2538 assert!(matches!(error, RoutingError::StaleMetadata { .. }));
2539 let guidance = retry_guidance(&error);
2540 assert_eq!(guidance.category(), ErrorCategory::StaleMetadata);
2541
2542 let policy = RetryPolicy::default();
2545 let cache = RoutingCache::new();
2546 let operation = OperationDescriptor {
2547 idempotent: true,
2548 idempotency_key: None,
2549 read_only: true,
2550 deadline: Duration::from_secs(30),
2551 max_attempts: 3,
2552 };
2553 for error in [
2554 RoutingError::TabletSplit {
2555 tablet_id: tablet_id(1),
2556 used_generation: 5,
2557 current_generation: 6,
2558 },
2559 RoutingError::TabletMoved {
2560 tablet_id: tablet_id(1),
2561 used_generation: 5,
2562 current_generation: 7,
2563 },
2564 RoutingError::StaleMetadata {
2565 tablet_id: tablet_id(1),
2566 used_generation: 5,
2567 current_generation: 7,
2568 },
2569 ] {
2570 let mut state = RetryState::default();
2571 let action = policy.decide(
2572 GroupKey::Tablet(tablet_id(1)),
2573 &operation,
2574 &mut state,
2575 &retry_guidance(&error).failure(),
2576 &cache,
2577 Duration::ZERO,
2578 );
2579 assert!(
2580 matches!(action, RetryAction::RefreshMetadata { .. }),
2581 "{error} did not map onto a metadata refresh"
2582 );
2583 }
2584 }
2585
2586 #[test]
2589 fn progress_record_round_trips_and_fails_closed() {
2590 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2591 let fixture = split_fixture();
2592 let executor = begin_executor(&fixture);
2593 let path = fixture
2594 .source_layout
2595 .tablet_dir()
2596 .join(SPLIT_PROGRESS_FILENAME);
2597 assert!(path.is_file());
2598 drop(executor);
2599
2600 std::fs::write(&path, b"{ not json").unwrap();
2602 assert!(matches!(
2603 SplitExecutor::resume(
2604 fixture.source_layout.clone(),
2605 fixture.meta.clone(),
2606 fixture.keyspace.clone(),
2607 fixture.sinks.clone(),
2608 ),
2609 Err(SplitError::Tablet(TabletError::Metadata(
2610 ClusterError::CorruptMetadata { .. }
2611 )))
2612 ));
2613
2614 let executor = begin_executor(&fixture);
2616 drop(executor);
2617 let mut value: serde_json::Value =
2618 serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
2619 value["format_version"] = serde_json::json!(99);
2620 std::fs::write(&path, serde_json::to_vec(&value).unwrap()).unwrap();
2621 assert!(matches!(
2622 SplitExecutor::resume(
2623 fixture.source_layout.clone(),
2624 fixture.meta.clone(),
2625 fixture.keyspace.clone(),
2626 fixture.sinks.clone(),
2627 ),
2628 Err(SplitError::Tablet(TabletError::Metadata(
2629 ClusterError::UnsupportedFormatVersion { found: 99, .. }
2630 )))
2631 ));
2632
2633 let executor = begin_executor(&fixture);
2635 drop(executor);
2636 let mut value: serde_json::Value =
2637 serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
2638 value["progress"]["split_ts"]["physical_micros"] = serde_json::json!(999);
2639 std::fs::write(&path, serde_json::to_vec(&value).unwrap()).unwrap();
2640 assert!(matches!(
2641 SplitExecutor::resume(
2642 fixture.source_layout.clone(),
2643 fixture.meta.clone(),
2644 fixture.keyspace.clone(),
2645 fixture.sinks.clone(),
2646 ),
2647 Err(SplitError::Tablet(TabletError::Metadata(
2648 ClusterError::CorruptMetadata { .. }
2649 )))
2650 ));
2651 }
2652
2653 #[test]
2654 fn begin_rejects_a_mismatched_source_layout_and_missing_replica() {
2655 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2656 let fixture = split_fixture();
2657 let foreign = TabletLayout::new(
2659 fixture._dir.path(),
2660 tablet_id(9),
2661 fixture.source.raft_group_id,
2662 );
2663 assert!(matches!(
2664 SplitExecutor::begin(
2665 fixture.plan.clone(),
2666 foreign,
2667 fixture.meta.clone(),
2668 fixture.keyspace.clone(),
2669 fixture.sinks.clone(),
2670 ),
2671 Err(SplitError::Tablet(TabletError::TabletMismatch { .. }))
2672 ));
2673 let other_dir = tempfile::tempdir().unwrap();
2675 let missing = TabletLayout::new(
2676 other_dir.path(),
2677 fixture.source.tablet_id,
2678 fixture.source.raft_group_id,
2679 );
2680 assert!(matches!(
2681 SplitExecutor::begin(
2682 fixture.plan.clone(),
2683 missing,
2684 fixture.meta.clone(),
2685 fixture.keyspace.clone(),
2686 fixture.sinks.clone(),
2687 ),
2688 Err(SplitError::Tablet(TabletError::MissingMetadata(_)))
2689 ));
2690 }
2691
2692 #[test]
2695 fn abort_before_publish_restores_the_source_and_removes_the_children() {
2696 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2697 let fixture = split_fixture();
2698 let mut executor = begin_executor(&fixture);
2699 executor.run_until(SplitPhase::SnapshotPinned).unwrap();
2700 drop(executor);
2701 let marked = fixture.meta.tablet(fixture.source.tablet_id).unwrap();
2704 assert_eq!(marked.state, TabletState::Splitting);
2705 assert_eq!(marked.generation, 6);
2706 for id in [tablet_id(2), tablet_id(3)] {
2707 assert!(fixture.meta.tablet(id).is_some());
2708 }
2709 assert!(fixture
2710 .source_layout
2711 .tablet_dir()
2712 .join(SPLIT_PROGRESS_FILENAME)
2713 .is_file());
2714
2715 let mut meta = fixture.meta.clone();
2716 let report = abort_split(&fixture.source_layout, &mut meta).unwrap();
2717 assert_eq!(report.source, fixture.source.tablet_id);
2718 assert_eq!(report.phase, Some(SplitPhase::SnapshotPinned));
2719 assert_eq!(report.children_removed, vec![tablet_id(2), tablet_id(3)]);
2720 let restored = report.source_after.unwrap();
2723 assert_eq!(restored.state, TabletState::Active);
2724 assert_eq!(restored.generation, 7);
2725 assert_eq!(
2726 meta.tablet(fixture.source.tablet_id),
2727 Some(restored.clone())
2728 );
2729 assert_eq!(fixture.source_layout.load_metadata().unwrap(), restored);
2730 for (index, id) in [tablet_id(2), tablet_id(3)].into_iter().enumerate() {
2733 assert!(meta.tablet(id).is_none());
2734 assert!(!fixture.plan.children[index].layout.tablet_dir().exists());
2735 }
2736 assert!(!fixture
2737 .source_layout
2738 .tablet_dir()
2739 .join(SPLIT_PROGRESS_FILENAME)
2740 .exists());
2741 assert_eq!(fixture.keyspace.rows_at(ts(u64::MAX)).len(), 14);
2743 }
2744
2745 #[test]
2746 fn abort_is_idempotent_across_a_mid_abort_crash() {
2747 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2748 let fixture = split_fixture();
2749 let mut executor = begin_executor(&fixture);
2750 executor.run_until(SplitPhase::ChildrenCreated).unwrap();
2751 drop(executor);
2752
2753 let mut meta = fixture.meta.clone();
2754 let first = abort_split(&fixture.source_layout, &mut meta).unwrap();
2755 assert_eq!(first.phase, Some(SplitPhase::ChildrenCreated));
2756 let second = abort_split(&fixture.source_layout, &mut meta).unwrap();
2758 assert_eq!(second.phase, None);
2759 assert!(second.children_removed.is_empty());
2760 assert!(second.source_after.is_none());
2761 let restored = meta.tablet(fixture.source.tablet_id).unwrap();
2763 assert_eq!(restored.state, TabletState::Active);
2764 assert_eq!(restored.generation, 7);
2765 }
2766
2767 #[test]
2768 fn abort_at_started_unwinds_before_any_meta_write() {
2769 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2770 let fixture = split_fixture();
2771 let executor = begin_executor(&fixture);
2772 drop(executor); let mut meta = fixture.meta.clone();
2775 let report = abort_split(&fixture.source_layout, &mut meta).unwrap();
2776 assert_eq!(report.phase, Some(SplitPhase::Started));
2777 assert!(report.children_removed.is_empty());
2778 let restored = report.source_after.unwrap();
2780 assert_eq!(restored, fixture.source);
2781 assert!(!fixture
2782 .source_layout
2783 .tablet_dir()
2784 .join(SPLIT_PROGRESS_FILENAME)
2785 .exists());
2786 }
2787
2788 #[test]
2789 fn abort_after_publish_fails_closed() {
2790 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2791 let fixture = split_fixture();
2792 let mut executor = begin_executor(&fixture);
2793 executor.run_until(SplitPhase::Published).unwrap();
2794 drop(executor);
2795
2796 let mut meta = fixture.meta.clone();
2797 let error = abort_split(&fixture.source_layout, &mut meta).unwrap_err();
2798 assert!(matches!(
2799 error,
2800 SplitError::CannotAbort { tablet, phase }
2801 if tablet == fixture.source.tablet_id && phase == SplitPhase::Published
2802 ));
2803 assert_eq!(
2805 meta.tablet(fixture.source.tablet_id).unwrap().state,
2806 TabletState::Retiring
2807 );
2808 for id in [tablet_id(2), tablet_id(3)] {
2809 assert_eq!(meta.tablet(id).unwrap().state, TabletState::Active);
2810 }
2811 }
2812
2813 #[test]
2814 fn the_source_can_split_again_after_an_abort() {
2815 let _lock = EXECUTOR_TEST_LOCK.lock().unwrap();
2816 let fixture = split_fixture();
2817 let mut executor = begin_executor(&fixture);
2818 executor.run_until(SplitPhase::ChildrenCreated).unwrap();
2819 drop(executor);
2820 let mut meta = fixture.meta.clone();
2821 abort_split(&fixture.source_layout, &mut meta).unwrap();
2822
2823 let restored = meta.tablet(fixture.source.tablet_id).unwrap();
2826 let planner = TabletSplitPlanner::new(fixture._dir.path());
2827 let plan = planner
2828 .plan(
2829 &restored,
2830 SplitKeySelection::Explicit(text_key("n")),
2831 ts(300),
2832 [allocation(4, 4, 41), allocation(5, 5, 51)],
2833 )
2834 .unwrap();
2835 let sinks = [MapChildSink::new(), MapChildSink::new()];
2836 let mut executor = SplitExecutor::begin(
2837 plan,
2838 fixture.source_layout.clone(),
2839 meta.clone(),
2840 fixture.keyspace.clone(),
2841 sinks.clone(),
2842 )
2843 .unwrap();
2844 executor.run().unwrap();
2845 assert!(meta.tablet(fixture.source.tablet_id).is_none());
2846 let mut union = sinks[0].rows();
2847 union.extend(sinks[1].rows());
2848 assert_eq!(union, fixture.keyspace.rows_at(ts(u64::MAX)));
2849 }
2850}