use crate::Key;
use crate::compaction::{CompactionIndex, compact_shard};
use crate::config::Config;
use crate::disk_loc::DiskLoc;
use crate::durability::{Bitcask, Durability, DurabilityInner};
use crate::engine::Engine;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, WriteHook};
use crate::key::Location;
use crate::recovery::recover_const_tree;
use crate::skiplist::node::{ConstNode, SkipNode, random_height};
use crate::skiplist::{InsertResult, SkipList};
use crate::sync::MutexGuard;
use std::mem::size_of;
use std::ops::Bound;
pub struct ConstTree<K: Key, const V: usize, H: WriteHook<K> = NoHook, D: Durability = Bitcask> {
index: SkipList<ConstNode<K, V, D::Loc>>,
durability: D,
shard_prefix_bits: usize,
reversed: bool,
hook: H,
}
impl<K: Key, const V: usize> ConstTree<K, V, NoHook, Bitcask> {
pub fn open(path: impl AsRef<std::path::Path>, config: Config) -> DbResult<Self> {
Self::open_hooked(path, config, NoHook)
}
}
impl<K: Key, const V: usize, H: WriteHook<K>> ConstTree<K, V, H, Bitcask> {
pub fn open_hooked(
path: impl AsRef<std::path::Path>,
config: Config,
hook: H,
) -> DbResult<Self> {
let config = config.with_resolved_hints(false);
Self::open_inner(path, config, hook)
}
fn open_inner(path: impl AsRef<std::path::Path>, config: Config, hook: H) -> DbResult<Self> {
let compaction_threshold = config.compaction_threshold;
let shard_prefix_bits = config.shard_prefix_bits;
let reversed = config.reversed;
let engine = Engine::open(path, config)?;
let durability = Bitcask {
engine,
compaction_threshold,
};
let tree = Self {
index: SkipList::new(reversed),
shard_prefix_bits,
reversed,
hook,
durability,
};
let shard_dirs = tree.durability.engine.shard_dirs();
let shard_dir_refs = Engine::shard_dir_refs(&shard_dirs);
let shard_ids = tree.durability.engine.shard_ids();
let hints = tree.durability.engine.hints();
let outcome = recover_const_tree::<K, V>(
&shard_dir_refs,
&shard_ids,
tree.index(),
hints,
#[cfg(feature = "encryption")]
tree.durability.engine.cipher(),
)?;
for tail in &outcome.active_tails {
tree.durability.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
}
for (shard_idx, dead) in outcome.shard_dead_bytes {
tree.durability.engine.shards()[shard_idx].install_dead_bytes(dead);
}
let max_gsn = outcome.max_gsn;
tree.durability
.engine
.gsn()
.fetch_max(max_gsn + 1, std::sync::atomic::Ordering::Relaxed);
if hints {
for shard in tree.durability.engine.shards().iter() {
shard.set_key_len(size_of::<K>());
}
}
tracing::info!(
key_size = size_of::<K>(),
V,
entries = tree.len(),
"const_tree recovered"
);
Ok(tree)
}
pub fn close(self) -> DbResult<()> {
if self.durability.engine.hints() {
self.sync_hints()?;
}
self.durability.engine.flush()
}
pub fn flush_buffers(&self) -> DbResult<()> {
self.durability.engine.flush_buffers()
}
pub fn config(&self) -> &Config {
self.durability.engine.config()
}
pub fn compact(&self) -> DbResult<usize> {
let mut total_compacted = 0;
for shard in self.durability.engine.shards().iter() {
total_compacted += compact_shard(shard, self, self.durability.compaction_threshold)?;
}
Ok(total_compacted)
}
pub fn sync_hints(&self) -> DbResult<()> {
for shard in self.durability.engine.shards().iter() {
shard.write_active_hint(size_of::<K>())?;
}
Ok(())
}
}
impl<K: Key, const V: usize, H: WriteHook<K>> CompactionIndex<K> for ConstTree<K, V, H, Bitcask> {
fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
let guard = self.index.collector().enter();
if let Some(node) = self.index.get(key.as_bytes(), &guard)
&& node.read_loc() == old_loc
{
node.write_loc(new_loc);
return true;
}
false
}
fn contains_key(&self, key: &K) -> bool {
self.contains(key)
}
fn is_live(&self, _shard_id: u8, key: &K, loc: DiskLoc) -> bool {
let guard = self.index.collector().enter();
self.index
.get(key.as_bytes(), &guard)
.is_some_and(|node| node.read_loc() == loc)
}
}
use crate::durability::Fixed;
use crate::fixed::config::FixedConfig;
impl<K: Key, const V: usize> ConstTree<K, V, NoHook, Fixed> {
pub fn open(path: impl AsRef<std::path::Path>, config: FixedConfig) -> DbResult<Self> {
Self::open_fixed_inner(path, config, NoHook)
}
}
impl<K: Key, const V: usize, H: WriteHook<K>> ConstTree<K, V, H, Fixed> {
pub fn open_with_hook(
path: impl AsRef<std::path::Path>,
config: FixedConfig,
hook: H,
) -> DbResult<Self> {
Self::open_fixed_inner(path, config, hook)
}
fn open_fixed_inner(
path: impl AsRef<std::path::Path>,
config: FixedConfig,
hook: H,
) -> DbResult<Self> {
let shard_prefix_bits = config.shard_prefix_bits;
let reversed = config.reversed;
let dur = Fixed::open(path, config, size_of::<K>(), V)?;
let index = SkipList::new(reversed);
let total_recovered =
dur.recover_entries(|_shard_id, key_bytes, value_bytes, slot_id| {
let key = K::from_bytes(key_bytes);
let mut value = [0u8; V];
value.copy_from_slice(value_bytes);
let height = random_height();
let node = ConstNode::<K, V, u32>::alloc(key, value, slot_id, height);
let guard = index.collector().enter();
let _ = index.insert(node, &guard);
})?;
tracing::info!(
key_size = size_of::<K>(),
V,
entries = total_recovered,
"fixed_tree recovered"
);
Ok(Self {
index,
durability: dur,
shard_prefix_bits,
reversed,
hook,
})
}
pub fn close(self) -> DbResult<()> {
self.durability.close()
}
}
#[cfg(feature = "replication")]
impl<K: Key, const V: usize, H: WriteHook<K>> ConstTree<K, V, H, Fixed> {
pub(crate) fn fixed_durability(&self) -> &Fixed {
&self.durability
}
pub fn fixed_engine_access(
&self,
) -> std::sync::Arc<dyn crate::fixed_replication::FixedEngineAccess> {
self.durability.engine.clone()
}
pub(crate) fn get_slot_id(&self, key: &K) -> Option<u32> {
let guard = self.index.collector().enter();
let node = self.index.get(key.as_bytes(), &guard)?;
Some(node.read_loc())
}
pub(crate) fn remove_key_if_slot_matches(&self, key: &K, slot_id: u32) -> bool {
let guard = self.index.collector().enter();
let node = match self.index.get(key.as_bytes(), &guard) {
Some(n) => n,
None => return false,
};
if node.read_loc() != slot_id {
return false;
}
self.index.remove(key.as_bytes(), &guard);
true
}
pub(crate) fn upsert_replicated(&self, key: &K, value: [u8; V], slot_id: u32) {
let guard = self.index.collector().enter();
if let Some(existing) = self.index.get(key.as_bytes(), &guard) {
if existing.read_loc() == slot_id {
existing.write_data(slot_id, &value);
return;
}
self.index.remove(key.as_bytes(), &guard);
}
let height = random_height();
let node_ptr = ConstNode::<K, V, u32>::alloc(*key, value, slot_id, height);
match self.index.insert(node_ptr, &guard) {
InsertResult::Inserted => {}
InsertResult::Exists(existing) => {
debug_assert!(
false,
"upsert_replicated: unexpected concurrent insert for key"
);
existing.write_data(slot_id, &value);
unsafe {
ConstNode::<K, V, u32>::dealloc_node(node_ptr);
}
}
}
}
}
impl<K: Key, const V: usize, H: WriteHook<K>, D: Durability> ConstTree<K, V, H, D> {
pub fn get(&self, key: &K) -> Option<[u8; V]> {
metrics::counter!("armdb.ops", "op" => "get", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.get");
let guard = self.index.collector().enter();
let node = self.index.get(key.as_bytes(), &guard)?;
Some(node.read_value())
}
pub fn get_or_err(&self, key: &K) -> DbResult<[u8; V]> {
self.get(key).ok_or(DbError::KeyNotFound)
}
pub fn put(&self, key: &K, value: &[u8; V]) -> DbResult<Option<[u8; V]>> {
metrics::counter!("armdb.ops", "op" => "put", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.put");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let old = self.put_locked(shard_id, &mut *inner, &guard, key, value)?;
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
self.hook
.on_write(key, old.as_ref().map(|v| &v[..]), Some(&value[..]));
Ok(old)
}
pub fn insert(&self, key: &K, value: &[u8; V]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "insert", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.insert");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
self.insert_locked(shard_id, &mut *inner, &guard, key, value)?;
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
self.hook.on_write(key, None, Some(&value[..]));
Ok(())
}
pub fn delete(&self, key: &K) -> DbResult<Option<[u8; V]>> {
metrics::counter!("armdb.ops", "op" => "delete", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.delete");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let old = self.delete_locked(shard_id, &mut *inner, &guard, key)?;
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
if let Some(ref old_val) = old {
self.hook.on_write(key, Some(&old_val[..]), None);
}
Ok(old)
}
pub fn atomic<R>(
&self,
shard_key: &K,
f: impl FnOnce(&mut ConstShard<'_, K, V, H, D>) -> DbResult<R>,
) -> DbResult<R> {
let shard_id = self.shard_for(shard_key);
let inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let mut shard = ConstShard {
tree: self,
inner,
shard_id,
guard,
events: Vec::new(),
};
let result = f(&mut shard);
let ConstShard {
inner,
guard,
events,
..
} = shard;
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
if H::NEEDS_WRITE {
for (k, old, new) in &events {
self.hook.on_write(
k,
old.as_ref().map(|v| &v[..]),
new.as_ref().map(|v| &v[..]),
);
}
}
drop(guard);
result
}
fn put_locked(
&self,
shard_id: usize,
inner: &mut D::Inner,
guard: &seize::LocalGuard<'_>,
key: &K,
value: &[u8; V],
) -> DbResult<Option<[u8; V]>> {
if let Some(existing) = self.index.get(key.as_bytes(), guard) {
let old_value = existing.read_value();
let old_loc = existing.read_loc();
let new_loc = inner.write_update(shard_id as u8, old_loc, key.as_bytes(), value)?;
existing.write_data(new_loc, value);
return Ok(Some(old_value));
}
let loc = inner.write_new(shard_id as u8, key.as_bytes(), value)?;
let height = random_height();
let node_ptr = ConstNode::<K, V, D::Loc>::alloc(*key, *value, loc, height);
match self.index.insert(node_ptr, guard) {
InsertResult::Inserted => Ok(None),
InsertResult::Exists(existing) => {
inner.write_discard(loc, key.as_bytes())?;
let old_value = existing.read_value();
let old_loc = existing.read_loc();
let new_loc = inner.write_update(shard_id as u8, old_loc, key.as_bytes(), value)?;
existing.write_data(new_loc, value);
unsafe {
ConstNode::<K, V, D::Loc>::dealloc_node(node_ptr);
}
Ok(Some(old_value))
}
}
}
fn insert_locked(
&self,
shard_id: usize,
inner: &mut D::Inner,
guard: &seize::LocalGuard<'_>,
key: &K,
value: &[u8; V],
) -> DbResult<()> {
if self.index.get(key.as_bytes(), guard).is_some() {
return Err(DbError::KeyExists);
}
let loc = inner.write_new(shard_id as u8, key.as_bytes(), value)?;
let height = random_height();
let node_ptr = ConstNode::<K, V, D::Loc>::alloc(*key, *value, loc, height);
match self.index.insert(node_ptr, guard) {
InsertResult::Inserted => Ok(()),
InsertResult::Exists(_) => {
inner.write_discard(loc, key.as_bytes())?;
unsafe {
ConstNode::<K, V, D::Loc>::dealloc_node(node_ptr);
}
Err(DbError::KeyExists)
}
}
}
fn delete_locked(
&self,
shard_id: usize,
inner: &mut D::Inner,
guard: &seize::LocalGuard<'_>,
key: &K,
) -> DbResult<Option<[u8; V]>> {
if let Some(existing) = self.index.get(key.as_bytes(), guard) {
let old_value = existing.read_value();
let old_loc = existing.read_loc();
inner.write_tombstone(shard_id as u8, old_loc, key.as_bytes())?;
self.index.remove(key.as_bytes(), guard);
return Ok(Some(old_value));
}
Ok(None)
}
fn put_no_hook(&self, key: &K, value: &[u8; V]) -> DbResult<Option<[u8; V]>> {
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let result = self.put_locked(shard_id, &mut *inner, &guard, key, value);
let needs_sync = result.is_ok() && inner.should_sync();
drop(inner);
drop(guard);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
result
}
fn delete_no_hook(&self, key: &K) -> DbResult<Option<[u8; V]>> {
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let result = self.delete_locked(shard_id, &mut *inner, &guard, key);
let needs_sync = result.is_ok() && inner.should_sync();
drop(inner);
drop(guard);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
result
}
pub fn cas(&self, key: &K, expected: &[u8; V], new_value: &[u8; V]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "cas", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.cas");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let existing = self
.index
.get(key.as_bytes(), &guard)
.ok_or(DbError::KeyNotFound)?;
let current_value = existing.read_value();
if current_value != *expected {
return Err(DbError::CasMismatch);
}
let old_loc = existing.read_loc();
let new_loc = inner.write_update(shard_id as u8, old_loc, key.as_bytes(), new_value)?;
existing.write_data(new_loc, new_value);
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
self.hook
.on_write(key, Some(&expected[..]), Some(&new_value[..]));
Ok(())
}
pub fn compare_delete(&self, key: &K, expected: &[u8; V]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "const_tree")
.increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.compare_delete");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let existing = self
.index
.get(key.as_bytes(), &guard)
.ok_or(DbError::KeyNotFound)?;
let current_value = existing.read_value();
if current_value != *expected {
return Err(DbError::CasMismatch);
}
let old_loc = existing.read_loc();
inner.write_tombstone(shard_id as u8, old_loc, key.as_bytes())?;
self.index.remove(key.as_bytes(), &guard);
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
self.hook.on_write(key, Some(¤t_value[..]), None);
Ok(())
}
pub fn update(
&self,
key: &K,
f: impl FnOnce(&[u8; V]) -> [u8; V],
) -> DbResult<Option<[u8; V]>> {
self.update_inner(key, f, false)
}
pub fn fetch_update(
&self,
key: &K,
f: impl FnOnce(&[u8; V]) -> [u8; V],
) -> DbResult<Option<[u8; V]>> {
self.update_inner(key, f, true)
}
fn update_inner(
&self,
key: &K,
f: impl FnOnce(&[u8; V]) -> [u8; V],
return_old: bool,
) -> DbResult<Option<[u8; V]>> {
metrics::counter!("armdb.ops", "op" => "update", "tree" => "const_tree").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("const_tree.update");
let shard_id = self.shard_for(key);
let mut inner = self.durability.lock_shard(shard_id);
let guard = self.index.collector().enter();
let existing = match self.index.get(key.as_bytes(), &guard) {
Some(n) => n,
None => return Ok(None),
};
let old_value = existing.read_value();
let new_value = f(&old_value);
let old_loc = existing.read_loc();
let new_loc = inner.write_update(shard_id as u8, old_loc, key.as_bytes(), &new_value)?;
existing.write_data(new_loc, &new_value);
let needs_sync = inner.should_sync();
drop(inner);
if needs_sync {
self.durability.lock_shard(shard_id).sync()?;
}
self.hook
.on_write(key, Some(&old_value[..]), Some(&new_value[..]));
Ok(Some(if return_old { old_value } else { new_value }))
}
pub fn contains(&self, key: &K) -> bool {
self.get(key).is_some()
}
pub fn first(&self) -> Option<(K, [u8; V])> {
let guard = self.index.collector().enter();
let mut ptr = crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
});
while !ptr.is_null() {
let node = unsafe { &*ptr };
if !node.is_marked() {
return Some((node.key, node.read_value()));
}
ptr = crate::skiplist::strip_mark(
node.tower(0).load(std::sync::atomic::Ordering::Acquire),
);
}
let _ = guard;
None
}
pub fn last(&self) -> Option<(K, [u8; V])> {
self.iter().next_back()
}
fn resolve_front_asc(
&self,
bound: &Bound<&K>,
guard: &seize::LocalGuard<'_>,
) -> *mut ConstNode<K, V, D::Loc> {
match bound {
Bound::Included(k) => self.index.find_first_ge(k.as_bytes(), guard),
Bound::Excluded(k) => {
let ge = self.index.find_first_ge(k.as_bytes(), guard);
if !ge.is_null()
&& !unsafe { &*ge }.is_marked()
&& unsafe { &*ge }.key_bytes() == k.as_bytes()
{
crate::skiplist::strip_mark(unsafe {
(*ge).tower(0).load(std::sync::atomic::Ordering::Acquire)
})
} else {
ge
}
}
Bound::Unbounded => crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
}),
}
}
fn resolve_front_rev(
&self,
bound: &Bound<&K>,
guard: &seize::LocalGuard<'_>,
) -> *mut ConstNode<K, V, D::Loc> {
match bound {
Bound::Included(k) => self.index.find_first_ge(k.as_bytes(), guard),
Bound::Excluded(k) => {
let ge = self.index.find_first_ge(k.as_bytes(), guard);
if !ge.is_null()
&& !unsafe { &*ge }.is_marked()
&& unsafe { &*ge }.key_bytes() == k.as_bytes()
{
crate::skiplist::strip_mark(unsafe {
(*ge).tower(0).load(std::sync::atomic::Ordering::Acquire)
})
} else {
ge
}
}
Bound::Unbounded => crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
}),
}
}
fn prefix_bounds(&self, prefix: &[u8]) -> (K, Bound<K>) {
if self.reversed {
let mut search = K::zeroed();
search.as_bytes_mut().fill(0xFF);
search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
let mut end_key = K::zeroed();
end_key.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
(search, Bound::Included(end_key))
} else {
let mut search = K::zeroed();
search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
let end = prefix_to_end_bound::<K>(prefix);
(search, end)
}
}
pub fn prefix_iter(&self, prefix: &[u8]) -> ConstIter<'_, K, V, D::Loc> {
if prefix.len() > size_of::<K>() {
let guard = self.index.collector().enter();
return ConstIter {
list: &self.index,
front: std::ptr::null_mut(),
back: Some(std::ptr::null_mut()),
end: Bound::Unbounded,
start: Bound::Unbounded,
reversed: self.reversed,
done: true,
_guard: guard,
};
}
let guard = self.index.collector().enter();
let (search_key, end) = self.prefix_bounds(prefix);
let front = self.index.find_first_ge(search_key.as_bytes(), &guard);
ConstIter {
list: &self.index,
front,
back: None,
end,
start: Bound::Included(search_key),
reversed: self.reversed,
done: false,
_guard: guard,
}
}
pub fn iter(&self) -> ConstIter<'_, K, V, D::Loc> {
let guard = self.index.collector().enter();
let front = crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
});
ConstIter {
list: &self.index,
front,
back: None,
end: Bound::Unbounded,
start: Bound::Unbounded,
reversed: self.reversed,
done: false,
_guard: guard,
}
}
pub fn range(&self, start: &K, end: &K) -> ConstIter<'_, K, V, D::Loc> {
self.range_bounds(Bound::Included(start), Bound::Excluded(end))
}
pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> ConstIter<'_, K, V, D::Loc> {
let guard = self.index.collector().enter();
if self.reversed {
let front = self.resolve_front_rev(&end, &guard);
ConstIter {
list: &self.index,
front,
back: None,
end: bound_owned(&start),
start: bound_owned(&end),
reversed: true,
done: false,
_guard: guard,
}
} else {
let front = self.resolve_front_asc(&start, &guard);
ConstIter {
list: &self.index,
front,
back: None,
end: bound_owned(&end),
start: bound_owned(&start),
reversed: false,
done: false,
_guard: guard,
}
}
}
pub fn len(&self) -> usize {
self.index.len()
}
pub fn is_empty(&self) -> bool {
self.index.is_empty()
}
pub fn migrate(
&self,
f: impl Fn(&K, &[u8; V]) -> crate::MigrateAction<[u8; V]>,
) -> DbResult<usize> {
use crate::MigrateAction;
let _guard = self.index.collector().enter();
let mut current = crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
});
let mut count = 0;
while !current.is_null() {
let node = unsafe { &*current };
current = crate::skiplist::strip_mark(
node.tower(0).load(std::sync::atomic::Ordering::Acquire),
);
if node.is_marked() {
continue;
}
let value = node.read_value();
match f(&node.key, &value) {
MigrateAction::Keep => {
if H::NEEDS_INIT {
self.hook.on_init(&node.key, &value[..]);
}
}
MigrateAction::Update(value) => {
self.put_no_hook(&node.key, &value)?;
if H::NEEDS_INIT {
self.hook.on_init(&node.key, &value[..]);
}
count += 1;
}
MigrateAction::Delete => {
self.delete_no_hook(&node.key)?;
count += 1;
}
}
}
tracing::info!(mutations = count, "const_tree migration complete");
Ok(count)
}
pub(crate) fn replay_init(&self) {
if !H::NEEDS_INIT {
return;
}
let _guard = self.index.collector().enter();
let mut current = crate::skiplist::strip_mark(unsafe {
(*self.index.head_ptr())
.tower(0)
.load(std::sync::atomic::Ordering::Acquire)
});
while !current.is_null() {
let node = unsafe { &*current };
current = crate::skiplist::strip_mark(
node.tower(0).load(std::sync::atomic::Ordering::Acquire),
);
if !node.is_marked() {
self.hook.on_init(&node.key, &node.read_value()[..]);
}
}
}
pub(crate) fn index(&self) -> &SkipList<ConstNode<K, V, D::Loc>> {
&self.index
}
pub fn shard_for(&self, key: &K) -> usize {
if self.shard_prefix_bits == 0 || self.shard_prefix_bits >= size_of::<K>() * 8 {
let hash = xxhash_rust::xxh3::xxh3_64(key.as_bytes());
return (hash as usize) % self.durability.shard_count();
}
let full_bytes = self.shard_prefix_bits / 8;
let extra_bits = self.shard_prefix_bits % 8;
let hash = if extra_bits == 0 {
xxhash_rust::xxh3::xxh3_64(&key.as_bytes()[..full_bytes])
} else {
let mut buf = K::zeroed();
buf.as_bytes_mut()[..full_bytes].copy_from_slice(&key.as_bytes()[..full_bytes]);
let mask = !((1u8 << (8 - extra_bits)) - 1);
buf.as_bytes_mut()[full_bytes] = key.as_bytes()[full_bytes] & mask;
xxhash_rust::xxh3::xxh3_64(&buf.as_bytes()[..full_bytes + 1])
};
(hash as usize) % self.durability.shard_count()
}
pub fn flush(&self) -> DbResult<()> {
self.durability.flush()
}
}
#[cfg(feature = "replication")]
impl<K: Key, const V: usize, H: WriteHook<K>> crate::replication::ReplicationTarget
for ConstTree<K, V, H, Bitcask>
{
fn apply_entry(
&self,
_shard_inner: &mut crate::shard::ShardInner,
_shard_id: u8,
file_id: u32,
entry_offset: u64,
header: &crate::entry::EntryHeader,
key: &[u8],
value: &[u8],
) -> DbResult<crate::replication::ApplyOutcome> {
use crate::replication::ApplyOutcome;
let key: K = K::from_bytes(key);
let value_offset =
entry_offset + size_of::<crate::entry::EntryHeader>() as u64 + size_of::<K>() as u64;
let disk = DiskLoc::new(file_id, value_offset as u32, header.value_len);
if header.is_tombstone() {
let guard = self.index.collector().enter();
let removed = self.index.remove(key.as_bytes(), &guard);
match removed {
Some(node_ptr) => {
let old_disk = unsafe { &*node_ptr }.read_loc();
Ok(ApplyOutcome::TombstoneRemoved(old_disk))
}
None => Ok(ApplyOutcome::Inserted), }
} else {
let value: [u8; V] = value.try_into().map_err(|_| DbError::CorruptedEntry {
offset: entry_offset,
})?;
let guard = self.index.collector().enter();
let height = random_height();
let node_ptr = ConstNode::alloc(key, value, disk, height);
match self.index.insert(node_ptr, &guard) {
InsertResult::Inserted => Ok(ApplyOutcome::Inserted),
InsertResult::Exists(existing) => {
let old_disk = existing.read_loc();
existing.write_data(disk, &value);
unsafe {
ConstNode::<K, V>::dealloc_node(node_ptr);
}
Ok(ApplyOutcome::Replaced(old_disk))
}
}
}
}
fn try_apply_entry(
&self,
shard_inner: &mut crate::shard::ShardInner,
shard_id: u8,
file_id: u32,
entry_offset: u64,
header: &crate::entry::EntryHeader,
raw_after_header: &[u8],
) -> DbResult<crate::replication::ApplyOutcome> {
use crate::replication::ApplyOutcome;
if raw_after_header.len() < size_of::<K>() + header.value_len as usize {
return Ok(ApplyOutcome::NotMatched);
}
let key = &raw_after_header[..size_of::<K>()];
let value = &raw_after_header[size_of::<K>()..size_of::<K>() + header.value_len as usize];
let crc = crate::entry::compute_crc32(header.gsn, header.value_len, key, value);
if crc != header.crc32 {
return Ok(ApplyOutcome::NotMatched);
}
self.apply_entry(
shard_inner,
shard_id,
file_id,
entry_offset,
header,
key,
value,
)
}
fn key_len(&self) -> usize {
size_of::<K>()
}
}
#[cfg(feature = "replication")]
impl<K: Key, const V: usize, H: WriteHook<K>> ConstTree<K, V, H, Bitcask> {
pub fn start_replication_server(
&self,
bind_addr: std::net::SocketAddr,
signal: crate::shutdown::ShutdownSignal,
) -> crate::error::DbResult<crate::replication::ReplicationServer> {
let consumers = self.install_replication_producers()?;
crate::replication::ReplicationServer::start(
bind_addr,
self.durability.engine.shards().clone(),
consumers,
self.durability.engine.config().max_file_size,
signal,
)
}
fn install_replication_producers(
&self,
) -> crate::error::DbResult<Vec<rtrb::Consumer<crate::replication::ReplicationEntry>>> {
const SPSC_CAPACITY: usize = 4096;
let shards = self.durability.engine.shards();
let mut consumers = Vec::with_capacity(shards.len());
for shard in shards.iter() {
let (p, c) = rtrb::RingBuffer::new(SPSC_CAPACITY);
shard.set_replication_producer(p);
consumers.push(c);
}
Ok(consumers)
}
pub fn start_replication_client(
&self,
leader_addr: std::net::SocketAddr,
registry: std::sync::Arc<crate::replication::ReplicationRegistry>,
signal: crate::shutdown::ShutdownSignal,
) -> crate::error::DbResult<crate::replication::ReplicationClient> {
crate::replication::ReplicationClient::start(
leader_addr,
self.durability.engine.shards().clone(),
registry,
size_of::<K>() as u16,
signal,
)
}
}
#[cfg(feature = "replication")]
impl<K, const V: usize, H> ConstTree<K, V, H, Bitcask>
where
K: Key + Send + Sync + 'static,
H: WriteHook<K> + Send + Sync + 'static,
{
pub fn as_replication_target(
self: &std::sync::Arc<Self>,
) -> Box<dyn crate::replication::ReplicationTarget> {
Box::new(std::sync::Arc::clone(self))
}
}
type ConstShardEvent<K, const V: usize> = (K, Option<[u8; V]>, Option<[u8; V]>);
pub struct ConstShard<'a, K: Key, const V: usize, H: WriteHook<K> = NoHook, D: Durability = Bitcask>
{
tree: &'a ConstTree<K, V, H, D>,
inner: MutexGuard<'a, D::Inner>,
shard_id: usize,
guard: seize::LocalGuard<'a>,
events: Vec<ConstShardEvent<K, V>>,
}
impl<K: Key, const V: usize, H: WriteHook<K>, D: Durability> ConstShard<'_, K, V, H, D> {
pub fn put(&mut self, key: &K, value: &[u8; V]) -> DbResult<Option<[u8; V]>> {
self.check_shard(key)?;
let old = self
.tree
.put_locked(self.shard_id, &mut *self.inner, &self.guard, key, value)?;
if H::NEEDS_WRITE {
self.events.push((*key, old, Some(*value)));
}
Ok(old)
}
pub fn insert(&mut self, key: &K, value: &[u8; V]) -> DbResult<()> {
self.check_shard(key)?;
self.tree
.insert_locked(self.shard_id, &mut *self.inner, &self.guard, key, value)?;
if H::NEEDS_WRITE {
self.events.push((*key, None, Some(*value)));
}
Ok(())
}
pub fn delete(&mut self, key: &K) -> DbResult<Option<[u8; V]>> {
self.check_shard(key)?;
let old = self
.tree
.delete_locked(self.shard_id, &mut *self.inner, &self.guard, key)?;
if H::NEEDS_WRITE
&& let Some(ref old_val) = old
{
self.events.push((*key, Some(*old_val), None));
}
Ok(old)
}
pub fn get(&self, key: &K) -> Option<[u8; V]> {
let node = self.tree.index.get(key.as_bytes(), &self.guard)?;
Some(node.read_value())
}
pub fn get_or_err(&self, key: &K) -> DbResult<[u8; V]> {
self.get(key).ok_or(DbError::KeyNotFound)
}
pub fn contains(&self, key: &K) -> bool {
self.tree.index.get(key.as_bytes(), &self.guard).is_some()
}
fn check_shard(&self, key: &K) -> DbResult<()> {
if self.tree.shard_for(key) != self.shard_id {
return Err(DbError::ShardMismatch);
}
Ok(())
}
}
#[cfg(feature = "armour")]
pub struct ConstTx<'a, K: Key, const V: usize, H: WriteHook<K> = NoHook, D: Durability = Bitcask> {
tree: &'a ConstTree<K, V, H, D>,
inners: Vec<(usize, MutexGuard<'a, D::Inner>)>,
seize: seize::LocalGuard<'a>,
log: Vec<ConstShardEvent<K, V>>,
}
#[cfg(feature = "armour")]
impl<'a, K: Key, const V: usize, H: WriteHook<K>, D: Durability> ConstTx<'a, K, V, H, D> {
fn position(&self, key: &K) -> DbResult<usize> {
let sid = self.tree.shard_for(key);
self.inners
.iter()
.position(|(s, _)| *s == sid)
.ok_or(DbError::ShardMismatch)
}
pub fn try_get(&self, key: &K) -> DbResult<Option<[u8; V]>> {
self.position(key)?;
Ok(self
.tree
.index
.get(key.as_bytes(), &self.seize)
.map(|n| n.read_value()))
}
pub fn try_contains(&self, key: &K) -> DbResult<bool> {
self.position(key)?;
Ok(self.tree.index.get(key.as_bytes(), &self.seize).is_some())
}
pub fn get_or_err(&self, key: &K) -> DbResult<[u8; V]> {
self.try_get(key)?.ok_or(DbError::KeyNotFound)
}
pub fn put(&mut self, key: &K, value: &[u8; V]) -> DbResult<Option<[u8; V]>> {
let i = self.position(key)?;
let (sid, inner) = &mut self.inners[i];
let old = self
.tree
.put_locked(*sid, &mut **inner, &self.seize, key, value)?;
if H::NEEDS_WRITE {
self.log.push((*key, old, Some(*value)));
}
Ok(old)
}
pub fn insert(&mut self, key: &K, value: &[u8; V]) -> DbResult<()> {
let i = self.position(key)?;
let (sid, inner) = &mut self.inners[i];
self.tree
.insert_locked(*sid, &mut **inner, &self.seize, key, value)?;
if H::NEEDS_WRITE {
self.log.push((*key, None, Some(*value)));
}
Ok(())
}
pub fn delete(&mut self, key: &K) -> DbResult<Option<[u8; V]>> {
let i = self.position(key)?;
let (sid, inner) = &mut self.inners[i];
let old = self
.tree
.delete_locked(*sid, &mut **inner, &self.seize, key)?;
if H::NEEDS_WRITE
&& let Some(ref old_val) = old
{
self.log.push((*key, Some(*old_val), None));
}
Ok(old)
}
}
#[cfg(feature = "armour")]
impl<K: Key, const V: usize, H: WriteHook<K>, D: Durability> crate::armour::multi_tx::MultiTx
for ConstTree<K, V, H, D>
{
type Key = K;
type Tx<'a>
= ConstTx<'a, K, V, H, D>
where
Self: 'a;
fn shard_for_key(&self, key: &K) -> usize {
self.shard_for(key)
}
fn begin_tx(&self) -> ConstTx<'_, K, V, H, D> {
ConstTx {
tree: self,
inners: Vec::new(),
seize: self.index.collector().enter(),
log: Vec::new(),
}
}
fn lock_shard_into<'a>(&'a self, shard_id: usize, tx: &mut ConstTx<'a, K, V, H, D>) {
tx.inners
.push((shard_id, self.durability.lock_shard(shard_id)));
}
fn release_locks(
&self,
tx: &mut ConstTx<'_, K, V, H, D>,
) -> crate::armour::multi_tx::SyncNeeds {
let mut needs = crate::armour::multi_tx::SyncNeeds::none();
for (sid, inner) in &tx.inners {
if inner.should_sync() {
needs.push(*sid);
}
}
tx.inners.clear(); needs
}
fn run_sync(&self, needs: crate::armour::multi_tx::SyncNeeds) -> DbResult<()> {
for &sid in needs.shards() {
self.durability.lock_shard(sid).sync()?;
}
Ok(())
}
fn replay_hooks(&self, tx: ConstTx<'_, K, V, H, D>) {
if H::NEEDS_WRITE {
for (k, old, new) in &tx.log {
self.hook.on_write(
k,
old.as_ref().map(|v| &v[..]),
new.as_ref().map(|v| &v[..]),
);
}
}
}
}
fn bound_owned<K: Copy>(b: &Bound<&K>) -> Bound<K> {
match b {
Bound::Included(k) => Bound::Included(**k),
Bound::Excluded(k) => Bound::Excluded(**k),
Bound::Unbounded => Bound::Unbounded,
}
}
fn prefix_to_end_bound<K: Key>(prefix: &[u8]) -> Bound<K> {
let mut incremented = prefix.to_vec();
let mut carry = true;
for byte in incremented.iter_mut().rev() {
if carry {
if *byte == 0xFF {
*byte = 0x00;
} else {
*byte += 1;
carry = false;
break;
}
}
}
if carry {
Bound::Unbounded
} else {
let mut end = K::zeroed();
end.as_bytes_mut()[..incremented.len()].copy_from_slice(&incremented);
Bound::Excluded(end)
}
}
pub struct ConstIter<'a, K: Key, const V: usize, L: Location = DiskLoc> {
list: &'a SkipList<ConstNode<K, V, L>>,
front: *mut ConstNode<K, V, L>,
back: Option<*mut ConstNode<K, V, L>>,
end: Bound<K>,
start: Bound<K>,
reversed: bool,
done: bool,
_guard: seize::LocalGuard<'a>,
}
impl<K: Key, const V: usize, L: Location> Iterator for ConstIter<'_, K, V, L> {
type Item = (K, [u8; V]);
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.done || self.front.is_null() {
return None;
}
let node = unsafe { &*self.front };
let converged = self.back.is_some_and(|back| std::ptr::eq(self.front, back));
self.front = crate::skiplist::strip_mark(
node.tower(0).load(std::sync::atomic::Ordering::Acquire),
);
if converged {
self.done = true;
}
if node.is_marked() {
if converged {
return None;
}
continue;
}
if !self.check_end(&node.key) {
self.done = true;
return None;
}
return Some((node.key, node.read_value()));
}
}
}
impl<K: Key, const V: usize, L: Location> DoubleEndedIterator for ConstIter<'_, K, V, L> {
fn next_back(&mut self) -> Option<Self::Item> {
if self.back.is_none() {
self.back = Some(self.resolve_back());
if self.front.is_null() {
self.done = true;
}
}
loop {
let back = self.back.unwrap_or(std::ptr::null_mut());
if self.done || back.is_null() {
return None;
}
let node = unsafe { &*back };
let key = node.key;
let converged = std::ptr::eq(self.front, back);
self.back = Some(self.list.find_last_lt(key.as_bytes(), &self._guard));
if converged {
self.done = true;
}
if node.is_marked() {
if converged {
return None;
}
continue;
}
if !self.check_start(&key) {
self.done = true;
return None;
}
return Some((key, node.read_value()));
}
}
}
impl<K: Key, const V: usize, L: Location> ConstIter<'_, K, V, L> {
fn resolve_back(&self) -> *mut ConstNode<K, V, L> {
match &self.end {
Bound::Unbounded => self.list.find_last(&self._guard),
Bound::Excluded(k) => self.list.find_last_lt(k.as_bytes(), &self._guard),
Bound::Included(k) => {
let ge = self.list.find_first_ge(k.as_bytes(), &self._guard);
if !ge.is_null()
&& !unsafe { &*ge }.is_marked()
&& unsafe { &*ge }.key_bytes() == k.as_bytes()
{
ge
} else {
self.list.find_last_lt(k.as_bytes(), &self._guard)
}
}
}
}
#[inline(always)]
fn check_end(&self, key: &K) -> bool {
match &self.end {
Bound::Unbounded => true,
Bound::Excluded(end) => {
if self.reversed {
key.as_bytes() > end.as_bytes()
} else {
key.as_bytes() < end.as_bytes()
}
}
Bound::Included(end) => {
if self.reversed {
key.as_bytes() >= end.as_bytes()
} else {
key.as_bytes() <= end.as_bytes()
}
}
}
}
#[inline(always)]
fn check_start(&self, key: &K) -> bool {
match &self.start {
Bound::Unbounded => true,
Bound::Excluded(s) => {
if self.reversed {
key.as_bytes() < s.as_bytes()
} else {
key.as_bytes() > s.as_bytes()
}
}
Bound::Included(s) => {
if self.reversed {
key.as_bytes() <= s.as_bytes()
} else {
key.as_bytes() >= s.as_bytes()
}
}
}
}
}
impl<K: Key, const V: usize, L: Location> ConstIter<'_, K, V, L> {
pub fn collect_vec(&mut self) -> Vec<(K, [u8; V])> {
self.collect()
}
}
#[cfg(test)]
mod tests {
use super::ConstTree;
use crate::fixed::FixedTree;
use crate::hook::WriteHook;
use crate::{Config, DbError, DbResult, FixedConfig};
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use tempfile::tempdir;
#[derive(Default)]
struct RecHook {
writes: AtomicUsize,
#[allow(clippy::type_complexity)]
seq: crate::sync::Mutex<Vec<(u64, Option<Vec<u8>>, Option<Vec<u8>>)>>,
}
impl WriteHook<[u8; 8]> for RecHook {
const NEEDS_OLD_VALUE: bool = true;
fn on_write(&self, key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
self.writes.fetch_add(1, AtomicOrdering::Relaxed);
crate::sync::lock(&self.seq).push((
u64::from_be_bytes(*key),
old.map(<[u8]>::to_vec),
new.map(<[u8]>::to_vec),
));
}
}
fn open_const_hooked(dir: &std::path::Path, hook: RecHook) -> ConstTree<[u8; 8], 4, RecHook> {
let mut cfg = Config::test();
cfg.shard_count = 1;
ConstTree::open_hooked(dir, cfg, hook).expect("open hooked")
}
#[test]
fn compaction_reclaims_dead_bytes_under_overwrite_churn() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 64 * 1024; cfg.write_buffer_size = 16 * 1024;
cfg.compaction_threshold = 0.30; let max_file_size = cfg.max_file_size;
let tree = ConstTree::<[u8; 8], 64>::open(dir.path(), cfg).unwrap();
const N: u64 = 500; const M: u64 = 40; let entry_sz: u64 = 16 + 8 + 64; let live_set = N * entry_sz;
let total_written = N * (M + 1) * entry_sz;
for k in 0..N {
tree.put(&k.to_be_bytes(), &[0u8; 64]).unwrap();
}
for round in 1..=M {
let mut v = [0u8; 64];
v[0] = round as u8;
for k in 0..N {
tree.put(&k.to_be_bytes(), &v).unwrap();
}
}
tree.flush_buffers().unwrap();
let shard = &tree.durability.engine.shards()[0];
let disk = || {
let inner = shard.lock();
inner.active.write_offset + inner.immutable.iter().map(|f| f.total_bytes).sum::<u64>()
};
assert!(
disk() > total_written / 2,
"expected dead-byte buildup before compaction: disk={} total_written={total_written}",
disk()
);
let before = disk();
tree.compact().unwrap();
assert!(
disk() < before,
"one compaction pass reclaimed nothing: {before} -> {}",
disk()
);
for _ in 0..200 {
let before = disk();
tree.compact().unwrap();
if disk() == before {
break;
}
}
assert!(
disk() <= live_set + max_file_size,
"compaction failed to reclaim dead bytes: final disk={} live_set={live_set} \
total_written={total_written}",
disk()
);
let mut v = [0u8; 64];
v[0] = M as u8;
assert_eq!(tree.get(&0u64.to_be_bytes()), Some(v));
assert_eq!(tree.get(&(N - 1).to_be_bytes()), Some(v));
assert_eq!(tree.len(), N as usize);
}
#[test]
fn dead_bytes_credited_on_immutable_files_without_flush() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 64 * 1024;
cfg.write_buffer_size = 16 * 1024;
cfg.compaction_threshold = 0.30;
let threshold = cfg.compaction_threshold;
let tree = ConstTree::<[u8; 8], 64>::open(dir.path(), cfg).unwrap();
const N: u64 = 500;
const M: u64 = 40;
for k in 0..N {
tree.put(&k.to_be_bytes(), &[0u8; 64]).unwrap();
}
for round in 1..=M {
let mut v = [0u8; 64];
v[0] = round as u8;
for k in 0..N {
tree.put(&k.to_be_bytes(), &v).unwrap();
}
}
let shard = &tree.durability.engine.shards()[0];
let inner = shard.lock();
assert!(
!inner.immutable.is_empty(),
"churn should have rotated immutable files"
);
let mut crossed = 0usize;
for f in &inner.immutable {
let dead = inner.dead_bytes.get(&f.file_id).copied().unwrap_or(0);
if f.total_bytes > 0 && dead as f64 / f.total_bytes as f64 > threshold {
crossed += 1;
}
}
assert_eq!(
crossed,
inner.immutable.len(),
"every rotated immutable file must cross the {threshold} dead-ratio \
threshold WITHOUT a flush (dead bytes are credited at write_update \
time, not at flush): {crossed}/{} crossed",
inner.immutable.len(),
);
}
#[test]
fn compaction_max_files_per_pass_controls_backlog_drain() {
const N: u64 = 500;
const M: u64 = 200; let entry_sz: u64 = 16 + 8 + 64;
let live_set = N * entry_sz;
let build = |cap: usize| {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 64 * 1024;
cfg.write_buffer_size = 16 * 1024;
cfg.compaction_threshold = 0.30;
cfg.compaction_max_files_per_pass = cap;
let tree = ConstTree::<[u8; 8], 64>::open(dir.path(), cfg).unwrap();
for k in 0..N {
tree.put(&k.to_be_bytes(), &[0u8; 64]).unwrap();
}
for round in 1..=M {
let mut v = [0u8; 64];
v[0] = round as u8;
for k in 0..N {
tree.put(&k.to_be_bytes(), &v).unwrap();
}
}
tree.flush_buffers().unwrap();
(dir, tree)
};
let disk = |tree: &ConstTree<[u8; 8], 64>| {
let inner = tree.durability.engine.shards()[0].lock();
inner.active.write_offset + inner.immutable.iter().map(|f| f.total_bytes).sum::<u64>()
};
let default_cap = Config::test().compaction_max_files_per_pass;
assert!(default_cap > 0, "default cap must be a finite throttle");
let (_dir_d, tree_d) = build(default_cap);
let before_d = tree_d.durability.engine.shards()[0].lock().immutable.len();
tree_d.compact().unwrap();
let after_d = tree_d.durability.engine.shards()[0].lock().immutable.len();
assert!(
before_d.saturating_sub(after_d) <= default_cap,
"default-capped pass retired more than {default_cap} files: {before_d} -> {after_d}"
);
assert!(
disk(&tree_d) > 10 * live_set,
"one default-capped pass should NOT collapse the backlog: disk={} live_set={live_set}",
disk(&tree_d)
);
let (_dir_u, tree_u) = build(0);
tree_u.compact().unwrap();
assert!(
disk(&tree_u) <= live_set + 64 * 1024,
"unlimited pass must drain the backlog in one call: disk={} live_set={live_set}",
disk(&tree_u)
);
let mut v = [0u8; 64];
v[0] = M as u8;
assert_eq!(tree_u.get(&0u64.to_be_bytes()), Some(v));
assert_eq!(tree_u.get(&(N - 1).to_be_bytes()), Some(v));
assert_eq!(tree_u.len(), N as usize);
}
#[test]
fn compaction_aborts_on_corrupt_key_of_live_entry() {
use std::os::unix::fs::FileExt;
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.compaction_threshold = 0.0; let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
let victim = 1u64.to_be_bytes();
let victim_val = 7u64.to_be_bytes();
let other = 2u64.to_be_bytes();
tree.put(&victim, &victim_val).unwrap();
tree.put(&other, &9u64.to_be_bytes()).unwrap();
tree.put(&other, &10u64.to_be_bytes()).unwrap();
let shard = &tree.durability.engine.shards()[0];
shard.rotate_active_for_test(8).unwrap();
let data_path = dir.path().join("shard_000").join("000001.data");
let f = std::fs::OpenOptions::new()
.write(true)
.open(&data_path)
.unwrap();
f.write_at(&[0xFF], 16).unwrap();
f.sync_all().unwrap();
drop(f);
let res = tree.compact();
assert!(
matches!(res, Err(crate::DbError::CrcMismatch { .. })),
"expected CrcMismatch, got {res:?}"
);
assert!(
data_path.exists(),
"source file must be preserved when compaction aborts on corruption"
);
assert_eq!(tree.get(&victim), Some(victim_val));
}
fn scan_disk_for_key(shard_dir: &std::path::Path, key: &[u8; 8]) -> Option<[u8; 8]> {
use crate::entry::EntryHeader;
use zerocopy::FromBytes;
let hsz = std::mem::size_of::<EntryHeader>();
let mut ids: Vec<u32> = std::fs::read_dir(shard_dir)
.unwrap()
.filter_map(|e| {
let n = e.ok()?.file_name().to_string_lossy().into_owned();
n.strip_suffix(".data")?.parse::<u32>().ok()
})
.collect();
ids.sort_unstable();
let mut found = None;
for id in ids {
let bytes = std::fs::read(shard_dir.join(format!("{id:06}.data"))).unwrap();
let mut off = 0usize;
while off + hsz <= bytes.len() {
let header = match EntryHeader::read_from_bytes(&bytes[off..off + hsz]) {
Ok(h) => h,
Err(_) => break,
};
if header.gsn == 0 && header.crc32 == 0 && header.value_len == 0 {
break; }
let total = hsz + 8 + header.value_len as usize;
if off + total > bytes.len() {
break;
}
if &bytes[off + hsz..off + hsz + 8] == key {
found = if header.is_tombstone() {
None
} else {
Some(bytes[off + hsz + 8..off + total].try_into().unwrap())
};
}
off += total;
}
}
found
}
#[test]
fn compaction_preserves_durable_value_when_overwrite_unflushed() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.compaction_threshold = 0.0;
let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
let key = 1u64.to_be_bytes();
let v1 = 11u64.to_be_bytes();
let v2 = 22u64.to_be_bytes();
tree.put(&key, &v1).unwrap();
let shard = &tree.durability.engine.shards()[0];
shard.rotate_active_for_test(8).unwrap();
tree.put(&key, &v2).unwrap();
tree.compact().unwrap();
let shard_dir = dir.path().join("shard_000");
let durable = scan_disk_for_key(&shard_dir, &key);
assert!(
durable.is_some(),
"compaction dropped the only durable copy of a live key (overwrite was unflushed)"
);
}
#[test]
fn open_resolves_hints_false_for_const() {
let dir = tempfile::tempdir().unwrap();
let cfg = crate::Config::test_no_hints();
let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
assert_eq!(tree.config().hints, Some(false));
}
#[test]
fn const_tree_atomic_fires_hooks_in_order() {
let dir = tempfile::tempdir().unwrap();
let tree = open_const_hooked(dir.path(), RecHook::default());
let k = 7u64.to_be_bytes();
tree.atomic(&k, |s| {
s.put(&k, &[1, 1, 1, 1])?; s.put(&k, &[2, 2, 2, 2])?; s.delete(&k)?; Ok(())
})
.expect("atomic");
assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 3);
let seq = crate::sync::lock(&tree.hook.seq).clone();
assert_eq!(seq[0], (7, None, Some(vec![1, 1, 1, 1])));
assert_eq!(seq[1], (7, Some(vec![1, 1, 1, 1]), Some(vec![2, 2, 2, 2])));
assert_eq!(seq[2], (7, Some(vec![2, 2, 2, 2]), None));
}
#[cfg(feature = "armour")]
#[test]
fn const_tx_routes_get_put_and_rejects_out_of_scope() {
use crate::armour::MultiTx;
let dir = tempfile::tempdir().unwrap();
let tree = open_const_hooked(dir.path(), RecHook::default());
let k1 = [1u8; 8];
let k2 = [2u8; 8];
tree.put(&k1, &[10u8; 4]).unwrap();
tree.hook.writes.store(0, AtomicOrdering::Relaxed);
let s1 = tree.shard_for(&k1);
let s2 = tree.shard_for(&k2);
let mut tx = tree.begin_tx();
tree.lock_shard_into(s1, &mut tx);
if s2 != s1 {
tree.lock_shard_into(s2, &mut tx);
}
assert_eq!(tx.try_get(&k1).unwrap(), Some([10u8; 4]));
assert!(!tx.try_contains(&k2).unwrap());
tx.put(&k2, &[20u8; 4]).unwrap();
assert_eq!(tx.try_get(&k2).unwrap(), Some([20u8; 4]));
let mut unlocked_key = [0u8; 8];
for b in 3u8..=255 {
unlocked_key[0] = b;
let s = tree.shard_for(&unlocked_key);
if s != s1 && s != s2 {
break;
}
}
let s_un = tree.shard_for(&unlocked_key);
if s_un != s1 && s_un != s2 {
assert!(matches!(
tx.try_get(&unlocked_key),
Err(DbError::ShardMismatch)
));
assert!(matches!(
tx.try_contains(&unlocked_key),
Err(DbError::ShardMismatch)
));
assert!(matches!(
tx.put(&unlocked_key, &[0u8; 4]),
Err(DbError::ShardMismatch)
));
}
let needs = tree.release_locks(&mut tx);
tree.run_sync(needs).unwrap();
tree.replay_hooks(tx);
assert_eq!(tree.get(&k2), Some([20u8; 4]));
assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 1);
}
#[cfg(feature = "armour")]
#[test]
fn cross_shard_hook_order_is_closure_order() {
use crate::armour::Db;
let dir = tempfile::tempdir().unwrap();
let db = Db::open_test(dir.path()).unwrap();
let mut cfg = Config::test();
cfg.shard_count = 4;
let tree: ConstTree<[u8; 8], 4, RecHook> =
ConstTree::open_hooked(dir.path().join("hooked"), cfg, RecHook::default())
.expect("open hooked");
let mut cfg2 = Config::test();
cfg2.shard_count = 4;
let other: ConstTree<[u8; 8], 4> =
ConstTree::<[u8; 8], 4>::open(dir.path().join("other"), cfg2).expect("open other");
let k_lo = [1u8; 8];
let mut k_hi = [2u8; 8];
for b in 2u8..=255 {
k_hi[0] = b;
if tree.shard_for(&k_hi) != tree.shard_for(&k_lo) {
break;
}
}
assert_ne!(
tree.shard_for(&k_hi),
tree.shard_for(&k_lo),
"need two shards"
);
let ko = [7u8; 8];
db.atomic2(&tree, &[k_lo, k_hi], &other, &[ko], |t, _o| {
t.put(&k_hi, &[1u8; 4])?;
t.put(&k_lo, &[2u8; 4])?;
Ok(())
})
.unwrap();
let seq = crate::sync::lock(&tree.hook.seq).clone();
let order: Vec<u64> = seq.iter().map(|(k, _, _)| *k).collect();
assert_eq!(
order,
vec![u64::from_be_bytes(k_hi), u64::from_be_bytes(k_lo)],
"hooks must replay in closure order, not shard order"
);
}
#[test]
fn const_tree_atomic_fires_for_applied_on_err() {
let dir = tempfile::tempdir().unwrap();
let tree = open_const_hooked(dir.path(), RecHook::default());
let k = 1u64.to_be_bytes();
let r: DbResult<()> = tree.atomic(&k, |s| {
s.put(&k, &[9, 9, 9, 9])?;
Err(DbError::KeyNotFound) });
assert!(r.is_err());
assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 1); }
#[test]
fn const_tree_atomic_nohook_applies_mutations() {
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let tree = ConstTree::<[u8; 8], 4>::open(dir.path(), cfg).unwrap();
let k = 3u64.to_be_bytes();
tree.atomic(&k, |s| {
s.put(&k, &[5, 5, 5, 5])?;
Ok(())
})
.unwrap();
assert_eq!(tree.get(&k), Some([5, 5, 5, 5]));
}
#[test]
fn const_tree_atomic_hook_reentrancy_no_deadlock() {
use std::sync::{Arc, OnceLock, Weak};
#[derive(Default)]
struct ReentrantHook {
tree: OnceLock<Weak<ConstTree<[u8; 8], 4, ReentrantHook>>>,
seen: AtomicUsize,
}
impl WriteHook<[u8; 8]> for ReentrantHook {
fn on_write(&self, key: &[u8; 8], _old: Option<&[u8]>, _new: Option<&[u8]>) {
if let Some(t) = self.tree.get().and_then(Weak::upgrade) {
let _ = t.get(key); self.seen.fetch_add(1, AtomicOrdering::Relaxed);
}
}
}
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let tree =
Arc::new(ConstTree::open_hooked(dir.path(), cfg, ReentrantHook::default()).unwrap());
tree.hook.tree.set(Arc::downgrade(&tree)).ok();
let k = 2u64.to_be_bytes();
tree.atomic(&k, |s| {
s.put(&k, &[1, 2, 3, 4])?;
Ok(())
})
.unwrap();
assert_eq!(tree.hook.seen.load(AtomicOrdering::Relaxed), 1);
}
#[test]
fn const_tree_atomic_fires_hooks_fixedstore_sync_seam() {
let dir = tempdir().unwrap();
let fixed_cfg = FixedConfig {
shard_count: 1,
grow_step: 64,
sync_batch_size: 1,
..FixedConfig::test()
};
let tree = FixedTree::<[u8; 8], 4, RecHook>::open_with_hook(
dir.path(),
fixed_cfg,
RecHook::default(),
)
.expect("open fixed hooked");
let k = 7u64.to_be_bytes();
tree.atomic(&k, |s| {
s.put(&k, &[1, 1, 1, 1])?;
s.put(&k, &[2, 2, 2, 2])?;
s.delete(&k)?;
Ok(())
})
.expect("atomic");
assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 3);
let seq = crate::sync::lock(&tree.hook.seq).clone();
assert_eq!(seq[0], (7, None, Some(vec![1, 1, 1, 1])));
assert_eq!(seq[1], (7, Some(vec![1, 1, 1, 1]), Some(vec![2, 2, 2, 2])));
assert_eq!(seq[2], (7, Some(vec![2, 2, 2, 2]), None));
}
fn cfg() -> FixedConfig {
FixedConfig {
shard_count: 2,
grow_step: 64,
..FixedConfig::test()
}
}
#[test]
fn compare_delete_fixed_store_match_mismatch_absent() {
let dir = tempdir().unwrap();
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg()).unwrap();
let key = 1u64.to_be_bytes();
let val = 42u64.to_be_bytes();
let other = 99u64.to_be_bytes();
tree.put(&key, &val).unwrap();
assert!(matches!(
tree.compare_delete(&key, &other),
Err(DbError::CasMismatch)
));
assert_eq!(tree.get(&key), Some(val));
assert!(tree.compare_delete(&key, &val).is_ok());
assert_eq!(tree.get(&key), None);
assert!(matches!(
tree.compare_delete(&key, &val),
Err(DbError::KeyNotFound)
));
}
#[test]
fn compare_delete_bitcask_match_mismatch_absent() {
let dir = tempdir().unwrap();
let tree = crate::ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
let key = 7u64.to_be_bytes();
let val = 7u64.to_be_bytes();
let other = 8u64.to_be_bytes();
tree.put(&key, &val).unwrap();
assert!(matches!(
tree.compare_delete(&key, &other),
Err(DbError::CasMismatch)
));
assert_eq!(tree.get(&key), Some(val));
assert!(tree.compare_delete(&key, &val).is_ok());
assert_eq!(tree.get(&key), None);
assert!(matches!(
tree.compare_delete(&key, &val),
Err(DbError::KeyNotFound)
));
}
#[test]
fn fixed_tree_default_desc() {
let dir = tempdir().unwrap();
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), FixedConfig::test()).unwrap();
for i in 1u64..=3 {
tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
}
let keys: Vec<u64> = tree.iter().map(|(k, _)| u64::from_be_bytes(k)).collect();
assert_eq!(keys, [3, 2, 1], "default is DESC (newest first)");
assert_eq!(tree.first().map(|(k, _)| u64::from_be_bytes(k)), Some(3));
assert_eq!(tree.last().map(|(k, _)| u64::from_be_bytes(k)), Some(1));
}
#[test]
fn fixed_tree_ascending_when_reversed_false() {
let dir = tempdir().unwrap();
let cfg = FixedConfig {
reversed: false,
..FixedConfig::test()
};
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
for i in 1u64..=3 {
tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
}
let keys: Vec<u64> = tree.iter().map(|(k, _)| u64::from_be_bytes(k)).collect();
assert_eq!(keys, [1, 2, 3], "reversed=false yields ASC");
assert_eq!(tree.first().map(|(k, _)| u64::from_be_bytes(k)), Some(1));
assert_eq!(tree.last().map(|(k, _)| u64::from_be_bytes(k)), Some(3));
}
#[test]
fn fixed_tree_reversed_is_tunable_across_reopen() {
let dir = tempdir().unwrap();
{
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), FixedConfig::test()).unwrap();
for i in 1u64..=3 {
tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
}
let keys: Vec<u64> = tree.iter().map(|(k, _)| u64::from_be_bytes(k)).collect();
assert_eq!(keys, [3, 2, 1], "default DESC before reopen");
tree.close().unwrap();
}
let cfg = FixedConfig {
reversed: false,
..FixedConfig::test()
};
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
let keys: Vec<u64> = tree.iter().map(|(k, _)| u64::from_be_bytes(k)).collect();
assert_eq!(
keys,
[1, 2, 3],
"ASC after reopen — direction is tunable, no migration"
);
assert_eq!(tree.get(&2u64.to_be_bytes()), Some(2u64.to_be_bytes()));
}
}
#[cfg(all(test, feature = "replication"))]
mod replication_helper_tests {
use crate::FixedConfig;
use crate::fixed::FixedTree;
use tempfile::tempdir;
fn cfg() -> FixedConfig {
FixedConfig {
shard_count: 2,
grow_step: 64,
..FixedConfig::test()
}
}
#[test]
fn get_slot_id_present_and_absent() {
let dir = tempdir().unwrap();
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg()).unwrap();
let key = 1u64.to_be_bytes();
let value = 42u64.to_be_bytes();
tree.put(&key, &value).unwrap();
let slot = tree.get_slot_id(&key).expect("present key must resolve");
assert_eq!(tree.get_slot_id(&key), Some(slot));
let missing = 9999u64.to_be_bytes();
assert_eq!(tree.get_slot_id(&missing), None);
}
#[test]
fn remove_key_if_slot_matches_matching_and_nonmatching() {
let dir = tempdir().unwrap();
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg()).unwrap();
let key = 7u64.to_be_bytes();
let value = 7u64.to_be_bytes();
tree.put(&key, &value).unwrap();
let slot = tree.get_slot_id(&key).unwrap();
assert!(!tree.remove_key_if_slot_matches(&key, slot.wrapping_add(1)));
assert!(tree.contains(&key));
assert!(tree.remove_key_if_slot_matches(&key, slot));
assert!(!tree.contains(&key));
assert!(!tree.remove_key_if_slot_matches(&key, slot));
}
#[test]
fn upsert_replicated_insert_and_update() {
let dir = tempdir().unwrap();
let tree = FixedTree::<[u8; 8], 8>::open(dir.path(), cfg()).unwrap();
let key = 3u64.to_be_bytes();
let value_a = 100u64.to_be_bytes();
let value_b = 200u64.to_be_bytes();
tree.upsert_replicated(&key, value_a, 77);
assert_eq!(tree.get(&key), Some(value_a));
assert_eq!(tree.get_slot_id(&key), Some(77));
tree.upsert_replicated(&key, value_b, 77);
assert_eq!(tree.get(&key), Some(value_b));
assert_eq!(tree.get_slot_id(&key), Some(77));
let value_c = 300u64.to_be_bytes();
tree.upsert_replicated(&key, value_c, 123);
assert_eq!(tree.get(&key), Some(value_c));
assert_eq!(tree.get_slot_id(&key), Some(123));
}
}