#[cfg(test)]
mod tests {
use super::{
FormatMarker, ShardSwapState, configured_shards, prepare_shard_swap,
read_collection_marker, read_shard_version, remove_collection_marker, swap_prepared_shard,
write_collection_marker, write_shard_version,
};
use crate::DbError;
use std::fs;
use tempfile::tempdir;
#[test]
fn format_marker_round_trips() {
let dir = tempdir().unwrap();
write_shard_version(dir.path(), 7).unwrap();
assert_eq!(read_shard_version(dir.path(), 1).unwrap(), 7);
}
#[test]
fn format_marker_publish_failure_preserves_the_durable_marker() {
let dir = tempdir().unwrap();
write_shard_version(dir.path(), 1).unwrap();
fs::create_dir(dir.path().join("format.v.tmp")).unwrap();
assert!(write_shard_version(dir.path(), 2).is_err());
assert_eq!(read_shard_version(dir.path(), 9).unwrap(), 1);
}
#[test]
fn format_marker_rejects_bad_magic_truncation_and_crc_corruption() {
let dir = tempdir().unwrap();
let marker = dir.path().join("format.v");
fs::write(&marker, b"BAD!\x07\0\0\0\0\0").unwrap();
assert!(matches!(
read_shard_version(dir.path(), 1),
Err(DbError::FormatMismatch(_))
));
fs::write(&marker, b"AFV1\x07").unwrap();
assert!(matches!(
read_shard_version(dir.path(), 1),
Err(DbError::FormatMismatch(_))
));
fs::write(&marker, b"AFV1\x07\0\0\0\0\0").unwrap();
assert!(matches!(
read_shard_version(dir.path(), 1),
Err(DbError::FormatMismatch(_))
));
}
#[test]
fn format_marker_collection_scope_round_trips_and_removes() {
let dir = tempdir().unwrap();
let marker = FormatMarker { from: 1, to: 3 };
write_collection_marker(dir.path(), marker).unwrap();
assert_eq!(read_collection_marker(dir.path()).unwrap(), Some(marker));
remove_collection_marker(dir.path()).unwrap();
assert_eq!(read_collection_marker(dir.path()).unwrap(), None);
}
#[test]
fn format_swap_state_recovers_every_durable_layout() {
let root = tempdir().unwrap();
let collection = root.path();
let shard = collection.join("000");
let mig = collection.join("000.mig");
let old = collection.join("000.old");
fs::create_dir(&shard).unwrap();
write_shard_version(&shard, 3).unwrap();
fs::create_dir(&old).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::DeleteOld
);
fs::remove_dir_all(&old).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::AlreadyTarget
);
fs::create_dir(&mig).unwrap();
write_shard_version(&mig, 3).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::DeleteMig
);
fs::remove_dir_all(&mig).unwrap();
fs::remove_dir_all(&shard).unwrap();
fs::create_dir(&shard).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::RewriteFrom(1)
);
fs::create_dir(&mig).unwrap();
write_shard_version(&mig, 3).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::DeleteMigThenRewrite(1)
);
fs::remove_dir_all(&shard).unwrap();
assert_eq!(
prepare_shard_swap(collection, "000", 1, 3).unwrap(),
ShardSwapState::FinishSwap
);
}
#[test]
fn format_swap_state_target_live_discards_mig_without_rewrite() {
let root = tempdir().unwrap();
let shard = root.path().join("000");
let migration = root.path().join("000.mig");
fs::create_dir(&shard).unwrap();
fs::create_dir(&migration).unwrap();
write_shard_version(&shard, 2).unwrap();
write_shard_version(&migration, 2).unwrap();
let state = prepare_shard_swap(root.path(), "000", 1, 2).unwrap();
assert_eq!(state, ShardSwapState::DeleteMig);
}
#[test]
fn format_swap_state_rewrites_from_intermediate_version_and_rejects_outside_chain() {
let root = tempdir().unwrap();
let shard = root.path().join("000");
fs::create_dir(&shard).unwrap();
write_shard_version(&shard, 2).unwrap();
assert_eq!(
prepare_shard_swap(root.path(), "000", 1, 3).unwrap(),
ShardSwapState::RewriteFrom(2)
);
write_shard_version(&shard, 4).unwrap();
assert!(matches!(
prepare_shard_swap(root.path(), "000", 1, 3),
Err(DbError::FormatMismatch(_))
));
}
#[test]
fn format_swap_state_finishes_prepared_swap() {
let root = tempdir().unwrap();
let shard = root.path().join("000");
let migration = root.path().join("000.mig");
fs::create_dir(&shard).unwrap();
fs::create_dir(&migration).unwrap();
write_shard_version(&migration, 3).unwrap();
swap_prepared_shard(root.path(), "000").unwrap();
assert_eq!(read_shard_version(&shard, 1).unwrap(), 3);
assert!(!migration.exists());
assert!(!root.path().join("000.old").exists());
}
#[test]
fn format_swap_state_discards_a_torn_migration_marker_when_source_is_recoverable() {
let root = tempdir().unwrap();
let shard = root.path().join("000");
let migration = root.path().join("000.mig");
fs::create_dir(&shard).unwrap();
fs::create_dir(&migration).unwrap();
fs::write(migration.join("format.v"), b"AFV1\x02").unwrap();
assert_eq!(
prepare_shard_swap(root.path(), "000", 1, 2).unwrap(),
ShardSwapState::RewriteFrom(1)
);
assert!(!migration.exists());
fs::rename(&shard, root.path().join("000.old")).unwrap();
fs::create_dir(&migration).unwrap();
fs::write(migration.join("format.v"), b"AFV1\x02").unwrap();
assert_eq!(
prepare_shard_swap(root.path(), "000", 1, 2).unwrap(),
ShardSwapState::RestoreOldThenRewrite(1)
);
assert!(!migration.exists());
}
#[test]
fn configured_shard_census_rejects_missing_and_out_of_range_identities() {
let root = tempdir().unwrap();
fs::create_dir(root.path().join("shard_000")).unwrap();
assert!(matches!(
configured_shards(root.path(), 2, 1, 2),
Err(DbError::FormatMismatch(message)) if message.contains("missing shard")
));
fs::create_dir(root.path().join("shard_001")).unwrap();
fs::create_dir(root.path().join("shard_002")).unwrap();
assert!(matches!(
configured_shards(root.path(), 2, 1, 2),
Err(DbError::FormatMismatch(message)) if message.contains("out-of-range")
));
}
#[test]
fn configured_shard_census_rejects_non_directory_and_propagates_read_dir_errors() {
let root = tempdir().unwrap();
fs::write(root.path().join("shard_000"), b"corrupt shard identity").unwrap();
assert!(matches!(
configured_shards(root.path(), 1, 1, 2),
Err(DbError::FormatMismatch(message)) if message.contains("not a directory")
));
let not_a_directory = root.path().join("ordinary-file");
fs::write(¬_a_directory, b"not a directory").unwrap();
assert!(matches!(
configured_shards(¬_a_directory, 1, 1, 2),
Err(DbError::Io(_))
));
}
}
use std::fs;
use std::io::Write;
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::io::direct;
use crate::{DbError, DbResult};
const FORMAT_MARKER: &str = "format.v";
const FORMAT_MARKER_TMP: &str = "format.v.tmp";
const MIGRATION_MARKER: &str = ".migration";
const MIGRATION_MARKER_TMP: &str = ".migration.tmp";
const FORMAT_MAGIC: [u8; 4] = *b"AFV1";
const FORMAT_MARKER_LEN: usize = 10;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct FormatMarker {
pub(crate) from: u16,
pub(crate) to: u16,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ShardSwapState {
RewriteFrom(u16),
AlreadyTarget,
DeleteOld,
DeleteMig,
DeleteOldThenRewrite(u16),
DeleteMigThenRewrite(u16),
RestoreOldThenRewrite(u16),
FinishSwap,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ConfiguredShard {
pub(crate) id: usize,
pub(crate) name: String,
pub(crate) state: ShardSwapState,
}
fn format_error(message: impl Into<String>) -> DbError {
DbError::FormatMismatch(message.into())
}
fn validate_chain(version: u16, source: u16, target: u16, location: &Path) -> DbResult<()> {
if source <= version && version <= target {
Ok(())
} else {
Err(format_error(format!(
"format version {version} in {} is outside declared chain {source}..={target}",
location.display()
)))
}
}
pub(crate) fn read_shard_version(dir: &Path, fallback: u16) -> DbResult<u16> {
let path = dir.join(FORMAT_MARKER);
let bytes = match fs::read(&path) {
Ok(bytes) => bytes,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(fallback),
Err(err) => return Err(err.into()),
};
if bytes.len() != FORMAT_MARKER_LEN {
return Err(format_error(format!(
"invalid format marker length in {}",
path.display()
)));
}
if bytes[..4] != FORMAT_MAGIC {
return Err(format_error(format!(
"invalid format marker magic in {}",
path.display()
)));
}
let expected = u32::from_le_bytes(bytes[6..10].try_into().expect("marker length checked"));
let actual = crc32fast::hash(&bytes[..6]);
if actual != expected {
return Err(format_error(format!(
"invalid format marker CRC in {}",
path.display()
)));
}
Ok(u16::from_le_bytes(
bytes[4..6].try_into().expect("marker length checked"),
))
}
pub(crate) fn write_shard_version(dir: &Path, version: u16) -> DbResult<()> {
let path = dir.join(FORMAT_MARKER);
let temp = dir.join(FORMAT_MARKER_TMP);
let mut bytes = [0_u8; FORMAT_MARKER_LEN];
bytes[..4].copy_from_slice(&FORMAT_MAGIC);
bytes[4..6].copy_from_slice(&version.to_le_bytes());
let crc = crc32fast::hash(&bytes[..6]).to_le_bytes();
bytes[6..].copy_from_slice(&crc);
let mut file = fs::File::create(&temp)?;
file.write_all(&bytes)?;
file.sync_all()?;
fs::rename(temp, path)?;
direct::sync_dir(dir)
}
pub(crate) fn write_collection_marker(collection_dir: &Path, marker: FormatMarker) -> DbResult<()> {
let marker_path = collection_dir.join(MIGRATION_MARKER);
let temp_path = collection_dir.join(MIGRATION_MARKER_TMP);
let bytes = serde_json::to_vec(&marker).map_err(|err| format_error(err.to_string()))?;
let mut file = fs::File::create(&temp_path)?;
file.write_all(&bytes)?;
file.sync_all()?;
fs::rename(temp_path, marker_path)?;
direct::sync_dir(collection_dir)
}
pub(crate) fn read_collection_marker(collection_dir: &Path) -> DbResult<Option<FormatMarker>> {
let path = collection_dir.join(MIGRATION_MARKER);
match fs::read(path) {
Ok(bytes) => serde_json::from_slice(&bytes)
.map(Some)
.map_err(|err| format_error(err.to_string())),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(err) => Err(err.into()),
}
}
pub(crate) fn remove_collection_marker(collection_dir: &Path) -> DbResult<()> {
let path = collection_dir.join(MIGRATION_MARKER);
match fs::remove_file(path) {
Ok(()) => direct::sync_dir(collection_dir),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(err) => Err(err.into()),
}
}
pub(crate) fn prepare_shard_swap(
collection_dir: &Path,
shard_name: &str,
source_version: u16,
target_version: u16,
) -> DbResult<ShardSwapState> {
if source_version > target_version {
return Err(format_error(format!(
"invalid format migration chain {source_version}..={target_version}"
)));
}
let shard = collection_dir.join(shard_name);
let migration = collection_dir.join(format!("{shard_name}.mig"));
let old = collection_dir.join(format!("{shard_name}.old"));
let shard_version = directory_version(&shard, source_version)?;
let old_version = directory_version(&old, source_version)?;
let migration_version = match directory_version(&migration, source_version) {
Ok(version) => version,
Err(error) if shard_version.is_some() || old_version.is_some() => {
fs::remove_dir_all(&migration)?;
direct::sync_dir(collection_dir)?;
tracing::warn!(
path = ?migration,
%error,
"discarded format output with an invalid marker"
);
None
}
Err(error) => return Err(error),
};
for (version, path) in [
(shard_version, &shard),
(migration_version, &migration),
(old_version, &old),
] {
if let Some(version) = version {
validate_chain(version, source_version, target_version, path)?;
}
}
match (shard_version, migration_version, old_version) {
(Some(version), _, Some(_)) if version == target_version => Ok(ShardSwapState::DeleteOld),
(Some(version), _, Some(_)) => Ok(ShardSwapState::DeleteOldThenRewrite(version)),
(Some(version), Some(_), None) if version == target_version => {
Ok(ShardSwapState::DeleteMig)
}
(Some(version), Some(_), None) => Ok(ShardSwapState::DeleteMigThenRewrite(version)),
(Some(version), None, None) if version == target_version => {
Ok(ShardSwapState::AlreadyTarget)
}
(Some(version), None, None) => Ok(ShardSwapState::RewriteFrom(version)),
(None, Some(version), _) if version == target_version => Ok(ShardSwapState::FinishSwap),
(None, Some(_), Some(old_version)) => {
Ok(ShardSwapState::RestoreOldThenRewrite(old_version))
}
(None, None, Some(old_version)) => Ok(ShardSwapState::RestoreOldThenRewrite(old_version)),
(None, Some(version), None) => Err(format_error(format!(
"cannot finish swap from non-target format version {version} in {}",
migration.display()
))),
(None, None, None) => Err(format_error(format!(
"missing shard and migration output for {}",
shard.display()
))),
}
}
fn remove_known_temp_file(path: &Path) -> DbResult<()> {
match fs::symlink_metadata(path) {
Ok(metadata) if metadata.file_type().is_file() => {
fs::remove_file(path)?;
if let Some(parent) = path.parent() {
direct::sync_dir(parent)?;
}
Ok(())
}
Ok(_) => Err(format_error(format!(
"format migration temp identity {} is not a regular file",
path.display()
))),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
}
}
fn remove_known_staging_dir(path: &Path) -> DbResult<()> {
match fs::symlink_metadata(path) {
Ok(metadata) if metadata.file_type().is_dir() => {
fs::remove_dir_all(path)?;
if let Some(parent) = path.parent() {
direct::sync_dir(parent)?;
}
Ok(())
}
Ok(_) => Err(format_error(format!(
"format migration staging identity {} is not a directory",
path.display()
))),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
}
}
pub(crate) fn cleanup_stale_format_temps(
collection_dir: &Path,
shard_count: usize,
) -> DbResult<()> {
match fs::symlink_metadata(collection_dir) {
Ok(metadata) if metadata.file_type().is_dir() => {}
Ok(_) => {
return Err(format_error(format!(
"format migration collection identity {} is not a directory",
collection_dir.display()
)));
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(error) => return Err(error.into()),
}
remove_known_temp_file(&collection_dir.join(MIGRATION_MARKER_TMP))?;
for shard in 0..shard_count {
remove_known_staging_dir(&collection_dir.join(format!(".format-stage-{shard:03}")))?;
for suffix in ["", ".mig", ".old"] {
let dir = collection_dir.join(format!("shard_{shard:03}{suffix}"));
match fs::symlink_metadata(&dir) {
Ok(metadata) if metadata.file_type().is_dir() => {}
Ok(_) => {
return Err(format_error(format!(
"format migration shard identity {} is not a directory",
dir.display()
)));
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
Err(error) => return Err(error.into()),
}
for name in [FORMAT_MARKER_TMP, "format.gsn.tmp", "format.file-id.tmp"] {
remove_known_temp_file(&dir.join(name))?;
}
}
}
Ok(())
}
pub(crate) fn configured_shards(
collection_dir: &Path,
shard_count: usize,
source_version: u16,
target_version: u16,
) -> DbResult<Vec<ConfiguredShard>> {
for entry in fs::read_dir(collection_dir)? {
let entry = entry?;
let name = entry.file_name();
let name = name.to_str().ok_or_else(|| {
format_error(format!(
"non-UTF-8 entry in format migration root {}",
collection_dir.display()
))
})?;
let Some(identity) = name.strip_prefix("shard_") else {
continue;
};
let (digits, suffix) = match identity.split_once('.') {
Some((digits, suffix)) => (digits, Some(suffix)),
None => (identity, None),
};
if suffix.is_some_and(|suffix| suffix != "mig" && suffix != "old") {
return Err(format_error(format!(
"unexpected shard identity `{name}` in {}",
collection_dir.display()
)));
}
let id = digits.parse::<usize>().map_err(|_| {
format_error(format!(
"unexpected shard identity `{name}` in {}",
collection_dir.display()
))
})?;
let canonical = match suffix {
Some(suffix) => format!("shard_{id:03}.{suffix}"),
None => format!("shard_{id:03}"),
};
if name != canonical || id >= shard_count {
return Err(format_error(format!(
"unexpected or out-of-range shard identity `{name}` for configured shard_count {shard_count}"
)));
}
if !entry.file_type()?.is_dir() {
return Err(format_error(format!(
"shard identity {} is not a directory",
entry.path().display()
)));
}
}
(0..shard_count)
.map(|id| {
let name = format!("shard_{id:03}");
let state = prepare_shard_swap(collection_dir, &name, source_version, target_version)?;
Ok(ConfiguredShard { id, name, state })
})
.collect()
}
fn directory_version(path: &Path, fallback: u16) -> DbResult<Option<u16>> {
match fs::symlink_metadata(path) {
Ok(metadata) if metadata.file_type().is_dir() => {
read_shard_version(path, fallback).map(Some)
}
Ok(_) => Err(format_error(format!(
"format migration shard variant {} is not a directory",
path.display()
))),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(error) => Err(error.into()),
}
}
pub(crate) fn swap_prepared_shard(collection_dir: &Path, shard_name: &str) -> DbResult<()> {
let shard = collection_dir.join(shard_name);
let migration = collection_dir.join(format!("{shard_name}.mig"));
let old = collection_dir.join(format!("{shard_name}.old"));
if directory_version(&shard, 0)?.is_some() {
hit(2);
fs::rename(&shard, &old)?;
direct::sync_dir(collection_dir)?;
hit(3);
}
hit(4);
fs::rename(&migration, &shard)?;
direct::sync_dir(collection_dir)?;
hit(5);
if directory_version(&old, 0)?.is_some() {
hit(6);
fs::remove_dir_all(old)?;
direct::sync_dir(collection_dir)?;
hit(7);
}
Ok(())
}
#[cfg(any(test, feature = "test-faults"))]
fn hit(stage: u8) {
crate::format_migration::failpoints::hit(stage);
}
#[cfg(not(any(test, feature = "test-faults")))]
fn hit(_: u8) {}