1use std::collections::BTreeMap;
64use std::path::{Path, PathBuf};
65use std::sync::{Arc, Mutex};
66use std::time::Duration;
67
68use mongreldb_consensus::engine_sink::{
69 build_tablet_write_envelope, open_engine_sink, EngineApplySink, EngineGroupConfig,
70 EngineSinkError, EngineTabletPin, TabletDataCommand, TabletDataCommandRecord,
71 TabletDataMutation, TabletTableBinding, TabletWriteOperation, COMMAND_TYPE_TABLET_DATA,
72};
73use mongreldb_consensus::error::ConsensusError;
74use mongreldb_consensus::group::{ConsensusGroup, GroupCommitReceipt, GroupConfig, GroupMetrics};
75use mongreldb_consensus::identity::{raft_node_id, CommandKind, RaftNodeId};
76use mongreldb_consensus::raft_log::RaftCommitLog;
77use mongreldb_consensus::state_machine::ApplySink;
78use mongreldb_log::commit_log::{CommitLog, ExecutionControl, LogPosition};
79use mongreldb_log::envelope::CommandEnvelope;
80use mongreldb_types::hlc::HlcTimestamp;
81use mongreldb_types::ids::{DatabaseId, MetadataVersion, NodeId, RaftGroupId, TabletId};
82use openraft::BasicNode;
83#[cfg(test)]
84use serde::{Deserialize, Serialize};
85
86use crate::merge::{
87 merge_progress, MergeExecutor, MergeInputs, MergeMetaPlane, MergePhase, MergePlan,
88 MergePlanner, MergePublishCommand,
89};
90use crate::meta::{MetaCommand, MetaError, MetaGroup, MetaGroupConfig, MetaRejectionReason};
91use crate::network::{
92 InternalRpcHandler, PeerEndpoint, TcpTransport, TransportConfig, TransportError,
93 TransportSecurity, TransportServer,
94};
95use crate::node::{ClusterError, NodeIdentity};
96use crate::split::{
97 abort_split, split_progress, ChildAllocation, ChildStateSink, SnapshotPin, SplitAbortReport,
98 SplitError, SplitExecutor, SplitKeySelection, SplitPhase, SplitPlan, SplitPublishCommand,
99 TabletDataError, TabletKeyspace, TabletMetaPlane, TabletMutation, TabletSplitPlanner,
100};
101use crate::tablet::{
102 Key, ReplicaDescriptor, TablePartitioningRecord, TabletDescriptor, TabletError, TabletLayout,
103 TabletOwnershipGuard, TabletOwnershipRegistry, TabletState, TABLETS_DIR, TABLET_META_FILENAME,
104};
105
106#[derive(Debug, thiserror::Error)]
108pub enum RuntimeError {
109 #[error(transparent)]
111 Cluster(#[from] ClusterError),
112 #[error(transparent)]
114 Meta(#[from] MetaError),
115 #[error(transparent)]
117 Consensus(#[from] ConsensusError),
118 #[error(transparent)]
120 Transport(#[from] TransportError),
121 #[error(transparent)]
123 Tablet(#[from] TabletError),
124 #[error(transparent)]
126 EngineSink(#[from] EngineSinkError),
127 #[error(transparent)]
129 Split(#[from] SplitError),
130 #[error(transparent)]
132 Merge(#[from] crate::merge::MergeError),
133 #[error("runtime I/O error: {0}")]
135 Io(#[from] std::io::Error),
136 #[error("invalid runtime request: {0}")]
139 InvalidRequest(String),
140}
141
142#[derive(Clone, Copy, Debug, PartialEq, Eq)]
146pub struct GroupTiming {
147 pub heartbeat_interval: Duration,
149 pub election_timeout_min: Duration,
151 pub election_timeout_max: Duration,
153 pub install_snapshot_timeout: Duration,
155}
156
157impl GroupTiming {
158 fn apply(&self, config: &mut GroupConfig) {
159 config.heartbeat_interval = self.heartbeat_interval;
160 config.election_timeout_min = self.election_timeout_min;
161 config.election_timeout_max = self.election_timeout_max;
162 config.install_snapshot_timeout = self.install_snapshot_timeout;
163 }
164}
165
166#[derive(Clone, Debug)]
170pub struct MetaMembership {
171 pub meta_group_id: RaftGroupId,
174 pub bootstrap_voters: Option<Vec<(NodeId, String)>>,
179}
180
181#[derive(Clone, Debug)]
183pub struct NodeRuntimeConfig {
184 pub node_data: PathBuf,
186 pub security: TransportSecurity,
188 pub transport: TransportConfig,
190 pub listen_address: String,
193 pub rpc_address: Option<String>,
196 pub peers: Vec<(NodeId, String)>,
200 pub meta: Option<MetaMembership>,
202 pub timing: Option<GroupTiming>,
204}
205
206impl NodeRuntimeConfig {
207 #[cfg(any(test, feature = "dangerous-test-transport"))]
209 pub fn plaintext_for_tests(node_data: PathBuf, listen_address: String) -> Self {
210 Self::new_inner(
211 node_data,
212 listen_address,
213 TransportSecurity::PlaintextForTesting,
214 )
215 }
216
217 pub fn production(
221 node_data: PathBuf,
222 listen_address: String,
223 security: TransportSecurity,
224 ) -> Result<Self, String> {
225 match security {
226 TransportSecurity::PlaintextForTesting => {
227 Err("production NodeRuntimeConfig rejects plaintext cluster transport".into())
228 }
229 other => Ok(Self::new_inner(node_data, listen_address, other)),
230 }
231 }
232
233 fn new_inner(node_data: PathBuf, listen_address: String, security: TransportSecurity) -> Self {
234 Self {
235 node_data,
236 security,
237 transport: TransportConfig::default(),
238 listen_address,
239 rpc_address: None,
240 peers: Vec::new(),
241 meta: None,
242 timing: None,
243 }
244 }
245
246 #[cfg(any(test, feature = "dangerous-test-transport"))]
248 pub fn new(node_data: PathBuf, listen_address: String) -> Self {
249 Self::plaintext_for_tests(node_data, listen_address)
250 }
251}
252
253#[cfg(test)]
254mod node_runtime_config_tests {
255 use super::*;
256 use std::path::PathBuf;
257
258 #[test]
259 fn production_rejects_plaintext_transport() {
260 let err = NodeRuntimeConfig::production(
261 PathBuf::from("/tmp/node"),
262 "127.0.0.1:0".into(),
263 TransportSecurity::PlaintextForTesting,
264 )
265 .unwrap_err();
266 assert!(err.contains("plaintext"), "{err}");
267 }
268
269 #[test]
270 fn plaintext_for_tests_is_usable() {
271 let cfg = NodeRuntimeConfig::plaintext_for_tests(
272 PathBuf::from("/tmp/node"),
273 "127.0.0.1:0".into(),
274 );
275 assert!(matches!(
276 cfg.security,
277 TransportSecurity::PlaintextForTesting
278 ));
279 }
280}
281
282#[derive(Clone, Debug)]
284pub struct MetaGroupStatus {
285 pub meta_group_id: RaftGroupId,
287 pub metadata_version: MetadataVersion,
289 pub metrics: GroupMetrics,
291}
292
293#[derive(Clone, Debug)]
295pub struct TabletGroupStatus {
296 pub tablet_id: TabletId,
298 pub raft_group_id: RaftGroupId,
300 pub state: TabletState,
302 pub replicas: Vec<ReplicaDescriptor>,
304 pub applied: LogPosition,
306 pub metrics: GroupMetrics,
308}
309
310#[derive(Clone, Debug)]
313pub struct RuntimeStatus {
314 pub identity: NodeIdentity,
316 pub rpc_address: String,
318 pub meta: Option<MetaGroupStatus>,
320 pub tablets: Vec<TabletGroupStatus>,
322}
323
324fn resolve_tablet_database_id(descriptor: &crate::tablet::TabletDescriptor) -> DatabaseId {
332 if descriptor.database_id != DatabaseId::ZERO {
333 return descriptor.database_id;
334 }
335 DatabaseId::from_bytes(*descriptor.raft_group_id.as_bytes())
336}
337
338fn tablet_group_name(tablet_id: TabletId) -> String {
341 format!("tablet-{}", tablet_id.to_hex())
342}
343
344pub const TABLET_LEDGER_FILENAME: &str = "tablet-ledger.json";
346#[cfg(test)]
348pub const TABLET_LEDGER_FORMAT_VERSION: u32 = 1;
349#[cfg(test)]
351pub const MIN_SUPPORTED_TABLET_LEDGER_FORMAT_VERSION: u32 = 1;
352
353#[cfg(test)]
356const MAX_TIMESTAMP: HlcTimestamp = HlcTimestamp {
357 physical_micros: u64::MAX,
358 logical: u32::MAX,
359 node_tiebreaker: u32::MAX,
360};
361
362#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
364#[cfg(test)]
365struct TabletLedgerCheckpoint {
366 format_version: u32,
368 position: LogPosition,
370 rows: BTreeMap<Key, Vec<(HlcTimestamp, Vec<u8>)>>,
372}
373
374#[derive(Debug)]
377#[cfg(test)]
378pub struct TabletLedger {
379 rows: BTreeMap<Key, Vec<(HlcTimestamp, Vec<u8>)>>,
380 pins: BTreeMap<HlcTimestamp, usize>,
382 position: LogPosition,
383 state_dir: PathBuf,
384}
385
386#[cfg(test)]
387impl TabletLedger {
388 pub fn open(group_dir: &Path) -> Result<Self, RuntimeError> {
390 let state_dir = group_dir.join("raft").join("state");
391 std::fs::create_dir_all(&state_dir).map_err(RuntimeError::Io)?;
392 let path = state_dir.join(TABLET_LEDGER_FILENAME);
393 let Some(bytes) = crate::node::read_meta_file(&path)? else {
394 return Ok(Self {
395 rows: BTreeMap::new(),
396 pins: BTreeMap::new(),
397 position: LogPosition::ZERO,
398 state_dir,
399 });
400 };
401 let checkpoint: TabletLedgerCheckpoint =
402 crate::node::decode_json(TABLET_LEDGER_FILENAME, &bytes)?;
403 if checkpoint.format_version < MIN_SUPPORTED_TABLET_LEDGER_FORMAT_VERSION
404 || checkpoint.format_version > TABLET_LEDGER_FORMAT_VERSION
405 {
406 return Err(ClusterError::UnsupportedFormatVersion {
407 file: TABLET_LEDGER_FILENAME,
408 found: checkpoint.format_version,
409 min: MIN_SUPPORTED_TABLET_LEDGER_FORMAT_VERSION,
410 max: TABLET_LEDGER_FORMAT_VERSION,
411 }
412 .into());
413 }
414 Ok(Self {
415 rows: checkpoint.rows,
416 pins: BTreeMap::new(),
417 position: checkpoint.position,
418 state_dir,
419 })
420 }
421
422 fn apply(
426 &mut self,
427 command: &TabletDataCommand,
428 commit_ts: HlcTimestamp,
429 position: LogPosition,
430 ) -> Result<(), RuntimeError> {
431 if position.index <= self.position.index {
432 return Ok(());
433 }
434 match command {
435 TabletDataCommand::Upsert { entries } => {
436 for (key, value) in entries {
437 self.insert_version(Key::from_bytes(key.clone()), commit_ts, value.clone());
438 }
439 }
440 TabletDataCommand::Delete { keys } => {
441 if !self.pins.is_empty() {
442 return Err(RuntimeError::InvalidRequest(
443 "legacy tablet ledger Delete with live snapshot pins".to_owned(),
444 ));
445 }
446 for key in keys {
447 self.rows.remove(&Key::from_bytes(key.clone()));
448 }
449 }
450 TabletDataCommand::Replace { rows } => {
451 if !self.pins.is_empty() {
456 return Err(RuntimeError::InvalidRequest(
457 "tablet ledger Replace with live snapshot pins".to_owned(),
458 ));
459 }
460 self.rows.clear();
461 for (key, value) in rows {
462 self.insert_version(Key::from_bytes(key.clone()), commit_ts, value.clone());
463 }
464 }
465 }
466 self.position = position;
467 self.persist()
468 }
469
470 fn insert_version(&mut self, key: Key, ts: HlcTimestamp, value: Vec<u8>) {
474 let chain = self.rows.entry(key).or_default();
475 chain.push((ts, value));
476 chain.sort_by_key(|(version, _)| *version);
477 match self.pins.keys().next() {
478 None => {
479 let newest = chain.pop().expect("just inserted");
480 chain.clear();
481 chain.push(newest);
482 }
483 Some(oldest_pin) => {
484 let baseline = chain.partition_point(|(version, _)| *version <= *oldest_pin);
485 if baseline > 1 {
486 chain.drain(..baseline - 1);
487 }
488 }
489 }
490 }
491
492 fn persist(&self) -> Result<(), RuntimeError> {
494 let checkpoint = TabletLedgerCheckpoint {
495 format_version: TABLET_LEDGER_FORMAT_VERSION,
496 position: self.position,
497 rows: self.rows.clone(),
498 };
499 let bytes = crate::node::encode_json(TABLET_LEDGER_FILENAME, &checkpoint)?;
500 crate::node::write_meta_atomic(&self.state_dir, TABLET_LEDGER_FILENAME, &bytes)
501 .map_err(ClusterError::Io)?;
502 Ok(())
503 }
504
505 pub fn applied_position(&self) -> LogPosition {
507 self.position
508 }
509
510 fn pin(&mut self, ts: HlcTimestamp) {
511 *self.pins.entry(ts).or_insert(0) += 1;
512 }
513
514 fn unpin(&mut self, ts: HlcTimestamp) {
515 if let Some(count) = self.pins.get_mut(&ts) {
516 *count -= 1;
517 if *count == 0 {
518 self.pins.remove(&ts);
519 }
520 }
521 }
522
523 pub fn pin_count(&self) -> usize {
525 self.pins.values().sum()
526 }
527
528 pub fn rows_at(&self, ts: HlcTimestamp) -> BTreeMap<Key, Vec<u8>> {
530 self.rows
531 .iter()
532 .filter_map(|(key, chain)| {
533 let visible = chain.iter().rfind(|(version, _)| *version <= ts)?;
534 Some((key.clone(), visible.1.clone()))
535 })
536 .collect()
537 }
538
539 pub fn current_rows(&self) -> BTreeMap<Key, Vec<u8>> {
541 self.rows_at(MAX_TIMESTAMP)
542 }
543
544 pub fn deltas_after(&self, ts: HlcTimestamp) -> Vec<(Key, Vec<u8>)> {
547 let mut deltas: Vec<(HlcTimestamp, Key, Vec<u8>)> = Vec::new();
548 for (key, chain) in &self.rows {
549 for (version, value) in chain {
550 if *version > ts {
551 deltas.push((*version, key.clone(), value.clone()));
552 }
553 }
554 }
555 deltas.sort_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)));
556 deltas
557 .into_iter()
558 .map(|(_, key, value)| (key, value))
559 .collect()
560 }
561
562 pub fn size_bytes(&self) -> u64 {
565 self.rows
566 .iter()
567 .map(|(key, chain)| {
568 let newest = chain.last().expect("non-empty chain");
569 (key.as_bytes().len() + newest.1.len()) as u64
570 })
571 .sum()
572 }
573
574 fn snapshot_bytes(
575 &self,
576 ) -> Result<Vec<u8>, mongreldb_consensus::state_machine::StateMachineError> {
577 serde_json::to_vec(&TabletLedgerCheckpoint {
578 format_version: TABLET_LEDGER_FORMAT_VERSION,
579 position: self.position,
580 rows: self.rows.clone(),
581 })
582 .map_err(|error| {
583 mongreldb_consensus::state_machine::StateMachineError::Sink(format!(
584 "tablet ledger snapshot: {error}"
585 ))
586 })
587 }
588
589 fn install_bytes(
590 &mut self,
591 bytes: &[u8],
592 ) -> Result<(), mongreldb_consensus::state_machine::StateMachineError> {
593 use mongreldb_consensus::state_machine::StateMachineError;
594
595 let checkpoint: TabletLedgerCheckpoint =
596 serde_json::from_slice(bytes).map_err(|error| {
597 StateMachineError::Corrupt(format!("tablet ledger snapshot: {error}"))
598 })?;
599 if checkpoint.format_version < MIN_SUPPORTED_TABLET_LEDGER_FORMAT_VERSION
600 || checkpoint.format_version > TABLET_LEDGER_FORMAT_VERSION
601 {
602 return Err(StateMachineError::Corrupt(format!(
603 "tablet ledger snapshot format version {} is outside \
604 {MIN_SUPPORTED_TABLET_LEDGER_FORMAT_VERSION}..={TABLET_LEDGER_FORMAT_VERSION}",
605 checkpoint.format_version
606 )));
607 }
608 self.rows = checkpoint.rows;
609 self.position = checkpoint.position;
610 self.persist()
611 .map_err(|error| StateMachineError::Sink(error.to_string()))
612 }
613}
614
615#[cfg(test)]
616fn encode_group_snapshot(engine: &[u8], ledger: &[u8]) -> Vec<u8> {
617 let mut out = Vec::with_capacity(12 + engine.len() + ledger.len());
618 out.extend_from_slice(&1_u32.to_le_bytes());
619 out.extend_from_slice(&(engine.len() as u64).to_le_bytes());
620 out.extend_from_slice(engine);
621 out.extend_from_slice(ledger);
622 out
623}
624
625#[cfg(test)]
626fn decode_group_snapshot(
627 bytes: &[u8],
628) -> Result<(Vec<u8>, Vec<u8>), mongreldb_consensus::state_machine::StateMachineError> {
629 use mongreldb_consensus::state_machine::StateMachineError;
630
631 if bytes.len() < 12 {
632 return Err(StateMachineError::Corrupt(
633 "tablet group snapshot: truncated frame".to_owned(),
634 ));
635 }
636 let version = u32::from_le_bytes(bytes[..4].try_into().expect("4 bytes"));
637 if version != 1 {
638 return Err(StateMachineError::Corrupt(format!(
639 "tablet group snapshot format version {version} is not 1"
640 )));
641 }
642 let engine_len = u64::from_le_bytes(bytes[4..12].try_into().expect("8 bytes")) as usize;
643 if bytes.len() < 12 + engine_len {
644 return Err(StateMachineError::Corrupt(
645 "tablet group snapshot: truncated engine payload".to_owned(),
646 ));
647 }
648 Ok((
649 bytes[12..12 + engine_len].to_vec(),
650 bytes[12 + engine_len..].to_vec(),
651 ))
652}
653
654struct TabletGroup {
657 group: Arc<ConsensusGroup<TcpTransport>>,
658 descriptor: TabletDescriptor,
659 sink: Arc<Mutex<EngineApplySink>>,
660 _ownership: TabletOwnershipGuard<'static>,
661}
662
663#[derive(serde::Deserialize)]
667struct TabletFileProbe {
668 tablet: TabletProbe,
669}
670
671#[derive(serde::Deserialize)]
672struct TabletProbe {
673 raft_group_id: RaftGroupId,
674}
675
676fn scan_tablet_layouts(node_data: &Path) -> Result<Vec<(TabletId, RaftGroupId)>, RuntimeError> {
681 let root = node_data.join(TABLETS_DIR);
682 if !root.is_dir() {
683 return Ok(Vec::new());
684 }
685 let mut found = Vec::new();
686 for entry in std::fs::read_dir(&root)? {
687 let entry = entry?;
688 if !entry.file_type()?.is_dir() {
689 continue;
690 }
691 let name = entry.file_name();
692 let Some(name) = name.to_str() else {
693 return Err(RuntimeError::InvalidRequest(format!(
694 "tablet directory name {name:?} is not UTF-8"
695 )));
696 };
697 let tablet_id: TabletId = name.parse().map_err(|_| {
698 RuntimeError::InvalidRequest(format!("tablet directory `{name}` is not a tablet id"))
699 })?;
700 let probe_path = entry.path().join(TABLET_META_FILENAME);
701 let Some(bytes) = crate::node::read_meta_file(&probe_path)? else {
702 return Err(TabletError::MissingMetadata(probe_path).into());
703 };
704 let probe: TabletFileProbe = serde_json::from_slice(&bytes).map_err(|error| {
705 RuntimeError::InvalidRequest(format!(
706 "tablet metadata probe at {} failed: {error}",
707 probe_path.display()
708 ))
709 })?;
710 found.push((tablet_id, probe.tablet.raft_group_id));
711 }
712 found.sort();
713 Ok(found)
714}
715
716fn active_snapshot_timestamp(
717 node_data: &Path,
718 tablet_id: TabletId,
719) -> Result<Option<HlcTimestamp>, RuntimeError> {
720 let mut found = None;
721 for (source_id, raft_group_id) in scan_tablet_layouts(node_data)? {
722 let layout = TabletLayout::new(node_data.to_path_buf(), source_id, raft_group_id);
723 let timestamp = if let Some(progress) = split_progress(&layout)?.filter(|progress| {
724 progress.source.tablet_id == tablet_id && progress.phase < SplitPhase::Published
725 }) {
726 Some(progress.split_ts)
727 } else {
728 merge_progress(&layout)?
729 .filter(|progress| {
730 progress.phase < MergePhase::Published
731 && progress
732 .sources
733 .iter()
734 .any(|source| source.tablet_id == tablet_id)
735 })
736 .map(|progress| progress.merge_ts)
737 };
738 if let Some(timestamp) = timestamp {
739 if found
740 .replace(timestamp)
741 .is_some_and(|prior| prior != timestamp)
742 {
743 return Err(RuntimeError::InvalidRequest(format!(
744 "tablet {tablet_id} appears in conflicting split/merge progress"
745 )));
746 }
747 }
748 }
749 Ok(found)
750}
751
752fn peer_endpoint(security: &TransportSecurity, node_id: NodeId, address: &str) -> PeerEndpoint {
754 match security {
755 TransportSecurity::Mtls(_) => PeerEndpoint::mtls(address, node_id),
756 TransportSecurity::PlaintextForTesting => PeerEndpoint::plaintext(address),
757 }
758}
759
760struct TabletOpenContext<'a> {
763 node_data: &'a Path,
764 security: &'a TransportSecurity,
765 timing: Option<GroupTiming>,
766 identity: &'a NodeIdentity,
767 peers: &'a BTreeMap<NodeId, String>,
768}
769
770pub struct NodeRuntime {
773 identity: NodeIdentity,
774 node_data: PathBuf,
775 rpc_address: String,
776 security: TransportSecurity,
777 peers: BTreeMap<NodeId, String>,
778 timing: Option<GroupTiming>,
779 transport: Arc<TcpTransport>,
780 server: Option<TransportServer>,
781 meta: Option<Arc<MetaGroup<TcpTransport>>>,
782 tablets: BTreeMap<TabletId, TabletGroup>,
783}
784
785#[derive(Clone)]
787pub struct NodeInternalRpcClient {
788 transport: Arc<TcpTransport>,
789}
790
791impl NodeInternalRpcClient {
792 pub async fn call(
794 &self,
795 target: NodeId,
796 service_id: u32,
797 body: Vec<u8>,
798 ) -> Result<Vec<u8>, RuntimeError> {
799 self.transport
800 .internal_rpc(raft_node_id(&target), service_id, body)
801 .await
802 .map_err(Into::into)
803 }
804}
805
806impl NodeRuntime {
807 pub async fn start(config: NodeRuntimeConfig) -> Result<Self, RuntimeError> {
809 let identity =
810 NodeIdentity::load(&config.node_data)?.ok_or(ClusterError::NotInitialized)?;
811 let transport = Arc::new(TcpTransport::new(
812 config.transport.clone(),
813 config.security.clone(),
814 ));
815 let mut peers = BTreeMap::new();
816 for (node_id, address) in &config.peers {
817 transport.upsert_peer(
818 raft_node_id(node_id),
819 peer_endpoint(&config.security, *node_id, address),
820 );
821 peers.insert(*node_id, address.clone());
822 }
823 let server = TransportServer::bind_shared(
824 &config.listen_address,
825 transport.security_handle(),
826 transport.registry(),
827 config.transport.clone(),
828 )
829 .await?;
830 let rpc_address = config
831 .rpc_address
832 .clone()
833 .unwrap_or_else(|| server.local_addr().to_string());
834
835 let meta = match &config.meta {
836 Some(membership) => {
837 let group =
838 Self::start_meta_group(&config, &identity, membership, transport.clone())
839 .await?;
840 Some(Arc::new(group))
841 }
842 None => None,
843 };
844 let tablets =
845 Self::open_tablet_groups(&config, &identity, &peers, transport.clone()).await?;
846 Ok(Self {
847 identity,
848 node_data: config.node_data.clone(),
849 rpc_address,
850 security: config.security.clone(),
851 peers,
852 timing: config.timing,
853 transport,
854 server: Some(server),
855 meta,
856 tablets,
857 })
858 }
859
860 async fn start_meta_group(
863 config: &NodeRuntimeConfig,
864 identity: &NodeIdentity,
865 membership: &MetaMembership,
866 transport: Arc<TcpTransport>,
867 ) -> Result<MetaGroup<TcpTransport>, RuntimeError> {
868 let meta_config = MetaGroupConfig::new(
869 config.node_data.clone(),
870 membership.meta_group_id,
871 identity.node_id,
872 );
873 let mut group_config = meta_config.group_config();
874 if let Some(timing) = &config.timing {
875 timing.apply(&mut group_config);
876 }
877 let group = MetaGroup::create(meta_config, group_config, transport).await?;
878 if let Some(voters) = &membership.bootstrap_voters {
879 if !voters
880 .iter()
881 .any(|(node_id, _)| *node_id == identity.node_id)
882 {
883 return Err(RuntimeError::InvalidRequest(
884 "meta bootstrap voter set does not include this node".to_owned(),
885 ));
886 }
887 if !group.is_initialized().await? {
888 group.bootstrap(voters).await?;
889 }
890 }
891 Ok(group)
892 }
893
894 async fn open_tablet_groups(
896 config: &NodeRuntimeConfig,
897 identity: &NodeIdentity,
898 peers: &BTreeMap<NodeId, String>,
899 transport: Arc<TcpTransport>,
900 ) -> Result<BTreeMap<TabletId, TabletGroup>, RuntimeError> {
901 let mut groups = BTreeMap::new();
902 let context = TabletOpenContext {
903 node_data: &config.node_data,
904 security: &config.security,
905 timing: config.timing,
906 identity,
907 peers,
908 };
909 for (tablet_id, raft_group_id) in scan_tablet_layouts(&config.node_data)? {
910 let layout = TabletLayout::new(config.node_data.clone(), tablet_id, raft_group_id);
911 let descriptor = layout.validate()?;
912 let group =
913 Self::open_tablet_group(&context, transport.clone(), &layout, &descriptor).await?;
914 groups.insert(tablet_id, group);
915 }
916 Ok(groups)
917 }
918
919 async fn open_tablet_group(
924 context: &TabletOpenContext<'_>,
925 transport: Arc<TcpTransport>,
926 layout: &TabletLayout,
927 descriptor: &TabletDescriptor,
928 ) -> Result<TabletGroup, RuntimeError> {
929 let replica = *descriptor
930 .replica_on(context.identity.node_id)
931 .ok_or_else(|| {
932 RuntimeError::InvalidRequest(format!(
933 "tablet {} descriptor does not list this node ({}) as a replica",
934 descriptor.tablet_id, context.identity.node_id
935 ))
936 })?;
937 if transport.registry().get(replica.raft_node_id).is_some() {
938 return Err(RuntimeError::InvalidRequest(format!(
939 "raft id {} is already attached to this node's transport registry",
940 replica.raft_node_id
941 )));
942 }
943 for other in &descriptor.replicas {
946 let Some(address) = context.peers.get(&other.node_id) else {
947 return Err(RuntimeError::InvalidRequest(format!(
948 "no membership-directory address for replica node {} of tablet {}",
949 other.node_id, descriptor.tablet_id
950 )));
951 };
952 transport.upsert_peer(
953 other.raft_node_id,
954 peer_endpoint(context.security, other.node_id, address),
955 );
956 }
957 let ownership = TabletOwnershipRegistry::global().try_reserve(layout)?;
958 let database_id = resolve_tablet_database_id(descriptor);
959 let engine_config = EngineGroupConfig::new(
960 context.node_data.to_path_buf(),
961 descriptor.raft_group_id,
962 context.identity.cluster_id,
963 context.identity.node_id,
964 database_id,
965 );
966 let sink = open_engine_sink(&engine_config)?;
967 {
968 let mut engine = sink
969 .lock()
970 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?;
971 match engine.initialize_tablet_keyspace() {
977 Ok(()) => {}
978 Err(error) => {
979 let message = error.to_string();
980 if message.contains("opaque tablet keyspace is disabled") {
981 } else {
983 return Err(RuntimeError::from(error));
984 }
985 }
986 }
987 let legacy_path = engine_config
988 .group_dir()
989 .join("raft")
990 .join("state")
991 .join(TABLET_LEDGER_FILENAME);
992 if legacy_path.is_file() {
993 let bytes = std::fs::read(&legacy_path)?;
994 engine.migrate_legacy_tablet_ledger(
995 &bytes,
996 active_snapshot_timestamp(context.node_data, descriptor.tablet_id)?,
997 )?;
998 std::fs::remove_file(&legacy_path)?;
999 if let Some(parent) = legacy_path.parent() {
1000 std::fs::File::open(parent)?.sync_all()?;
1001 }
1002 }
1003 engine.finish_legacy_tablet_migration()?;
1004 }
1005 let mut group_config = GroupConfig::new(
1006 tablet_group_name(descriptor.tablet_id),
1007 replica.raft_node_id,
1008 engine_config.group_dir(),
1009 );
1010 group_config.storage = engine_config.storage.clone();
1011 group_config.idempotency_retention = engine_config.idempotency_retention;
1012 if let Some(timing) = &context.timing {
1013 timing.apply(&mut group_config);
1014 }
1015 let dyn_sink: Arc<Mutex<dyn ApplySink>> = sink.clone();
1016 let group = Arc::new(ConsensusGroup::create(group_config, transport, dyn_sink).await?);
1017 let commit_log: Arc<dyn CommitLog> = Arc::new(RaftCommitLog::new(group.clone()));
1018 sink.lock()
1019 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
1020 .bind_commit_log(commit_log)?;
1021 Ok(TabletGroup {
1022 group,
1023 descriptor: descriptor.clone(),
1024 sink,
1025 _ownership: ownership,
1026 })
1027 }
1028
1029 pub fn identity(&self) -> &NodeIdentity {
1031 &self.identity
1032 }
1033
1034 pub fn reload_trust(
1038 &mut self,
1039 trust: crate::bootstrap::TrustConfig,
1040 ) -> Result<(), RuntimeError> {
1041 let tls = crate::network::TlsConfig::from_trust(&trust)
1042 .map_err(|e| RuntimeError::InvalidRequest(format!("reload trust: {e}")))?;
1043 let security = crate::network::TransportSecurity::Mtls(tls);
1044 self.transport.reload_security(security.clone());
1045 self.security = security;
1046 Ok(())
1047 }
1048
1049 pub fn rpc_address(&self) -> &str {
1051 &self.rpc_address
1052 }
1053
1054 pub fn attach_internal_rpc_handler(
1059 &self,
1060 service_id: u32,
1061 handler: Arc<dyn InternalRpcHandler>,
1062 ) {
1063 self.transport.registry().attach_internal(
1064 raft_node_id(&self.identity.node_id),
1065 service_id,
1066 handler,
1067 );
1068 }
1069
1070 pub async fn internal_rpc(
1073 &self,
1074 target: NodeId,
1075 service_id: u32,
1076 body: Vec<u8>,
1077 ) -> Result<Vec<u8>, RuntimeError> {
1078 self.transport
1079 .internal_rpc(raft_node_id(&target), service_id, body)
1080 .await
1081 .map_err(Into::into)
1082 }
1083
1084 pub fn internal_rpc_client(&self) -> NodeInternalRpcClient {
1086 NodeInternalRpcClient {
1087 transport: Arc::clone(&self.transport),
1088 }
1089 }
1090
1091 pub fn meta_group(&self) -> Option<&MetaGroup<TcpTransport>> {
1093 self.meta.as_deref()
1094 }
1095
1096 pub fn tablet_group(&self, tablet_id: TabletId) -> Option<&ConsensusGroup<TcpTransport>> {
1098 self.tablets.get(&tablet_id).map(|tablet| &*tablet.group)
1099 }
1100
1101 pub fn tablet_sink(&self, tablet_id: TabletId) -> Option<Arc<Mutex<EngineApplySink>>> {
1107 self.tablets
1108 .get(&tablet_id)
1109 .map(|tablet| tablet.sink.clone())
1110 }
1111
1112 pub fn tablet_descriptor(&self, tablet_id: TabletId) -> Option<&TabletDescriptor> {
1114 self.tablets
1115 .get(&tablet_id)
1116 .map(|tablet| &tablet.descriptor)
1117 }
1118
1119 pub fn tablet_ids(&self) -> Vec<TabletId> {
1121 self.tablets.keys().copied().collect()
1122 }
1123
1124 fn hosted_tablet(&self, tablet_id: TabletId) -> Result<&TabletGroup, RuntimeError> {
1127 self.tablets.get(&tablet_id).ok_or_else(|| {
1128 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
1129 })
1130 }
1131
1132 pub async fn create_tablet(
1149 &mut self,
1150 descriptor: &TabletDescriptor,
1151 partitioning: &TablePartitioningRecord,
1152 bootstrap_voters: Option<&[(NodeId, String)]>,
1153 publish_to_meta: bool,
1154 control: &ExecutionControl,
1155 ) -> Result<(), RuntimeError> {
1156 descriptor.validate()?;
1157 partitioning.validate()?;
1158 if partitioning.table_id != descriptor.table_id {
1159 return Err(RuntimeError::InvalidRequest(format!(
1160 "partitioning record names table {} but the descriptor partitions table {}",
1161 partitioning.table_id, descriptor.table_id
1162 )));
1163 }
1164 self.create_hosted_replica(descriptor, bootstrap_voters)
1165 .await?;
1166 if publish_to_meta {
1167 self.publish_tablet_descriptor(descriptor, control).await?;
1168 }
1169 Ok(())
1170 }
1171
1172 async fn create_hosted_replica(
1180 &mut self,
1181 descriptor: &TabletDescriptor,
1182 bootstrap_voters: Option<&[(NodeId, String)]>,
1183 ) -> Result<(), RuntimeError> {
1184 descriptor.validate()?;
1185 if let Some(existing) = self.tablets.get(&descriptor.tablet_id) {
1186 if existing.descriptor == *descriptor {
1187 return Ok(());
1188 }
1189 return Err(RuntimeError::InvalidRequest(format!(
1190 "tablet {} is already open with a different descriptor",
1191 descriptor.tablet_id
1192 )));
1193 }
1194 if descriptor.replica_on(self.identity.node_id).is_none() {
1197 return Err(RuntimeError::InvalidRequest(format!(
1198 "tablet {} descriptor does not list this node ({}) as a replica",
1199 descriptor.tablet_id, self.identity.node_id
1200 )));
1201 }
1202 let layout = TabletLayout::new(
1203 self.node_data.clone(),
1204 descriptor.tablet_id,
1205 descriptor.raft_group_id,
1206 );
1207 layout.create(descriptor)?;
1208 let context = TabletOpenContext {
1209 node_data: &self.node_data,
1210 security: &self.security,
1211 timing: self.timing,
1212 identity: &self.identity,
1213 peers: &self.peers,
1214 };
1215 let group =
1216 Self::open_tablet_group(&context, self.transport.clone(), &layout, descriptor).await?;
1217 if let Some(voters) = bootstrap_voters {
1218 let mut members = BTreeMap::new();
1219 for (node_id, address) in voters {
1220 let Some(replica) = descriptor.replica_on(*node_id) else {
1221 return Err(RuntimeError::InvalidRequest(format!(
1222 "bootstrap voter {node_id} is not a replica of tablet {}",
1223 descriptor.tablet_id
1224 )));
1225 };
1226 members.insert(replica.raft_node_id, BasicNode::new(address.clone()));
1227 }
1228 if !members.contains_key(&group.group.node_id()) {
1229 return Err(RuntimeError::InvalidRequest(
1230 "tablet bootstrap voter set does not include this node".to_owned(),
1231 ));
1232 }
1233 if !group.group.is_initialized().await? {
1234 group.group.bootstrap(members).await?;
1235 }
1236 }
1237 self.tablets.insert(descriptor.tablet_id, group);
1238 Ok(())
1239 }
1240
1241 pub async fn add_tablet_replica(
1251 &self,
1252 tablet_id: TabletId,
1253 replica: ReplicaDescriptor,
1254 address: &str,
1255 ) -> Result<(), RuntimeError> {
1256 let tablet = self.hosted_tablet(tablet_id)?;
1257 if tablet.descriptor.replicas.iter().any(|existing| {
1258 existing.node_id == replica.node_id || existing.raft_node_id == replica.raft_node_id
1259 }) {
1260 return Err(RuntimeError::InvalidRequest(format!(
1261 "replica node {} / raft id {} already belongs to tablet {}",
1262 replica.node_id, replica.raft_node_id, tablet_id
1263 )));
1264 }
1265 self.transport.upsert_peer(
1266 replica.raft_node_id,
1267 peer_endpoint(&self.security, replica.node_id, address),
1268 );
1269 for _ in 0..10 {
1276 tablet
1277 .group
1278 .wait_uniform_membership(Duration::from_secs(30))
1279 .await?;
1280 let (voters, learners) = tablet.group.members();
1281 if voters.contains(&replica.raft_node_id) {
1282 return Ok(());
1283 }
1284 if !learners.contains(&replica.raft_node_id) {
1285 match tablet
1289 .group
1290 .add_learner(replica.raft_node_id, BasicNode::new(address.to_owned()))
1291 .await
1292 {
1293 Ok(()) => {}
1294 Err(ConsensusError::MembershipInProgress) => continue,
1295 Err(error) => return Err(error.into()),
1296 }
1297 }
1298 match tablet.group.promote(replica.raft_node_id).await {
1299 Ok(()) => return Ok(()),
1300 Err(ConsensusError::MembershipInProgress) => continue,
1301 Err(error) => return Err(error.into()),
1302 }
1303 }
1304 Err(RuntimeError::InvalidRequest(format!(
1305 "membership change for tablet {tablet_id} did not settle"
1306 )))
1307 }
1308
1309 pub async fn publish_tablet_descriptor(
1314 &self,
1315 descriptor: &TabletDescriptor,
1316 control: &ExecutionControl,
1317 ) -> Result<MetadataVersion, RuntimeError> {
1318 let meta = self.meta.as_ref().ok_or_else(|| {
1319 RuntimeError::InvalidRequest("this node hosts no meta group".to_owned())
1320 })?;
1321 let receipt = meta
1322 .propose(
1323 crate::meta::new_command_id()?,
1324 MetaCommand::SetTabletDescriptor {
1325 descriptor: descriptor.clone(),
1326 },
1327 control,
1328 )
1329 .await?;
1330 Ok(receipt.metadata_version)
1331 }
1332
1333 fn meta_group_or_err(&self) -> Result<Arc<MetaGroup<TcpTransport>>, RuntimeError> {
1341 self.meta
1342 .clone()
1343 .ok_or_else(|| RuntimeError::InvalidRequest("this node hosts no meta group".to_owned()))
1344 }
1345
1346 async fn plan_split(
1351 &self,
1352 tablet_id: TabletId,
1353 split_key: Option<Key>,
1354 control: &ExecutionControl,
1355 ) -> Result<SplitPlan, RuntimeError> {
1356 let meta = self.meta_group_or_err()?;
1357 let source = meta.state().tablet(tablet_id).cloned().ok_or_else(|| {
1358 RuntimeError::InvalidRequest(format!("tablet {tablet_id} is not in the meta state"))
1359 })?;
1360 if source.state != TabletState::Active {
1361 return Err(RuntimeError::InvalidRequest(format!(
1362 "tablet {tablet_id} is in state {}, expected Active",
1363 source.state
1364 )));
1365 }
1366 if source.replica_on(self.identity.node_id).is_none() {
1367 return Err(RuntimeError::InvalidRequest(format!(
1368 "this node hosts no replica of tablet {tablet_id}"
1369 )));
1370 }
1371 let replica_count = source.replicas.len();
1372 let raft_ids = meta
1373 .allocate_raft_node_ids(
1374 2 * u32::try_from(replica_count).unwrap_or(u32::MAX),
1375 control,
1376 )
1377 .await?;
1378 let allocation = |ids: &[RaftNodeId]| ChildAllocation {
1379 tablet_id: TabletId::new_random(),
1380 raft_group_id: RaftGroupId::new_random(),
1381 replicas: source
1382 .replicas
1383 .iter()
1384 .zip(ids)
1385 .map(|(replica, raft_node_id)| ReplicaDescriptor {
1386 node_id: replica.node_id,
1387 role: replica.role,
1388 raft_node_id: *raft_node_id,
1389 })
1390 .collect(),
1391 };
1392 let selection = match split_key {
1393 Some(key) => SplitKeySelection::Explicit(key),
1394 None => SplitKeySelection::Midpoint,
1395 };
1396 let plan = TabletSplitPlanner::new(self.node_data.clone()).plan(
1397 &source,
1398 selection,
1399 now_timestamp(),
1400 [
1401 allocation(&raft_ids[..replica_count]),
1402 allocation(&raft_ids[replica_count..]),
1403 ],
1404 )?;
1405 Ok(plan)
1406 }
1407
1408 async fn child_sinks(
1414 &mut self,
1415 children: &[TabletDescriptor; 2],
1416 control: &ExecutionControl,
1417 ) -> Result<[RuntimeChildSink; 2], RuntimeError> {
1418 for child in children {
1419 if child.replica_on(self.identity.node_id).is_some()
1422 && !self.tablets.contains_key(&child.tablet_id)
1423 {
1424 let voters: Vec<(NodeId, String)> = child
1425 .replicas
1426 .iter()
1427 .filter_map(|replica| {
1428 self.peers
1429 .get(&replica.node_id)
1430 .map(|address| (replica.node_id, address.clone()))
1431 })
1432 .collect();
1433 self.create_hosted_replica(child, Some(voters.as_slice()))
1434 .await?;
1435 }
1436 }
1437 let mut sinks = Vec::with_capacity(2);
1438 for child in children {
1439 let group = self
1440 .tablets
1441 .get(&child.tablet_id)
1442 .ok_or_else(|| {
1443 RuntimeError::InvalidRequest(format!(
1444 "child tablet {} is not hosted on this node",
1445 child.tablet_id
1446 ))
1447 })?
1448 .group
1449 .clone();
1450 sinks.push(RuntimeChildSink {
1451 group,
1452 handle: tokio::runtime::Handle::current(),
1453 control: control.clone(),
1454 staged: None,
1455 });
1456 }
1457 let [lower, upper]: [RuntimeChildSink; 2] = sinks.try_into().map_err(|_| {
1458 RuntimeError::InvalidRequest("expected exactly two child sinks".to_owned())
1459 })?;
1460 Ok([lower, upper])
1461 }
1462
1463 async fn drop_hosted_tablet(&mut self, tablet_id: TabletId) -> Result<(), RuntimeError> {
1466 if let Some(tablet) = self.tablets.remove(&tablet_id) {
1467 tablet.group.shutdown().await?;
1468 }
1469 Ok(())
1470 }
1471
1472 fn refresh_local_descriptor(&mut self, layout: &TabletLayout) -> Result<(), RuntimeError> {
1476 let descriptor = layout.load_metadata()?;
1477 if let Some(tablet) = self.tablets.get_mut(&layout.tablet_id()) {
1478 if descriptor.generation > tablet.descriptor.generation {
1479 tablet.descriptor = descriptor;
1480 }
1481 }
1482 Ok(())
1483 }
1484
1485 pub async fn split_step(
1498 &mut self,
1499 tablet_id: TabletId,
1500 split_key: Option<Key>,
1501 control: &ExecutionControl,
1502 ) -> Result<(SplitPhase, Option<SplitPublishCommand>), RuntimeError> {
1503 let meta = self.meta_group_or_err()?;
1504 let raft_group_id = match self.tablets.get(&tablet_id) {
1505 Some(tablet) => tablet.descriptor.raft_group_id,
1506 None => meta
1510 .state()
1511 .tablet(tablet_id)
1512 .map(|descriptor| descriptor.raft_group_id)
1513 .ok_or_else(|| {
1514 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
1515 })?,
1516 };
1517 let source_layout = TabletLayout::new(self.node_data.clone(), tablet_id, raft_group_id);
1518 let progress = split_progress(&source_layout)?;
1519 let hosted = self.tablets.contains_key(&tablet_id);
1520 if !hosted
1521 && progress
1522 .as_ref()
1523 .is_none_or(|record| record.phase < SplitPhase::Published)
1524 {
1525 return Err(RuntimeError::InvalidRequest(format!(
1526 "this node hosts no tablet {tablet_id}"
1527 )));
1528 }
1529 let keyspace = RuntimeKeyspace {
1530 sink: if progress
1531 .as_ref()
1532 .is_some_and(|record| record.phase >= SplitPhase::Published)
1533 {
1534 None
1535 } else {
1536 Some(self.tablet_sink(tablet_id).ok_or_else(|| {
1537 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
1538 })?)
1539 },
1540 pinned: None,
1541 };
1542 let plane = RuntimeMetaPlane {
1543 meta: meta.clone(),
1544 handle: tokio::runtime::Handle::current(),
1545 control: control.clone(),
1546 };
1547 let mut executor = match progress {
1548 Some(progress) => {
1549 let children = progress.plan().child_descriptors();
1550 let sinks = self.child_sinks(&children, control).await?;
1551 SplitExecutor::resume(source_layout.clone(), plane, keyspace, sinks)?.ok_or_else(
1552 || {
1553 RuntimeError::InvalidRequest(format!(
1554 "split progress of tablet {tablet_id} vanished mid-resume"
1555 ))
1556 },
1557 )?
1558 }
1559 None => {
1560 let plan = self.plan_split(tablet_id, split_key, control).await?;
1561 let sinks = self.child_sinks(&plan.child_descriptors(), control).await?;
1562 SplitExecutor::begin(plan, source_layout.clone(), plane, keyspace, sinks)?
1563 }
1564 };
1565 if executor.phase() == SplitPhase::Published {
1568 if let Some(sink) = self.tablet_sink(tablet_id) {
1569 sink.lock()
1570 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
1571 .release_tablet_snapshot()?;
1572 }
1573 self.drop_hosted_tablet(tablet_id).await?;
1574 }
1575 let phase = tokio::task::spawn_blocking(move || executor.step())
1576 .await
1577 .map_err(|error| {
1578 RuntimeError::InvalidRequest(format!("split driver task failed: {error}"))
1579 })??;
1580 if phase == SplitPhase::Published {
1581 if let Some(sink) = self.tablet_sink(tablet_id) {
1582 sink.lock()
1583 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
1584 .release_tablet_snapshot()?;
1585 }
1586 }
1587 if matches!(phase, SplitPhase::MarkedSplitting | SplitPhase::Published) {
1588 self.refresh_local_descriptor(&source_layout)?;
1589 }
1590 let published = if phase >= SplitPhase::Published {
1591 match split_progress(&source_layout)? {
1596 Some(record) => Some(SplitPublishCommand::from_plan(&record.plan())?),
1597 None => None,
1598 }
1599 } else {
1600 None
1601 };
1602 Ok((phase, published))
1603 }
1604
1605 pub async fn split_tablet(
1611 &mut self,
1612 tablet_id: TabletId,
1613 split_key: Option<Key>,
1614 control: &ExecutionControl,
1615 ) -> Result<SplitPublishCommand, RuntimeError> {
1616 let mut published = None;
1617 loop {
1618 let (phase, command) = self
1619 .split_step(tablet_id, split_key.clone(), control)
1620 .await?;
1621 if command.is_some() {
1622 published = command;
1623 }
1624 if phase == SplitPhase::SourceRetired {
1625 break;
1626 }
1627 }
1628 published.ok_or_else(|| {
1629 RuntimeError::InvalidRequest(format!(
1630 "split of tablet {tablet_id} completed without a publication"
1631 ))
1632 })
1633 }
1634
1635 pub async fn abort_split(
1641 &mut self,
1642 tablet_id: TabletId,
1643 control: &ExecutionControl,
1644 ) -> Result<SplitAbortReport, RuntimeError> {
1645 let meta = self.meta_group_or_err()?;
1646 let source_layout = {
1647 let source = self.hosted_tablet(tablet_id)?;
1648 TabletLayout::new(
1649 self.node_data.clone(),
1650 tablet_id,
1651 source.descriptor.raft_group_id,
1652 )
1653 };
1654 if let Some(progress) = split_progress(&source_layout)? {
1657 if progress.phase >= SplitPhase::Published {
1658 return Err(SplitError::CannotAbort {
1659 tablet: tablet_id,
1660 phase: progress.phase,
1661 }
1662 .into());
1663 }
1664 for child in &progress.children {
1665 self.drop_hosted_tablet(child.tablet_id).await?;
1666 }
1667 }
1668 let mut plane = RuntimeMetaPlane {
1669 meta,
1670 handle: tokio::runtime::Handle::current(),
1671 control: control.clone(),
1672 };
1673 let layout = source_layout.clone();
1674 let report = tokio::task::spawn_blocking(move || abort_split(&layout, &mut plane))
1675 .await
1676 .map_err(|error| {
1677 RuntimeError::InvalidRequest(format!("split abort task failed: {error}"))
1678 })??;
1679 if let Some(sink) = self.tablet_sink(tablet_id) {
1680 sink.lock()
1681 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
1682 .release_tablet_snapshot()?;
1683 }
1684 self.refresh_local_descriptor(&source_layout)?;
1687 Ok(report)
1688 }
1689
1690 pub async fn merge_step(
1698 &mut self,
1699 first: TabletId,
1700 second: TabletId,
1701 control: &ExecutionControl,
1702 ) -> Result<(MergePhase, Option<MergePublishCommand>), RuntimeError> {
1703 let meta = self.meta_group_or_err()?;
1704 let layout_of = |tablet_id: TabletId| -> Result<TabletLayout, RuntimeError> {
1707 let raft_group_id = match self.tablets.get(&tablet_id) {
1708 Some(tablet) => tablet.descriptor.raft_group_id,
1709 None => meta
1710 .state()
1711 .tablet(tablet_id)
1712 .map(|descriptor| descriptor.raft_group_id)
1713 .ok_or_else(|| {
1714 RuntimeError::InvalidRequest(format!(
1715 "this node hosts no tablet {tablet_id}"
1716 ))
1717 })?,
1718 };
1719 Ok(TabletLayout::new(
1720 self.node_data.clone(),
1721 tablet_id,
1722 raft_group_id,
1723 ))
1724 };
1725 let first_layout = layout_of(first)?;
1726 let second_layout = layout_of(second)?;
1727 let phase_in_flight = match (
1732 merge_progress(&first_layout)?,
1733 merge_progress(&second_layout)?,
1734 ) {
1735 (Some(progress), None) => Some(progress.phase),
1736 (None, None) => None,
1737 _ => {
1738 return Err(RuntimeError::InvalidRequest(
1739 "merge progress records on both sources: corrupt state".to_owned(),
1740 ));
1741 }
1742 };
1743 let retired_only = matches!(
1744 phase_in_flight,
1745 Some(MergePhase::Published | MergePhase::SourcesRetired)
1746 );
1747 if !retired_only
1748 && (!self.tablets.contains_key(&first) || !self.tablets.contains_key(&second))
1749 {
1750 return Err(RuntimeError::InvalidRequest(format!(
1751 "this node must host both merge sources ({first}, {second})"
1752 )));
1753 }
1754 let keyspace_of = |layout: &TabletLayout| -> Result<RuntimeKeyspace, RuntimeError> {
1755 Ok(RuntimeKeyspace {
1756 sink: if retired_only {
1757 None
1758 } else {
1759 Some(self.tablet_sink(layout.tablet_id()).ok_or_else(|| {
1760 RuntimeError::InvalidRequest(format!(
1761 "this node hosts no tablet {}",
1762 layout.tablet_id()
1763 ))
1764 })?)
1765 },
1766 pinned: None,
1767 })
1768 };
1769 let (first_keyspace, second_keyspace) =
1770 (keyspace_of(&first_layout)?, keyspace_of(&second_layout)?);
1771 let ordered_layouts = |lower_first: bool| {
1773 if lower_first {
1774 [first_layout.clone(), second_layout.clone()]
1775 } else {
1776 [second_layout.clone(), first_layout.clone()]
1777 }
1778 };
1779 let plane = RuntimeMetaPlane {
1780 meta: meta.clone(),
1781 handle: tokio::runtime::Handle::current(),
1782 control: control.clone(),
1783 };
1784 let mut executor = match phase_in_flight {
1785 Some(_) => {
1786 let progress = merge_progress(&first_layout)?
1787 .or(merge_progress(&second_layout)?)
1788 .expect("phase_in_flight is Some");
1789 let lower_first = progress.sources[0].tablet_id == first;
1790 let sink = self.replacement_sink(&progress.plan(), control).await?;
1791 let keyspaces = if lower_first {
1792 [first_keyspace, second_keyspace]
1793 } else {
1794 [second_keyspace, first_keyspace]
1795 };
1796 MergeExecutor::resume(ordered_layouts(lower_first), plane, keyspaces, sink)?
1797 .ok_or_else(|| {
1798 RuntimeError::InvalidRequest(
1799 "merge progress vanished mid-resume".to_owned(),
1800 )
1801 })?
1802 }
1803 None => {
1804 let plan = self.plan_merge(first, second, control).await?;
1805 let lower_first = plan.sources[0].tablet_id == first;
1806 let sink = self.replacement_sink(&plan, control).await?;
1807 let keyspaces = if lower_first {
1808 [first_keyspace, second_keyspace]
1809 } else {
1810 [second_keyspace, first_keyspace]
1811 };
1812 MergeExecutor::begin(plan, ordered_layouts(lower_first), plane, keyspaces, sink)?
1813 }
1814 };
1815 if executor.phase() == MergePhase::Published {
1818 for tablet_id in [first, second] {
1819 if let Some(sink) = self.tablet_sink(tablet_id) {
1820 sink.lock()
1821 .map_err(|_| {
1822 RuntimeError::InvalidRequest("engine sink lock poisoned".into())
1823 })?
1824 .release_tablet_snapshot()?;
1825 }
1826 }
1827 self.drop_hosted_tablet(first).await?;
1828 self.drop_hosted_tablet(second).await?;
1829 }
1830 let phase = tokio::task::spawn_blocking(move || executor.step())
1831 .await
1832 .map_err(|error| {
1833 RuntimeError::InvalidRequest(format!("merge driver task failed: {error}"))
1834 })??;
1835 if phase == MergePhase::Published {
1836 for tablet_id in [first, second] {
1837 if let Some(sink) = self.tablet_sink(tablet_id) {
1838 sink.lock()
1839 .map_err(|_| {
1840 RuntimeError::InvalidRequest("engine sink lock poisoned".into())
1841 })?
1842 .release_tablet_snapshot()?;
1843 }
1844 }
1845 }
1846 if matches!(phase, MergePhase::MarkedMerging | MergePhase::Published) {
1847 self.refresh_local_descriptor(&first_layout)?;
1848 self.refresh_local_descriptor(&second_layout)?;
1849 }
1850 let published = if phase >= MergePhase::Published {
1851 let progress = merge_progress(&first_layout)?.or(merge_progress(&second_layout)?);
1852 match progress {
1853 Some(record) => Some(MergePublishCommand::from_plan(&record.plan())?),
1854 None => None,
1855 }
1856 } else {
1857 None
1858 };
1859 Ok((phase, published))
1860 }
1861
1862 pub async fn merge_tablets(
1866 &mut self,
1867 first: TabletId,
1868 second: TabletId,
1869 control: &ExecutionControl,
1870 ) -> Result<MergePublishCommand, RuntimeError> {
1871 let mut published = None;
1872 loop {
1873 let (phase, command) = self.merge_step(first, second, control).await?;
1874 if command.is_some() {
1875 published = command;
1876 }
1877 if phase == MergePhase::SourcesRetired {
1878 break;
1879 }
1880 }
1881 published.ok_or_else(|| {
1882 RuntimeError::InvalidRequest("merge completed without a publication".to_owned())
1883 })
1884 }
1885
1886 async fn plan_merge(
1892 &self,
1893 first: TabletId,
1894 second: TabletId,
1895 control: &ExecutionControl,
1896 ) -> Result<MergePlan, RuntimeError> {
1897 let meta = self.meta_group_or_err()?;
1898 let state = meta.state();
1899 let descriptor_of = |tablet_id: TabletId| -> Result<TabletDescriptor, RuntimeError> {
1900 state.tablet(tablet_id).cloned().ok_or_else(|| {
1901 RuntimeError::InvalidRequest(format!("tablet {tablet_id} is not in the meta state"))
1902 })
1903 };
1904 let (first_desc, second_desc) = (descriptor_of(first)?, descriptor_of(second)?);
1905 let schema = state.table(first_desc.table_id).ok_or_else(|| {
1906 RuntimeError::InvalidRequest(format!(
1907 "table {} of tablet {first} is not in the meta state",
1908 first_desc.table_id
1909 ))
1910 })?;
1911 let active_schema_job = state
1912 .schema_jobs
1913 .values()
1914 .find(|job| job.table_id == first_desc.table_id && !job.state.is_terminal())
1915 .map(|job| job.job_id);
1916 let size_of = |tablet_id: TabletId| -> Result<u64, RuntimeError> {
1917 self.tablet_sink(tablet_id)
1918 .ok_or_else(|| {
1919 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
1920 })?
1921 .lock()
1922 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
1923 .tablet_size_bytes()
1924 .map_err(RuntimeError::from)
1925 };
1926 let replica_count = first_desc.replicas.len();
1927 let raft_ids = meta
1928 .allocate_raft_node_ids(u32::try_from(replica_count).unwrap_or(u32::MAX), control)
1929 .await?;
1930 let allocation = ChildAllocation {
1931 tablet_id: TabletId::new_random(),
1932 raft_group_id: RaftGroupId::new_random(),
1933 replicas: first_desc
1934 .replicas
1935 .iter()
1936 .zip(&raft_ids)
1937 .map(|(replica, raft_node_id)| ReplicaDescriptor {
1938 node_id: replica.node_id,
1939 role: replica.role,
1940 raft_node_id: *raft_node_id,
1941 })
1942 .collect(),
1943 };
1944 let plan = MergePlanner::new(self.node_data.clone()).plan(
1945 MergeInputs {
1946 first: first_desc,
1947 second: second_desc,
1948 first_schema: schema.schema_version,
1949 second_schema: schema.schema_version,
1950 active_schema_job,
1951 first_size_bytes: size_of(first)?,
1952 second_size_bytes: size_of(second)?,
1953 max_merged_size_bytes: DEFAULT_MAX_MERGED_SIZE_BYTES,
1954 },
1955 now_timestamp(),
1956 allocation,
1957 )?;
1958 Ok(plan)
1959 }
1960
1961 async fn replacement_sink(
1965 &mut self,
1966 plan: &MergePlan,
1967 control: &ExecutionControl,
1968 ) -> Result<RuntimeChildSink, RuntimeError> {
1969 let descriptor = plan.replacement_descriptor();
1970 if descriptor.replica_on(self.identity.node_id).is_some()
1971 && !self.tablets.contains_key(&descriptor.tablet_id)
1972 {
1973 let voters: Vec<(NodeId, String)> = descriptor
1974 .replicas
1975 .iter()
1976 .filter_map(|replica| {
1977 self.peers
1978 .get(&replica.node_id)
1979 .map(|address| (replica.node_id, address.clone()))
1980 })
1981 .collect();
1982 self.create_hosted_replica(&descriptor, Some(voters.as_slice()))
1983 .await?;
1984 }
1985 let group = self
1986 .tablets
1987 .get(&descriptor.tablet_id)
1988 .ok_or_else(|| {
1989 RuntimeError::InvalidRequest(format!(
1990 "replacement tablet {} is not hosted on this node",
1991 descriptor.tablet_id
1992 ))
1993 })?
1994 .group
1995 .clone();
1996 Ok(RuntimeChildSink {
1997 group,
1998 handle: tokio::runtime::Handle::current(),
1999 control: control.clone(),
2000 staged: None,
2001 })
2002 }
2003
2004 pub async fn sync_hosted_tablets(
2014 &mut self,
2015 _control: &ExecutionControl,
2016 ) -> Result<HostedSyncReport, RuntimeError> {
2017 let meta = self.meta_group_or_err()?;
2018 let state = meta.state();
2019 let node_id = self.identity.node_id;
2020 let mut report = HostedSyncReport::default();
2021 for record in state.tablets.values() {
2023 let descriptor = &record.descriptor;
2024 if self.tablets.contains_key(&descriptor.tablet_id)
2025 || descriptor.replica_on(node_id).is_none()
2026 {
2027 continue;
2028 }
2029 self.create_hosted_replica(descriptor, None).await?;
2030 report.created.push(descriptor.tablet_id);
2031 }
2032 for record in state.tablets.values() {
2035 let published = &record.descriptor;
2036 let Some(tablet) = self.tablets.get(&published.tablet_id) else {
2037 continue;
2038 };
2039 if published.generation <= tablet.descriptor.generation {
2040 continue;
2041 }
2042 if published.raft_group_id != tablet.descriptor.raft_group_id {
2043 return Err(RuntimeError::InvalidRequest(format!(
2044 "meta descriptor for tablet {} names a different raft group than the \
2045 hosted replica",
2046 published.tablet_id
2047 )));
2048 }
2049 if published.replica_on(node_id).is_none() {
2050 continue; }
2052 let layout = TabletLayout::new(
2053 self.node_data.clone(),
2054 published.tablet_id,
2055 published.raft_group_id,
2056 );
2057 layout.store_metadata(published)?;
2058 self.tablets
2059 .get_mut(&published.tablet_id)
2060 .expect("checked above")
2061 .descriptor = published.clone();
2062 report.refreshed.push(published.tablet_id);
2063 }
2064 let outgoing: Vec<TabletId> = self
2067 .tablets
2068 .keys()
2069 .filter(|tablet_id| {
2070 state
2071 .tablets
2072 .get(tablet_id)
2073 .is_none_or(|record| record.descriptor.replica_on(node_id).is_none())
2074 })
2075 .copied()
2076 .collect();
2077 for tablet_id in outgoing {
2078 let Some(tablet) = self.tablets.remove(&tablet_id) else {
2079 continue;
2080 };
2081 let layout = TabletLayout::new(
2082 self.node_data.clone(),
2083 tablet_id,
2084 tablet.descriptor.raft_group_id,
2085 );
2086 tablet.group.shutdown().await?;
2087 drop(tablet);
2088 layout.teardown()?;
2089 report.torn_down.push(tablet_id);
2090 }
2091 Ok(report)
2092 }
2093
2094 pub async fn write_tablet_rows(
2103 &self,
2104 tablet_id: TabletId,
2105 entries: &[(Key, Vec<u8>)],
2106 control: &ExecutionControl,
2107 ) -> Result<GroupCommitReceipt, RuntimeError> {
2108 let tablet = self.hosted_tablet(tablet_id)?;
2109 let envelope = CommandEnvelope::new(
2110 COMMAND_TYPE_TABLET_DATA,
2111 new_data_command_id()?,
2112 TabletDataCommandRecord::new(TabletDataCommand::Upsert {
2113 entries: entries
2114 .iter()
2115 .map(|(key, value)| (key.as_bytes().to_vec(), value.clone()))
2116 .collect(),
2117 })
2118 .encode(),
2119 );
2120 let receipt = tablet
2121 .group
2122 .propose(CommandKind::Catalog, envelope, control)
2123 .await?;
2124 Ok(receipt)
2125 }
2126
2127 pub async fn write_tablet_ops(
2132 &self,
2133 tablet_id: TabletId,
2134 operations: Vec<TabletWriteOperation>,
2135 control: &ExecutionControl,
2136 ) -> Result<GroupCommitReceipt, RuntimeError> {
2137 let tablet = self.hosted_tablet(tablet_id)?;
2138 let envelope = build_tablet_write_envelope(new_data_command_id()?, operations);
2139 let receipt = tablet
2140 .group
2141 .propose(CommandKind::Catalog, envelope, control)
2142 .await?;
2143 Ok(receipt)
2144 }
2145
2146 pub fn bind_tablet_user_table(
2148 &self,
2149 tablet_id: TabletId,
2150 binding: TabletTableBinding,
2151 ) -> Result<TabletTableBinding, RuntimeError> {
2152 let sink = self.tablet_sink(tablet_id).ok_or_else(|| {
2153 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
2154 })?;
2155 let mut locked = sink
2156 .lock()
2157 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?;
2158 Ok(locked.bind_tablet_user_table(binding)?)
2159 }
2160
2161 pub fn tablet_table_binding(&self, tablet_id: TabletId) -> Option<TabletTableBinding> {
2163 let sink = self.tablet_sink(tablet_id)?;
2164 let locked = sink.lock().ok()?;
2165 locked.tablet_table_binding().cloned()
2166 }
2167
2168 pub fn tablet_typed_rows(
2170 &self,
2171 tablet_id: TabletId,
2172 ) -> Result<mongreldb_consensus::engine_sink::TypedTabletRows, RuntimeError> {
2173 let sink = self.tablet_sink(tablet_id).ok_or_else(|| {
2174 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
2175 })?;
2176 let locked = sink
2177 .lock()
2178 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?;
2179 Ok(locked.tablet_typed_rows()?)
2180 }
2181
2182 pub fn tablet_rows(&self, tablet_id: TabletId) -> Result<BTreeMap<Key, Vec<u8>>, RuntimeError> {
2188 let sink = self.tablet_sink(tablet_id).ok_or_else(|| {
2189 RuntimeError::InvalidRequest(format!("this node hosts no tablet {tablet_id}"))
2190 })?;
2191 let rows = sink
2192 .lock()
2193 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".into()))?
2194 .tablet_rows()?;
2195 Ok(rows
2196 .into_iter()
2197 .map(|(key, value)| (Key::from_bytes(key), value))
2198 .collect())
2199 }
2200
2201 pub fn status(&self) -> RuntimeStatus {
2204 RuntimeStatus {
2205 identity: self.identity.clone(),
2206 rpc_address: self.rpc_address.clone(),
2207 meta: self.meta.as_ref().map(|meta| MetaGroupStatus {
2208 meta_group_id: meta.meta_group_id(),
2209 metadata_version: meta.metadata_version(),
2210 metrics: meta.group().metrics(),
2211 }),
2212 tablets: self
2213 .tablets
2214 .values()
2215 .map(|tablet| TabletGroupStatus {
2216 tablet_id: tablet.descriptor.tablet_id,
2217 raft_group_id: tablet.descriptor.raft_group_id,
2218 state: tablet.descriptor.state,
2219 replicas: tablet.descriptor.replicas.clone(),
2220 applied: tablet.group.applied_position(),
2221 metrics: tablet.group.metrics(),
2222 })
2223 .collect(),
2224 }
2225 }
2226
2227 pub async fn shutdown(mut self) -> Result<(), RuntimeError> {
2231 if let Some(server) = self.server.take() {
2234 server.shutdown().await;
2235 }
2236 let mut first_error: Option<RuntimeError> = None;
2237 for (_, tablet) in std::mem::take(&mut self.tablets) {
2240 if let Err(error) = tablet.group.shutdown().await {
2241 if first_error.is_none() {
2242 first_error = Some(error.into());
2243 }
2244 }
2245 let close_result = tablet
2246 .sink
2247 .lock()
2248 .map_err(|_| RuntimeError::InvalidRequest("engine sink lock poisoned".to_owned()))
2249 .and_then(|mut sink| sink.close().map_err(RuntimeError::from));
2250 if let Err(error) = close_result {
2251 if first_error.is_none() {
2252 first_error = Some(error);
2253 }
2254 }
2255 }
2256 if let Some(meta) = self.meta.take() {
2257 if let Err(error) = meta.shutdown().await {
2258 if first_error.is_none() {
2259 first_error = Some(error.into());
2260 }
2261 }
2262 }
2263 match first_error {
2265 Some(error) => Err(error),
2266 None => Ok(()),
2267 }
2268 }
2269
2270 pub async fn crash(mut self) {
2276 drop(self.server.take());
2278 for (_, tablet) in std::mem::take(&mut self.tablets) {
2279 let TabletGroup { group, sink, .. } = tablet;
2280 match Arc::try_unwrap(group) {
2281 Ok(group) => group.crash().await,
2282 Err(group) => {
2283 let _ = group.shutdown().await;
2284 }
2285 }
2286 if let Ok(mut sink) = sink.lock() {
2287 let _ = sink.crash();
2288 };
2289 }
2290 if let Some(meta) = self.meta.take() {
2291 match Arc::try_unwrap(meta) {
2292 Ok(meta) => meta.crash().await,
2293 Err(meta) => {
2294 let _ = meta.shutdown().await;
2295 }
2296 }
2297 }
2298 }
2299}
2300
2301pub const DEFAULT_MAX_MERGED_SIZE_BYTES: u64 = 64 * 1024 * 1024;
2309
2310fn now_timestamp() -> HlcTimestamp {
2314 let micros = std::time::SystemTime::now()
2315 .duration_since(std::time::UNIX_EPOCH)
2316 .map(|duration| duration.as_micros() as u64)
2317 .unwrap_or(0);
2318 HlcTimestamp {
2319 physical_micros: micros,
2320 logical: 0,
2321 node_tiebreaker: 0,
2322 }
2323}
2324
2325fn new_data_command_id() -> Result<[u8; 16], RuntimeError> {
2327 let mut id = [0u8; 16];
2328 getrandom::getrandom(&mut id)
2329 .map_err(|error| RuntimeError::InvalidRequest(format!("CSPRNG failed: {error}")))?;
2330 Ok(id)
2331}
2332
2333struct RuntimeMetaPlane {
2341 meta: Arc<MetaGroup<TcpTransport>>,
2342 handle: tokio::runtime::Handle,
2343 control: ExecutionControl,
2344}
2345
2346impl RuntimeMetaPlane {
2347 fn propose(&self, command: MetaCommand) -> Result<(), MetaRejectionReason> {
2348 let meta = self.meta.clone();
2349 let control = self.control.clone();
2350 self.handle
2351 .block_on(async move {
2352 meta.propose(crate::meta::new_command_id()?, command, &control)
2353 .await
2354 })
2355 .map(|_| ())
2356 .map_err(|error| match error {
2357 MetaError::Rejected(reason) => reason,
2358 MetaError::Consensus(ConsensusError::NotLeader { leader }) => {
2359 MetaRejectionReason::NotLeader { leader }
2360 }
2361 other => MetaRejectionReason::ProposalFailed {
2362 reason: other.to_string(),
2363 },
2364 })
2365 }
2366}
2367
2368impl TabletMetaPlane for RuntimeMetaPlane {
2369 fn set_tablet(&mut self, descriptor: &TabletDescriptor) -> Result<(), MetaRejectionReason> {
2370 self.propose(MetaCommand::SetTabletDescriptor {
2371 descriptor: descriptor.clone(),
2372 })
2373 }
2374
2375 fn tablet(&self, tablet_id: TabletId) -> Option<TabletDescriptor> {
2376 self.meta.state().tablet(tablet_id).cloned()
2377 }
2378
2379 fn remove_tablet(
2380 &mut self,
2381 tablet_id: TabletId,
2382 generation: u64,
2383 ) -> Result<(), MetaRejectionReason> {
2384 self.propose(MetaCommand::RemoveTabletDescriptor {
2385 tablet_id,
2386 generation,
2387 })
2388 }
2389
2390 fn publish_split(&mut self, command: &SplitPublishCommand) -> Result<(), MetaRejectionReason> {
2391 self.propose(MetaCommand::PublishSplit {
2392 command: command.clone(),
2393 })
2394 }
2395}
2396
2397impl MergeMetaPlane for RuntimeMetaPlane {
2398 fn publish_merge(&mut self, command: &MergePublishCommand) -> Result<(), MetaRejectionReason> {
2399 self.propose(MetaCommand::PublishMerge {
2400 command: command.clone(),
2401 })
2402 }
2403}
2404
2405struct RuntimeKeyspace {
2407 sink: Option<Arc<Mutex<EngineApplySink>>>,
2408 pinned: Option<(HlcTimestamp, u64)>,
2409}
2410
2411struct RuntimeSnapshotPin {
2412 ts: HlcTimestamp,
2413 _engine: EngineTabletPin,
2414}
2415
2416impl SnapshotPin for RuntimeSnapshotPin {
2417 fn pinned_at(&self) -> HlcTimestamp {
2418 self.ts
2419 }
2420}
2421
2422impl TabletKeyspace for RuntimeKeyspace {
2423 fn pin_snapshot(&mut self, ts: HlcTimestamp) -> Result<Box<dyn SnapshotPin>, TabletDataError> {
2424 let pin = self
2425 .sink
2426 .as_ref()
2427 .ok_or_else(|| TabletDataError::Keyspace("tablet engine is not open".to_owned()))?
2428 .lock()
2429 .map_err(|_| TabletDataError::Keyspace("engine sink lock poisoned".to_owned()))?
2430 .pin_tablet_snapshot(ts)
2431 .map_err(|error| TabletDataError::Keyspace(error.to_string()))?;
2432 self.pinned = Some((ts, pin.epoch()));
2433 Ok(Box::new(RuntimeSnapshotPin { ts, _engine: pin }))
2434 }
2435
2436 fn snapshot_at(
2437 &self,
2438 ts: HlcTimestamp,
2439 ) -> Result<crate::split::RecordStream<'_>, TabletDataError> {
2440 let (_, epoch) = self
2441 .pinned
2442 .filter(|(pinned, _)| *pinned == ts)
2443 .ok_or_else(|| {
2444 TabletDataError::Keyspace(format!("tablet snapshot {ts:?} is not pinned"))
2445 })?;
2446 let rows = self
2447 .sink
2448 .as_ref()
2449 .ok_or_else(|| TabletDataError::Keyspace("tablet engine is not open".to_owned()))?
2450 .lock()
2451 .map_err(|_| TabletDataError::Keyspace("engine sink lock poisoned".to_owned()))?
2452 .tablet_rows_at_epoch(epoch)
2453 .map_err(|error| TabletDataError::Keyspace(error.to_string()))?;
2454 Ok(Box::new(
2455 rows.into_iter()
2456 .map(|(key, value)| (Key::from_bytes(key), value)),
2457 ))
2458 }
2459
2460 fn deltas_after(
2461 &self,
2462 ts: HlcTimestamp,
2463 ) -> Result<crate::split::MutationStream<'_>, TabletDataError> {
2464 let (_, epoch) = self
2465 .pinned
2466 .filter(|(pinned, _)| *pinned == ts)
2467 .ok_or_else(|| {
2468 TabletDataError::Keyspace(format!("tablet snapshot {ts:?} is not pinned"))
2469 })?;
2470 let deltas = self
2471 .sink
2472 .as_ref()
2473 .ok_or_else(|| TabletDataError::Keyspace("tablet engine is not open".to_owned()))?
2474 .lock()
2475 .map_err(|_| TabletDataError::Keyspace("engine sink lock poisoned".to_owned()))?
2476 .tablet_deltas_after_epoch(epoch)
2477 .map_err(|error| TabletDataError::Keyspace(error.to_string()))?;
2478 Ok(Box::new(deltas.into_iter().map(
2479 |mutation| match mutation {
2480 TabletDataMutation::Upsert(key, value) => {
2481 TabletMutation::Upsert(Key::from_bytes(key), value)
2482 }
2483 TabletDataMutation::Delete(key) => TabletMutation::Delete(Key::from_bytes(key)),
2484 },
2485 )))
2486 }
2487}
2488
2489struct RuntimeChildSink {
2492 group: Arc<ConsensusGroup<TcpTransport>>,
2493 handle: tokio::runtime::Handle,
2494 control: ExecutionControl,
2495 staged: Option<BTreeMap<Key, Vec<u8>>>,
2496}
2497
2498impl RuntimeChildSink {
2499 fn propose_data(&self, command: TabletDataCommand) -> Result<(), TabletDataError> {
2500 let envelope = CommandEnvelope::new(
2501 COMMAND_TYPE_TABLET_DATA,
2502 new_data_command_id().map_err(|error| TabletDataError::Sink(error.to_string()))?,
2503 TabletDataCommandRecord::new(command).encode(),
2504 );
2505 self.handle
2506 .block_on(
2507 self.group
2508 .propose(CommandKind::Catalog, envelope, &self.control),
2509 )
2510 .map(|_| ())
2511 .map_err(|error| TabletDataError::Sink(error.to_string()))
2512 }
2513}
2514
2515impl ChildStateSink for RuntimeChildSink {
2516 fn begin_build(&mut self) -> Result<(), TabletDataError> {
2517 self.staged = Some(BTreeMap::new());
2518 Ok(())
2519 }
2520
2521 fn stage(&mut self, key: &Key, value: &[u8]) -> Result<(), TabletDataError> {
2522 let staged = self.staged.as_mut().ok_or(TabletDataError::NoStagedBuild)?;
2523 staged.insert(key.clone(), value.to_vec());
2524 Ok(())
2525 }
2526
2527 fn install_staged(&mut self) -> Result<(), TabletDataError> {
2528 let rows = self.staged.take().ok_or(TabletDataError::NoStagedBuild)?;
2529 self.propose_data(TabletDataCommand::Replace {
2530 rows: rows
2531 .into_iter()
2532 .map(|(key, value)| (key.into_bytes(), value))
2533 .collect(),
2534 })
2535 }
2536
2537 fn apply_delta(&mut self, mutation: &TabletMutation) -> Result<(), TabletDataError> {
2538 match mutation {
2539 TabletMutation::Upsert(key, value) => self.propose_data(TabletDataCommand::Upsert {
2540 entries: vec![(key.as_bytes().to_vec(), value.clone())],
2541 }),
2542 TabletMutation::Delete(key) => self.propose_data(TabletDataCommand::Delete {
2543 keys: vec![key.as_bytes().to_vec()],
2544 }),
2545 }
2546 }
2547}
2548
2549#[derive(Clone, Debug, Default, PartialEq, Eq)]
2551pub struct HostedSyncReport {
2552 pub created: Vec<TabletId>,
2554 pub refreshed: Vec<TabletId>,
2557 pub torn_down: Vec<TabletId>,
2560}
2561
2562#[cfg(test)]
2563mod tests {
2564 use super::*;
2565
2566 #[test]
2567 fn layout_scan_tolerates_a_node_without_tablets() {
2568 let tmp = tempfile::tempdir().unwrap();
2569 assert_eq!(scan_tablet_layouts(tmp.path()).unwrap(), Vec::new());
2570 std::fs::create_dir_all(tmp.path().join(TABLETS_DIR)).unwrap();
2572 assert_eq!(scan_tablet_layouts(tmp.path()).unwrap(), Vec::new());
2573 }
2574
2575 #[test]
2576 fn layout_scan_fails_closed_on_garbage_directories() {
2577 let tmp = tempfile::tempdir().unwrap();
2578 let garbage = tmp.path().join(TABLETS_DIR).join("not-a-tablet");
2580 std::fs::create_dir_all(&garbage).unwrap();
2581 assert!(matches!(
2582 scan_tablet_layouts(tmp.path()),
2583 Err(RuntimeError::InvalidRequest(_))
2584 ));
2585 std::fs::remove_dir_all(&garbage).unwrap();
2586 let tablet_id = TabletId::new_random();
2589 std::fs::create_dir_all(tmp.path().join(TABLETS_DIR).join(tablet_id.to_hex())).unwrap();
2590 assert!(matches!(
2591 scan_tablet_layouts(tmp.path()),
2592 Err(RuntimeError::Tablet(TabletError::MissingMetadata(_)))
2593 ));
2594 }
2595
2596 fn ts(micros: u64) -> HlcTimestamp {
2599 HlcTimestamp {
2600 physical_micros: micros,
2601 logical: 0,
2602 node_tiebreaker: 0,
2603 }
2604 }
2605
2606 fn pos(index: u64) -> LogPosition {
2607 LogPosition { term: 1, index }
2608 }
2609
2610 fn upsert(key: &[u8], value: &[u8]) -> TabletDataCommand {
2611 TabletDataCommand::Upsert {
2612 entries: vec![(key.to_vec(), value.to_vec())],
2613 }
2614 }
2615
2616 #[test]
2617 fn tablet_ledger_applies_versions_and_partitions_the_timeline() {
2618 let tmp = tempfile::tempdir().unwrap();
2619 let mut ledger = TabletLedger::open(tmp.path()).unwrap();
2620 ledger
2621 .apply(&upsert(b"a", b"a@1"), ts(100), pos(1))
2622 .unwrap();
2623 ledger
2624 .apply(&upsert(b"b", b"b@1"), ts(100), pos(2))
2625 .unwrap();
2626 ledger.pin(ts(100));
2629 ledger
2630 .apply(&upsert(b"a", b"a@2"), ts(200), pos(3))
2631 .unwrap();
2632 assert_eq!(
2634 ledger.rows_at(ts(100)),
2635 BTreeMap::from([
2636 (Key::from_bytes(b"a".to_vec()), b"a@1".to_vec()),
2637 (Key::from_bytes(b"b".to_vec()), b"b@1".to_vec()),
2638 ])
2639 );
2640 assert_eq!(
2641 ledger.current_rows().get(&Key::from_bytes(b"a".to_vec())),
2642 Some(&b"a@2".to_vec())
2643 );
2644 assert_eq!(
2646 ledger.deltas_after(ts(100)),
2647 vec![(Key::from_bytes(b"a".to_vec()), b"a@2".to_vec())]
2648 );
2649 assert!(ledger.deltas_after(ts(200)).is_empty());
2650 ledger.unpin(ts(100));
2651 ledger
2653 .apply(&upsert(b"z", b"z@1"), ts(300), pos(3))
2654 .unwrap();
2655 assert!(!ledger
2656 .current_rows()
2657 .contains_key(&Key::from_bytes(b"z".to_vec())));
2658 let position = ledger.applied_position();
2660 drop(ledger);
2661 let reopened = TabletLedger::open(tmp.path()).unwrap();
2662 assert_eq!(reopened.applied_position(), position);
2663 assert_eq!(
2664 reopened.current_rows().get(&Key::from_bytes(b"a".to_vec())),
2665 Some(&b"a@2".to_vec())
2666 );
2667 std::fs::write(
2669 tmp.path()
2670 .join("raft")
2671 .join("state")
2672 .join(TABLET_LEDGER_FILENAME),
2673 b"junk",
2674 )
2675 .unwrap();
2676 assert!(TabletLedger::open(tmp.path()).is_err());
2677 }
2678
2679 #[test]
2680 fn tablet_ledger_compacts_against_the_oldest_pin() {
2681 let tmp = tempfile::tempdir().unwrap();
2682 let mut ledger = TabletLedger::open(tmp.path()).unwrap();
2683 ledger
2684 .apply(&upsert(b"a", b"a@1"), ts(100), pos(1))
2685 .unwrap();
2686 ledger.pin(ts(150));
2689 ledger
2690 .apply(&upsert(b"a", b"a@2"), ts(200), pos(2))
2691 .unwrap();
2692 ledger
2693 .apply(&upsert(b"a", b"a@3"), ts(300), pos(3))
2694 .unwrap();
2695 assert_eq!(
2696 ledger.rows_at(ts(150)).get(&Key::from_bytes(b"a".to_vec())),
2697 Some(&b"a@1".to_vec())
2698 );
2699 assert_eq!(
2700 ledger.deltas_after(ts(150)),
2701 vec![
2702 (Key::from_bytes(b"a".to_vec()), b"a@2".to_vec()),
2703 (Key::from_bytes(b"a".to_vec()), b"a@3".to_vec()),
2704 ]
2705 );
2706 ledger.unpin(ts(150));
2708 ledger
2709 .apply(&upsert(b"a", b"a@4"), ts(400), pos(4))
2710 .unwrap();
2711 assert_eq!(
2712 ledger.rows_at(ts(150)).get(&Key::from_bytes(b"a".to_vec())),
2713 None
2714 );
2715 assert_eq!(
2716 ledger.current_rows().get(&Key::from_bytes(b"a".to_vec())),
2717 Some(&b"a@4".to_vec())
2718 );
2719 assert_eq!(ledger.pin_count(), 0);
2720 }
2721
2722 #[test]
2723 fn tablet_ledger_replace_is_atomic_and_refused_under_a_live_pin() {
2724 let tmp = tempfile::tempdir().unwrap();
2725 let mut ledger = TabletLedger::open(tmp.path()).unwrap();
2726 ledger
2727 .apply(&upsert(b"a", b"a@1"), ts(100), pos(1))
2728 .unwrap();
2729 ledger.pin(ts(150));
2730 let replace = TabletDataCommand::Replace {
2731 rows: vec![(b"b".to_vec(), b"b@2".to_vec())],
2732 };
2733 assert!(ledger.apply(&replace, ts(200), pos(2)).is_err());
2734 ledger.unpin(ts(150));
2735 ledger.apply(&replace, ts(200), pos(2)).unwrap();
2736 assert_eq!(
2737 ledger.current_rows(),
2738 BTreeMap::from([(Key::from_bytes(b"b".to_vec()), b"b@2".to_vec())])
2739 );
2740 let bytes = ledger.snapshot_bytes().unwrap();
2742 let follower_dir = tempfile::tempdir().unwrap();
2743 let mut follower = TabletLedger::open(follower_dir.path()).unwrap();
2744 follower.install_bytes(&bytes).unwrap();
2745 assert_eq!(follower.current_rows(), ledger.current_rows());
2746 assert!(follower.install_bytes(b"junk").is_err());
2747 }
2748
2749 #[test]
2750 fn group_snapshot_frame_round_trips_and_fails_closed() {
2751 let engine = b"engine-half".to_vec();
2752 let ledger = b"ledger-half".to_vec();
2753 let frame = encode_group_snapshot(&engine, &ledger);
2754 let (engine_back, ledger_back) = decode_group_snapshot(&frame).unwrap();
2755 assert_eq!(engine_back, engine);
2756 assert_eq!(ledger_back, ledger);
2757 let (empty_engine, empty_ledger) =
2759 decode_group_snapshot(&encode_group_snapshot(&[], &[])).unwrap();
2760 assert!(empty_engine.is_empty() && empty_ledger.is_empty());
2761 assert!(decode_group_snapshot(&frame[..6]).is_err());
2765 assert!(decode_group_snapshot(&frame[..12 + engine.len() - 1]).is_err());
2766 let mut future = frame.clone();
2767 future[..4].copy_from_slice(&99_u32.to_le_bytes());
2768 assert!(decode_group_snapshot(&future).is_err());
2769 }
2770
2771 #[test]
2772 fn tablet_data_command_record_round_trips_and_fails_closed() {
2773 let record = TabletDataCommandRecord::new(upsert(b"k", b"v"));
2774 let decoded = TabletDataCommandRecord::decode(&record.encode()).unwrap();
2775 assert_eq!(decoded, record);
2776 assert!(TabletDataCommandRecord::decode(b"not json").is_err());
2777 let mut value: serde_json::Value = serde_json::from_slice(&record.encode()).unwrap();
2778 value["format_version"] = serde_json::json!(99);
2779 assert!(TabletDataCommandRecord::decode(&serde_json::to_vec(&value).unwrap()).is_err());
2780 }
2781}