use std::path::Path;
use std::sync::Arc;
use crate::disk_loc::DiskLoc;
use crate::entry::entry_size;
use crate::error::{DbError, DbResult};
use crate::fixed::config::FixedConfig;
use crate::fixed::shard::{RECOVERY_CHUNK_BYTES, RecoverMode};
use crate::key::Location;
use crate::sync::{self, MutexGuard};
pub trait DurabilityInner: Send {
type Loc: Location;
fn write_new(&mut self, shard_id: u8, key: &[u8], value: &[u8]) -> DbResult<Self::Loc>;
fn write_update(
&mut self,
shard_id: u8,
old_loc: Self::Loc,
key: &[u8],
value: &[u8],
) -> DbResult<Self::Loc>;
fn write_tombstone(&mut self, shard_id: u8, old_loc: Self::Loc, key: &[u8]) -> DbResult<()>;
fn write_discard(&mut self, loc: Self::Loc, key: &[u8]) -> DbResult<()>;
fn should_sync(&self) -> bool;
fn sync(&mut self) -> DbResult<()>;
}
pub trait Durability: Send + Sync + Sized {
type Loc: Location;
type Inner: DurabilityInner<Loc = Self::Loc>;
fn shard_count(&self) -> usize;
fn lock_shard(&self, shard_id: usize) -> MutexGuard<'_, Self::Inner>;
fn shard_prefix_bits(&self) -> usize;
fn flush(&self) -> DbResult<()>;
fn close(&self) -> DbResult<()>;
}
#[allow(dead_code)]
pub struct Bitcask {
pub(crate) engine: crate::engine::Engine,
pub(crate) compaction_threshold: f64,
}
impl DurabilityInner for crate::shard::ShardInner {
type Loc = DiskLoc;
fn write_new(&mut self, shard_id: u8, key: &[u8], value: &[u8]) -> DbResult<DiskLoc> {
let (loc, _gsn) = self.append_entry(shard_id, key, value, false)?;
Ok(loc)
}
fn write_update(
&mut self,
shard_id: u8,
old_loc: DiskLoc,
key: &[u8],
value: &[u8],
) -> DbResult<DiskLoc> {
let (new_loc, _gsn) = self.append_entry(shard_id, key, value, false)?;
self.add_dead_bytes(old_loc.file_id, entry_size(key.len(), old_loc.len));
Ok(new_loc)
}
fn write_tombstone(&mut self, shard_id: u8, old_loc: DiskLoc, key: &[u8]) -> DbResult<()> {
let (_loc, _gsn) = self.append_entry(shard_id, key, &[], true)?;
self.add_dead_bytes(old_loc.file_id, entry_size(key.len(), old_loc.len));
Ok(())
}
fn write_discard(&mut self, _loc: DiskLoc, _key: &[u8]) -> DbResult<()> {
Ok(())
}
fn should_sync(&self) -> bool {
false
}
fn sync(&mut self) -> DbResult<()> {
Ok(())
}
}
impl Durability for Bitcask {
type Loc = DiskLoc;
type Inner = crate::shard::ShardInner;
fn shard_count(&self) -> usize {
self.engine.shards().len()
}
fn lock_shard(&self, shard_id: usize) -> MutexGuard<'_, crate::shard::ShardInner> {
self.engine.shards()[shard_id].lock()
}
fn shard_prefix_bits(&self) -> usize {
self.engine.config().shard_prefix_bits
}
fn flush(&self) -> DbResult<()> {
self.engine.flush()
}
fn close(&self) -> DbResult<()> {
self.engine.flush()
}
}
#[allow(dead_code)]
pub struct Fixed {
pub(crate) engine: Arc<crate::fixed::engine::FixedEngine>,
}
impl Fixed {
pub fn open(
path: impl AsRef<Path>,
config: FixedConfig,
key_len: usize,
value_len: usize,
) -> DbResult<Self> {
if key_len > u16::MAX as usize {
return Err(DbError::Config("key_len exceeds u16::MAX"));
}
if value_len > u16::MAX as usize {
return Err(DbError::Config("value_len exceeds u16::MAX"));
}
let engine = crate::fixed::engine::FixedEngine::open(
path,
config,
key_len as u16,
value_len as u16,
)?;
Ok(Self {
engine: Arc::new(engine),
})
}
pub(crate) fn recover_entries<F, V>(&self, make_visitor: F) -> DbResult<u32>
where
F: Fn(usize) -> V + Sync,
V: FnMut(&[u8], &[u8], u32),
{
let shards = self.engine.shards();
let make_visitor = &make_visitor;
let max_parallel = crate::config::default_shard_count().max(1);
let mut results: Vec<DbResult<u32>> = Vec::with_capacity(shards.len());
for batch_start in (0..shards.len()).step_by(max_parallel) {
let batch_end = (batch_start + max_parallel).min(shards.len());
let batch: Vec<DbResult<u32>> = std::thread::scope(|s| {
let handles: Vec<_> = (batch_start..batch_end)
.map(|shard_idx| {
let shard = &shards[shard_idx];
s.spawn(move || Self::recover_one_shard(shard, shard_idx, make_visitor))
})
.collect();
handles
.into_iter()
.map(|h| match h.join() {
Ok(r) => r,
Err(_) => Err(DbError::RecoveryPanic),
})
.collect()
});
results.extend(batch);
}
let mut total = 0u32;
for r in results {
total += r?;
}
Ok(total)
}
fn recover_one_shard<F, V>(
shard: &crate::fixed::engine::FixedShard,
shard_idx: usize,
make_visitor: &F,
) -> DbResult<u32>
where
F: Fn(usize) -> V,
V: FnMut(&[u8], &[u8], u32),
{
let start = std::time::Instant::now();
let mut inner = sync::lock(&shard.inner);
let use_sidecar = if inner.has_clean_shutdown() {
match inner.load_versions_sidecar() {
Ok(()) => true,
Err(e) => {
tracing::warn!(
shard = shard_idx, error = %e,
"fixed.versions sidecar invalid; falling back to full scan"
);
let _ = std::fs::remove_file(inner.dir().join("fixed.versions"));
false
}
}
} else {
false
};
let mode = if use_sidecar {
RecoverMode::Clean
} else {
RecoverMode::Dirty
};
let mut visitor = make_visitor(shard_idx);
let recovered = inner.recover_slots(mode, RECOVERY_CHUNK_BYTES, |k, v, slot_id| {
visitor(k, v, slot_id)
})?;
inner.clear_clean_shutdown()?;
let elapsed = start.elapsed().as_secs_f64();
metrics::histogram!(
"armdb.recovery.duration_seconds",
"path" => if use_sidecar { "fixed_sidecar" } else { "fixed_scan" }
)
.record(elapsed);
Ok(recovered)
}
}
impl DurabilityInner for crate::fixed::shard::FixedShardInner {
type Loc = u32;
fn write_new(&mut self, _shard_id: u8, key: &[u8], value: &[u8]) -> DbResult<u32> {
let slot = self.alloc_slot()?;
let _ = self.write_slot(slot, key, value)?;
Ok(slot)
}
fn write_update(
&mut self,
_shard_id: u8,
old_loc: u32,
key: &[u8],
value: &[u8],
) -> DbResult<u32> {
let _ = self.write_slot(old_loc, key, value)?;
Ok(old_loc)
}
fn write_tombstone(&mut self, _shard_id: u8, old_loc: u32, key: &[u8]) -> DbResult<()> {
self.delete_slot(old_loc, key)?;
self.bitmap.clear(old_loc);
Ok(())
}
fn write_discard(&mut self, loc: u32, key: &[u8]) -> DbResult<()> {
self.delete_slot(loc, key)?;
self.bitmap.clear(loc);
Ok(())
}
fn should_sync(&self) -> bool {
self.should_sync()
}
fn sync(&mut self) -> DbResult<()> {
self.sync()
}
}
impl Durability for Fixed {
type Loc = u32;
type Inner = crate::fixed::shard::FixedShardInner;
fn shard_count(&self) -> usize {
self.engine.shards().len()
}
fn lock_shard(&self, shard_id: usize) -> MutexGuard<'_, crate::fixed::shard::FixedShardInner> {
sync::lock(&self.engine.shards()[shard_id].inner)
}
fn shard_prefix_bits(&self) -> usize {
self.engine.config().shard_prefix_bits
}
fn flush(&self) -> DbResult<()> {
self.engine.flush()
}
fn close(&self) -> DbResult<()> {
self.engine.close()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fixed::config::FixedConfig;
use crate::fixed::slot::{STATUS_DELETED, STATUS_FREE, STATUS_OCCUPIED, status_of, version_of};
use tempfile::tempdir;
fn test_fixed_config() -> FixedConfig {
FixedConfig {
shard_count: 1,
grow_step: 16,
..FixedConfig::test()
}
}
#[test]
fn test_recover_populates_versions_dirty() {
let dir = tempdir().unwrap();
let path = dir.path().join("db");
{
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
{
let mut inner = sync::lock(&fixed.engine.shards()[0].inner);
let id0 = inner.alloc_slot().unwrap();
inner
.write_slot(id0, b"keyaaaaa", b"val_0000_0000_00")
.unwrap();
let id1 = inner.alloc_slot().unwrap();
inner
.write_slot(id1, b"keybbbbb", b"val_0001_0000_00")
.unwrap();
inner.delete_slot(id1, b"keybbbbb").unwrap();
inner.sync().unwrap();
}
drop(fixed);
}
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
let entries = std::sync::Mutex::new(Vec::new());
let entries_ref = &entries;
fixed
.recover_entries(move |_shard_idx| {
move |k: &[u8], v: &[u8], slot_id: u32| {
entries_ref
.lock()
.unwrap()
.push((k.to_vec(), v.to_vec(), slot_id));
}
})
.unwrap();
let entries = entries.into_inner().unwrap();
assert_eq!(entries.len(), 1, "only one OCCUPIED entry");
assert_eq!(entries[0].0, b"keyaaaaa");
let inner = sync::lock(&fixed.engine.shards()[0].inner);
assert_eq!(status_of(inner.versions[0]), STATUS_OCCUPIED);
assert_eq!(status_of(inner.versions[1]), STATUS_DELETED);
}
#[test]
fn test_recover_populates_versions_clean() {
let dir = tempdir().unwrap();
let path = dir.path().join("db");
{
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
{
let mut inner = sync::lock(&fixed.engine.shards()[0].inner);
let id = inner.alloc_slot().unwrap();
inner
.write_slot(id, b"keyaaaaa", b"val_0000_0000_00")
.unwrap();
}
fixed.engine.close().unwrap();
}
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
let entries = std::sync::Mutex::new(Vec::new());
let entries_ref = &entries;
fixed
.recover_entries(move |_shard_idx| {
move |k: &[u8], _v: &[u8], _id: u32| {
entries_ref.lock().unwrap().push(k.to_vec());
}
})
.unwrap();
let entries = entries.into_inner().unwrap();
assert_eq!(entries.len(), 1);
let inner = sync::lock(&fixed.engine.shards()[0].inner);
assert_eq!(status_of(inner.versions[0]), STATUS_OCCUPIED);
assert_eq!(version_of(inner.versions[0]), 1);
}
#[test]
fn test_torn_slot_preserves_version() {
let dir = tempdir().unwrap();
let path = dir.path().join("db");
{
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
{
let mut inner = sync::lock(&fixed.engine.shards()[0].inner);
let id = inner.alloc_slot().unwrap();
inner
.write_slot(id, b"keyaaaaa", b"val_0000_0000_00")
.unwrap();
inner.sync().unwrap();
}
drop(fixed);
}
let shard_data = path.join("shard_000").join("fixed.data");
use std::os::unix::fs::FileExt;
let f = std::fs::OpenOptions::new()
.write(true)
.open(&shard_data)
.unwrap();
f.write_all_at(&[0xFFu8; 16], 4112).unwrap();
f.sync_data().unwrap();
drop(f);
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
let entries = std::sync::Mutex::new(Vec::new());
let entries_ref = &entries;
fixed
.recover_entries(move |_shard_idx| {
move |k: &[u8], _v: &[u8], _id: u32| {
entries_ref.lock().unwrap().push(k.to_vec());
}
})
.unwrap();
let entries = entries.into_inner().unwrap();
assert!(entries.is_empty(), "torn slot must be skipped");
let inner = sync::lock(&fixed.engine.shards()[0].inner);
assert_eq!(status_of(inner.versions[0]), STATUS_FREE);
assert_eq!(
version_of(inner.versions[0]),
1,
"version must be preserved so next bump continues monotonically"
);
}
#[test]
fn fixed_open_rejects_oversized_key_len() {
let dir = tempfile::tempdir().unwrap();
let result = Fixed::open(dir.path(), FixedConfig::test(), 65536, 16);
assert!(result.is_err());
let msg = result.err().unwrap().to_string();
assert!(
msg.contains("key_len"),
"expected key_len error, got: {msg}"
);
}
#[test]
fn fixed_open_rejects_oversized_value_len() {
let dir = tempfile::tempdir().unwrap();
let result = Fixed::open(dir.path(), FixedConfig::test(), 8, 65536);
assert!(result.is_err());
let msg = result.err().unwrap().to_string();
assert!(
msg.contains("value_len"),
"expected value_len error, got: {msg}"
);
}
#[cfg(feature = "replication")]
#[test]
fn fixed_write_discard_uses_real_key_for_replication_delete() {
use crate::durability::DurabilityInner;
use crate::fixed_replication::FixedReplicationEvent;
let dir = tempfile::tempdir().unwrap();
let fixed = Fixed::open(dir.path(), FixedConfig::test(), 8, 16).unwrap();
let mut inner = fixed.lock_shard(0);
let (producer, mut consumer) = rtrb::RingBuffer::new(4);
inner.replication_tx = Some(producer);
inner.write_discard(0, b"key_0001").unwrap();
match consumer.pop().unwrap() {
FixedReplicationEvent::Delete { key, .. } => assert_eq!(key, b"key_0001"),
_ => panic!("expected delete event"),
}
}
#[test]
fn clean_recovery_skips_merged_gap_slot() {
use crate::fixed::slot::pack_meta;
let dir = tempdir().unwrap();
let path = dir.path().join("db");
{
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
{
let mut inner = sync::lock(&fixed.engine.shards()[0].inner);
for i in 0..3u64 {
let id = inner.alloc_slot().unwrap();
assert_eq!(id, i as u32);
let mut val = [0u8; 16];
val[..8].copy_from_slice(&(i * 100).to_be_bytes());
inner.write_slot(id, &i.to_be_bytes(), &val).unwrap();
}
inner.versions[1] = pack_meta(STATUS_FREE, version_of(inner.versions[1]));
inner.bitmap.clear(1);
}
fixed.engine.close().unwrap();
}
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
let seen = std::sync::Mutex::new(Vec::new());
let seen_ref = &seen;
fixed
.recover_entries(move |_shard_idx| {
move |k: &[u8], _v: &[u8], _id: u32| seen_ref.lock().unwrap().push(k.to_vec())
})
.unwrap();
let mut seen = seen.into_inner().unwrap();
seen.sort();
assert_eq!(
seen,
vec![0u64.to_be_bytes().to_vec(), 2u64.to_be_bytes().to_vec()]
);
let inner = sync::lock(&fixed.engine.shards()[0].inner);
assert_eq!(
status_of(inner.versions[1]),
STATUS_FREE,
"gap slot untouched"
);
assert!(!inner.bitmap.is_set(1), "gap slot not added to bitmap");
}
#[test]
fn clean_recovery_crc_fail_frees_slot() {
use std::os::unix::fs::FileExt;
let dir = tempdir().unwrap();
let path = dir.path().join("db");
{
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
{
let mut inner = sync::lock(&fixed.engine.shards()[0].inner);
for i in 0..3u64 {
let id = inner.alloc_slot().unwrap();
inner.write_slot(id, &i.to_be_bytes(), &[9u8; 16]).unwrap();
}
}
fixed.engine.close().unwrap(); }
let data = path.join("shard_000").join("fixed.data");
let ss = crate::fixed::slot::slot_size(8, 16) as u64;
let off = crate::fixed::shard::HEADER_SIZE + ss + 8 + 8;
let f = std::fs::OpenOptions::new().write(true).open(&data).unwrap();
f.write_all_at(&[0xFFu8; 8], off).unwrap();
f.sync_data().unwrap();
drop(f);
let fixed = Fixed::open(&path, test_fixed_config(), 8, 16).unwrap();
let seen = std::sync::Mutex::new(Vec::new());
let seen_ref = &seen;
fixed
.recover_entries(move |_shard_idx| {
move |k: &[u8], _v: &[u8], _id: u32| seen_ref.lock().unwrap().push(k.to_vec())
})
.unwrap();
let mut seen = seen.into_inner().unwrap();
seen.sort();
assert_eq!(
seen,
vec![0u64.to_be_bytes().to_vec(), 2u64.to_be_bytes().to_vec()]
);
let inner = sync::lock(&fixed.engine.shards()[0].inner);
assert_eq!(
status_of(inner.versions[1]),
STATUS_FREE,
"CRC-fail slot -> FREE"
);
assert!(!inner.bitmap.is_set(1));
}
#[test]
fn recover_entries_factory_called_once_per_shard() {
use std::sync::atomic::{AtomicUsize, Ordering};
let dir = tempdir().unwrap();
let path = dir.path().join("db");
let cfg = FixedConfig {
shard_count: 2,
grow_step: 16,
..FixedConfig::test()
};
{
let fixed = Fixed::open(&path, cfg.clone(), 8, 16).unwrap();
{
let mut s0 = fixed.lock_shard(0);
for i in 0..2u64 {
let id = s0.alloc_slot().unwrap();
s0.write_slot(id, &(1000 + i).to_be_bytes(), &[1u8; 16])
.unwrap();
}
}
{
let mut s1 = fixed.lock_shard(1);
for i in 0..3u64 {
let id = s1.alloc_slot().unwrap();
s1.write_slot(id, &(2000 + i).to_be_bytes(), &[2u8; 16])
.unwrap();
}
}
fixed.engine.close().unwrap();
}
let fixed = Fixed::open(&path, cfg, 8, 16).unwrap();
let calls = AtomicUsize::new(0);
let per_shard: Vec<std::sync::Mutex<Vec<Vec<u8>>>> =
(0..2).map(|_| std::sync::Mutex::new(Vec::new())).collect();
let calls_ref = &calls;
let per_shard_ref = &per_shard;
fixed
.recover_entries(move |shard_idx| {
calls_ref.fetch_add(1, Ordering::Relaxed);
move |k: &[u8], _v: &[u8], _id: u32| {
per_shard_ref[shard_idx].lock().unwrap().push(k.to_vec());
}
})
.unwrap();
assert_eq!(
calls.load(Ordering::Relaxed),
2,
"factory called once per shard"
);
let mut s0: Vec<u64> = per_shard[0]
.lock()
.unwrap()
.iter()
.map(|k| u64::from_be_bytes(k[..8].try_into().unwrap()))
.collect();
let mut s1: Vec<u64> = per_shard[1]
.lock()
.unwrap()
.iter()
.map(|k| u64::from_be_bytes(k[..8].try_into().unwrap()))
.collect();
s0.sort_unstable();
s1.sort_unstable();
assert_eq!(
s0,
vec![1000, 1001],
"shard 0 visitor sees only shard 0 keys"
);
assert_eq!(
s1,
vec![2000, 2001, 2002],
"shard 1 visitor sees only shard 1 keys"
);
}
}