use std::path::{Path, PathBuf};
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use mongreldb_log::commit_log::LogPosition;
use mongreldb_types::hlc::HlcTimestamp;
use mongreldb_types::ids::{ClusterId, DatabaseId, NodeId, RaftGroupId};
use crate::storage_mode::StorageMode;
use crate::{MongrelError, Result};
pub const COMMAND_TYPE_CATALOG_COMMAND: u32 = 2;
pub const COMMAND_TYPE_MAINTENANCE: u32 = 3;
pub const REPLICATED_TXN_FORMAT_VERSION: u16 = 1;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReplicatedTxnPayload {
pub version: u16,
pub records: Vec<crate::wal::Record>,
}
impl ReplicatedTxnPayload {
pub fn new(records: Vec<crate::wal::Record>) -> Self {
Self {
version: REPLICATED_TXN_FORMAT_VERSION,
records,
}
}
pub fn encode(&self) -> Result<Vec<u8>> {
Ok(bincode::serialize(self)?)
}
pub fn decode(bytes: &[u8]) -> Result<Self> {
let payload: Self = bincode::deserialize(bytes)?;
if payload.version != REPLICATED_TXN_FORMAT_VERSION {
return Err(MongrelError::UnsupportedStorageVersion {
component: "replicated transaction payload",
found: payload.version,
supported: REPLICATED_TXN_FORMAT_VERSION,
});
}
Ok(payload)
}
}
pub const ENGINE_SNAPSHOT_FORMAT_VERSION: u16 = 1;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineSnapshotFile {
pub path: PathBuf,
pub sha256: [u8; 32],
pub data: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineSnapshotTable {
pub table_id: u64,
pub name: String,
pub visible_rows: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineSnapshot {
pub version: u16,
pub group_id: RaftGroupId,
pub last_included: LogPosition,
pub commit_ts: Option<HlcTimestamp>,
pub epoch: u64,
pub cluster_id: ClusterId,
pub database_id: DatabaseId,
pub catalog_version: u64,
pub tables: Vec<EngineSnapshotTable>,
pub files: Vec<EngineSnapshotFile>,
pub wal_format: u16,
pub storage_mode_format: u16,
}
impl EngineSnapshot {
pub fn capture(
db: &crate::Database,
group_id: RaftGroupId,
last_included: LogPosition,
commit_ts: Option<HlcTimestamp>,
) -> Result<Self> {
let mode = db.storage_mode()?.ok_or_else(|| {
MongrelError::Other("engine snapshot capture before marker write".into())
})?;
let (cluster_id, _node_id, database_id) = mode.cluster_identity().ok_or_else(|| {
MongrelError::InvalidArgument(format!(
"engine snapshots capture cluster replicas, got mode {mode:?}"
))
})?;
let epoch = db.visible_epoch();
let snapshot = crate::epoch::Snapshot::at(epoch);
let mut tables = Vec::new();
for name in db.table_names() {
let table_id = db.table_id(&name)?;
let handle = db.table(&name)?;
let visible_rows = handle.lock().visible_rows(snapshot)?.len() as u64;
tables.push(EngineSnapshotTable {
table_id,
name,
visible_rows,
});
}
tables.sort_by_key(|table| table.table_id);
let files = crate::replication::capture_files(db.root())?
.into_iter()
.map(|file| {
let sha256: [u8; 32] = Sha256::digest(&file.data).into();
Ok(EngineSnapshotFile {
path: file.path,
sha256,
data: file.data,
})
})
.collect::<Result<Vec<_>>>()?;
Ok(Self {
version: ENGINE_SNAPSHOT_FORMAT_VERSION,
group_id,
last_included,
commit_ts,
epoch: epoch.0,
cluster_id,
database_id,
catalog_version: db.catalog_version(),
tables,
files,
wal_format: crate::wal::WAL_VERSION,
storage_mode_format: crate::storage_mode::STORAGE_MODE_FORMAT_VERSION,
})
}
pub fn encode(&self) -> Result<Vec<u8>> {
Ok(bincode::serialize(self)?)
}
pub fn decode(bytes: &[u8]) -> Result<Self> {
let snapshot: Self = bincode::deserialize(bytes)?;
if snapshot.version != ENGINE_SNAPSHOT_FORMAT_VERSION {
return Err(MongrelError::UnsupportedStorageVersion {
component: "engine snapshot",
found: snapshot.version,
supported: ENGINE_SNAPSHOT_FORMAT_VERSION,
});
}
Ok(snapshot)
}
pub fn validate(
&self,
group_id: &RaftGroupId,
cluster_id: &ClusterId,
database_id: &DatabaseId,
) -> Result<()> {
if &self.group_id != group_id {
return Err(MongrelError::InvalidArgument(format!(
"engine snapshot group {:?} does not match this replica's group {:?}",
self.group_id, group_id
)));
}
if &self.cluster_id != cluster_id || &self.database_id != database_id {
return Err(MongrelError::InvalidArgument(
"engine snapshot database identity does not match this replica".into(),
));
}
if self.wal_format != crate::wal::WAL_VERSION {
return Err(MongrelError::UnsupportedStorageVersion {
component: "wal",
found: self.wal_format,
supported: crate::wal::WAL_VERSION,
});
}
if self.storage_mode_format != crate::storage_mode::STORAGE_MODE_FORMAT_VERSION {
return Err(MongrelError::UnsupportedStorageVersion {
component: "storage-mode marker",
found: self.storage_mode_format,
supported: crate::storage_mode::STORAGE_MODE_FORMAT_VERSION,
});
}
let mut seen = std::collections::HashSet::new();
for file in &self.files {
crate::replication::validate_relative_path(&file.path)?;
if !seen.insert(file.path.clone()) {
return Err(MongrelError::InvalidArgument(format!(
"duplicate engine snapshot path {:?}",
file.path
)));
}
let digest: [u8; 32] = Sha256::digest(&file.data).into();
if digest != file.sha256 {
return Err(MongrelError::Other(format!(
"engine snapshot file {:?} failed its content hash",
file.path
)));
}
}
Ok(())
}
pub fn stage_into(&self, staging: &Path, node_id: NodeId) -> Result<()> {
let marker_relative =
Path::new(crate::database::META_DIR).join(crate::storage_mode::STORAGE_MODE_FILENAME);
for file in &self.files {
crate::replication::validate_relative_path(&file.path)?;
if file.path == marker_relative {
continue;
}
let path = staging.join(&file.path);
let parent = path.parent().expect("validated file has parent");
crate::durable_file::create_directory_all(parent)?;
let mut output = std::fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(&path)?;
std::io::Write::write_all(&mut output, &file.data)?;
output.sync_all()?;
crate::durable_file::sync_directory(parent)?;
}
for table in &self.tables {
crate::durable_file::create_directory_all(
&staging
.join(crate::database::TABLES_DIR)
.join(table.table_id.to_string())
.join("_runs"),
)?;
}
if !staging.join(crate::catalog::CATALOG_FILENAME).is_file() {
return Err(MongrelError::InvalidArgument(
"engine snapshot has no CATALOG".into(),
));
}
let durable_stage = crate::durable_file::DurableRoot::open(staging)?;
crate::storage_mode::rewrite(
&durable_stage,
&StorageMode::ClusterReplica {
cluster_id: self.cluster_id,
node_id,
database_id: self.database_id,
},
)?;
crate::durable_file::sync_directory(staging)?;
Ok(())
}
pub fn validate_staged(&self, staging: &Path) -> Result<()> {
let db = crate::Database::open_offline_validation(staging)?;
let snapshot = crate::epoch::Snapshot::at(crate::epoch::Epoch(self.epoch));
for table in &self.tables {
let handle = db.table(&table.name).map_err(|error| {
MongrelError::Other(format!(
"staged engine snapshot is missing table {:?}: {error}",
table.name
))
})?;
let rows = handle.lock().visible_rows(snapshot)?.len() as u64;
if rows != table.visible_rows {
return Err(MongrelError::Other(format!(
"staged engine snapshot table {:?} has {rows} visible rows at epoch {}, expected {}",
table.name, self.epoch, table.visible_rows
)));
}
}
Ok(())
}
pub fn install(self, live: &mut Option<Arc<crate::Database>>, node_id: NodeId) -> Result<()> {
let db = live.as_ref().ok_or_else(|| {
MongrelError::Other("engine snapshot install without a live database".into())
})?;
if Arc::strong_count(db) > 1 {
return Err(MongrelError::Conflict(
"engine snapshot install refused over live state: database is busy".into(),
));
}
let destination = db.root().to_path_buf();
match db.storage_mode()? {
Some(StorageMode::ClusterReplica {
cluster_id,
database_id,
..
}) if cluster_id == self.cluster_id && database_id == self.database_id => {}
other => {
return Err(MongrelError::InvalidArgument(format!(
"refusing to install an engine snapshot over storage mode {other:?}"
)));
}
}
let parent = destination
.parent()
.filter(|path| !path.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."))
.to_path_buf();
let name = destination
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| MongrelError::InvalidArgument("invalid database root".into()))?;
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
let stage = parent.join(format!(
".{name}.engine-stage-{}-{nonce}",
std::process::id()
));
let backup = parent.join(format!(".{name}.engine-old-{}-{nonce}", std::process::id()));
if stage.exists() || backup.exists() {
return Err(MongrelError::Conflict(
"engine snapshot staging path already exists".into(),
));
}
let result = (|| -> Result<()> {
crate::durable_file::create_directory(&stage)?;
self.stage_into(&stage, node_id)?;
self.validate_staged(&stage)?;
let owned = live.take().expect("checked above");
owned.shutdown()?;
if let Err(failure) = crate::replication::rename_entry(&destination, &backup) {
return Err(failure.error.into());
}
if let Err(failure) = crate::replication::rename_entry(&stage, &destination) {
if let Err(rollback) = crate::replication::rename_entry(&backup, &destination) {
return Err(crate::replication::uncertain_install_error(
self.epoch,
&failure.error,
&rollback.error,
));
}
return Err(failure.error.into());
}
crate::replication::remove_directory(&backup)?;
Ok(())
})();
if let Err(error) = result {
if stage.exists() {
let _ = crate::replication::remove_directory(&stage);
}
return Err(error);
}
let expected = StorageMode::ClusterReplica {
cluster_id: self.cluster_id,
node_id,
database_id: self.database_id,
};
let reopened = crate::Database::open_cluster_replica(&destination, &expected)?;
*live = Some(Arc::new(reopened));
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn txn_payload_round_trip_and_version_gate() {
let payload = ReplicatedTxnPayload::new(vec![crate::wal::Record::new(
crate::epoch::Epoch(1),
7,
crate::wal::Op::TxnCommit {
epoch: 1,
added_runs: Vec::new(),
},
)]);
let bytes = payload.encode().unwrap();
let decoded = ReplicatedTxnPayload::decode(&bytes).unwrap();
assert_eq!(decoded.records.len(), 1);
let mut corrupt = payload.clone();
corrupt.version = REPLICATED_TXN_FORMAT_VERSION + 1;
let bytes = bincode::serialize(&corrupt).unwrap();
assert!(matches!(
ReplicatedTxnPayload::decode(&bytes),
Err(MongrelError::UnsupportedStorageVersion { .. })
));
}
#[test]
fn validate_rejects_bad_hash_and_foreign_identity() {
let snapshot = EngineSnapshot {
version: ENGINE_SNAPSHOT_FORMAT_VERSION,
group_id: RaftGroupId::from_bytes([1; 16]),
last_included: LogPosition { term: 1, index: 2 },
commit_ts: None,
epoch: 2,
cluster_id: ClusterId::from_bytes([2; 16]),
database_id: DatabaseId::from_bytes([3; 16]),
catalog_version: 0,
tables: Vec::new(),
files: vec![EngineSnapshotFile {
path: PathBuf::from("CATALOG"),
sha256: [9; 32],
data: b"catalog".to_vec(),
}],
wal_format: crate::wal::WAL_VERSION,
storage_mode_format: crate::storage_mode::STORAGE_MODE_FORMAT_VERSION,
};
let group = RaftGroupId::from_bytes([1; 16]);
let cluster = ClusterId::from_bytes([2; 16]);
let database = DatabaseId::from_bytes([3; 16]);
assert!(snapshot.validate(&group, &cluster, &database).is_err());
let wrong_group = RaftGroupId::from_bytes([9; 16]);
assert!(snapshot
.validate(&wrong_group, &cluster, &database)
.is_err());
let wrong_database = DatabaseId::from_bytes([9; 16]);
assert!(snapshot
.validate(&group, &cluster, &wrong_database)
.is_err());
let mut good = snapshot.clone();
good.files[0].sha256 = Sha256::digest(&good.files[0].data).into();
good.validate(&group, &cluster, &database).unwrap();
}
}