use std::hash::Hash;
use std::mem::size_of;
use std::ops::{Bound, ControlFlow};
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::sync::{Arc, OnceLock, Weak};
use crate::Key;
use crate::byte_view::ByteView;
use crate::cache::{BlockCache, BlockKey};
use crate::compaction::{CompactionIndex, compact_shard};
use crate::config::Config;
use crate::ctx::{Ctx, Family, Held, Owned, ShardGuards};
use crate::decoded_cache::DecodedCache;
use crate::disk_loc::DiskLoc;
use crate::engine::Engine;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, WriteHook};
use crate::io::aligned_buf::AlignedBuf;
use crate::map_index::{
Dir, KeyPage, MapIter, MapIterView, OrdKey, Page, ShardIndex, ShardSource, TryMapIter,
dir_from, key_index_batch,
};
use crate::recovery::recover_var_map;
use crate::shard::ShardInner;
use crate::sync::{self, Mutex, MutexGuard};
use crate::ttl::{SweepOutcome, TTL_RETRY_BACKOFF_MS, TtlSlots};
use crate::value_cache::{ValueCache, ValueKey};
use crate::var_tree::VarHookNew;
use crate::watch::{WatchEventInner, WatchWakeBatch};
enum VarMapSweepAction<K: Key> {
Delete(K, Option<ByteView>, u64),
Reschedule(K, u64),
Drop,
Retry(K),
}
const MAX_STALE_RETRIES: usize = 3;
#[cfg(test)]
thread_local! {
pub(crate) static TEST_VALUE_READS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
#[cfg(test)]
pub(crate) fn test_value_reads() -> usize {
TEST_VALUE_READS.with(|c| c.get())
}
#[doc(hidden)]
pub struct VarMapIterVal {
pub(crate) loc: DiskLoc,
pub(crate) value: ByteView,
}
enum UpdateOutcome {
Missing,
Unchanged(ByteView),
Updated { old: ByteView, new: ByteView },
}
pub struct VarMap<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
indexes: Vec<Mutex<ShardIndex<K, DiskLoc>>>,
engine: Engine,
block_cache: BlockCache,
value_cache: ValueCache,
compaction_threshold: f64,
shard_prefix_bits: usize,
pub(crate) hook: H,
iterable: bool,
reversed: bool,
entry_count: AtomicUsize,
tx_self: OnceLock<Weak<Self>>,
ttl: TtlSlots<K>,
pub(crate) counter_dirty: OnceLock<crate::counter_dirty::CounterDirty>,
}
impl<K: Key + Send + Sync + Hash + Eq> VarMap<K> {
#[doc = include_str!("../docs/exclusive-access.md")]
pub fn open(path: impl AsRef<std::path::Path>, config: Config) -> DbResult<Self> {
Self::open_hooked(path, config, NoHook)
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
pub(crate) fn floor_recovered_gsns(&self, gsn_high: u64, shard_gsn_high: &[u64]) {
self.engine.gsn().fetch_max(
gsn_high.saturating_add(1),
std::sync::atomic::Ordering::Relaxed,
);
for (shard_id, high) in shard_gsn_high.iter().copied().enumerate() {
if let Some(shard) = self.engine.shards().get(shard_id) {
shard.set_durable_recovered_gsn(shard.durable_recovered_gsn().max(high));
}
}
}
#[doc = include_str!("../docs/exclusive-access.md")]
pub fn open_hooked(
path: impl AsRef<std::path::Path>,
config: Config,
hook: H,
) -> DbResult<Self> {
#[cfg(feature = "armour")]
if H::EMITS {
return Err(DbError::ChangeStreamRequiresManagedOpen);
}
let config = config.with_resolved_hints(true);
if config.var_cache_warn_needed() {
tracing::warn!(
"Var collection opened with no cache (block_cache=0 and value_cache=0): every get hits disk"
);
}
if config.app_cache_without_direct_io() {
tracing::warn!(
"Var collection has an app cache enabled (block_cache/value_cache) but direct_io is off: \
values are cached twice (app cache + OS page cache). Enable direct_io to avoid redundant \
RAM — exception: keeping the app cache for plaintext caching under encryption."
);
}
Self::open_inner(path, config, hook)
}
pub(crate) 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 iterable = config.iterable;
let reversed = config.reversed;
let block_cache = BlockCache::new(&config.block_cache);
let value_cache = ValueCache::new(&config.value_cache);
let engine = Engine::open(path, config)?;
let shard_count = engine.shards().len();
let mut indexes = Vec::with_capacity(shard_count);
for _ in 0..shard_count {
indexes.push(Mutex::new(ShardIndex::new(false)));
}
let map = Self {
indexes,
engine,
block_cache,
value_cache,
compaction_threshold,
shard_prefix_bits,
hook,
iterable,
reversed,
entry_count: AtomicUsize::new(0),
tx_self: OnceLock::new(),
ttl: TtlSlots::new(),
counter_dirty: OnceLock::new(),
};
let shard_dirs = map.engine.shard_dirs();
let shard_dir_refs = Engine::shard_dir_refs(&shard_dirs);
let shard_ids = map.engine.shard_ids();
let hints = map.engine.hints();
let outcome = recover_var_map::<K>(
&shard_dir_refs,
&shard_ids,
map.indexes(),
hints,
#[cfg(feature = "encryption")]
map.engine.shard_ciphers(),
)?;
for tail in &outcome.active_tails {
map.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
}
for (shard_idx, dead) in outcome.shard_dead_bytes {
map.engine.shards()[shard_idx].install_dead_bytes(dead);
}
for (shard_idx, file_seqs) in outcome.shard_file_seqs {
map.engine.shards()[shard_idx].install_recovery_file_seqs(file_seqs);
}
for (shard_idx, max_gsn) in &outcome.shard_max_gsns {
map.engine.shards()[*shard_idx].set_durable_recovered_gsn(*max_gsn);
}
let max_gsn = outcome.max_gsn;
let initial_len: usize = map.indexes.iter().map(|m| sync::lock(m).len()).sum();
map.entry_count.store(initial_len, AtomicOrdering::Relaxed);
map.engine
.gsn()
.fetch_max(max_gsn + 1, AtomicOrdering::Relaxed);
if hints {
for shard in map.engine.shards().iter() {
shard.set_key_len(size_of::<K>());
}
}
if iterable {
for idx in &map.indexes {
sync::lock(idx).set_iterable_and_rebuild();
}
}
tracing::info!(
key_size = size_of::<K>(),
entries = map.len(),
"var_map recovered"
);
Ok(map)
}
pub(crate) fn inspection_range_bounds_dir(
&self,
start: Bound<&K>,
end: Bound<&K>,
dir: Dir,
batch: usize,
) -> VarMapTryIter<'_, K, H> {
VarMapTryIter {
inner: MapIterView::with_batch(self, dir, batch).try_range_bounds(start, end),
}
}
pub fn clean_shutdown(&self) -> DbResult<()> {
if self.engine.hints() {
self.sync_hints()?;
}
self.engine.close()
}
pub fn close(self) -> DbResult<()> {
self.clean_shutdown()
}
pub fn flush_buffers(&self) -> DbResult<()> {
self.engine.flush_buffers()
}
pub fn dirty_shards(&self) -> usize {
self.engine.dirty_shards()
}
pub(crate) fn durable_flush(&self) -> DbResult<()> {
self.engine.flush()
}
pub fn config(&self) -> &Config {
self.engine.config()
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> CompactionIndex<K> for VarMap<K, H> {
fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
let mut index = sync::lock(&self.indexes[self.shard_for(key)]);
if let Some(disk) = index.get_mut(key)
&& *disk == old_loc
{
*disk = new_loc;
return true;
}
false
}
fn invalidate_blocks(&self, shard_id: u8, file_id: u32, total_bytes: u64) {
self.block_cache
.invalidate_file(shard_id, file_id, total_bytes);
}
fn contains_key(&self, key: &K) -> bool {
self.contains(key)
}
fn is_live(&self, shard_id: u8, key: &K, loc: DiskLoc) -> bool {
let index = sync::lock(&self.indexes[shard_id as usize]);
index.get(key).is_some_and(|disk| *disk == loc)
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> Family for VarMap<K, H> {
type Locked<'a>
= (
MutexGuard<'a, ShardInner>,
MutexGuard<'a, ShardIndex<K, DiskLoc>>,
)
where
Self: 'a;
type Event<'v> = (K, Option<ByteView>, Option<VarHookNew<'v>>);
type Owned = (K, Option<ByteView>, Option<ByteView>);
#[inline]
fn lock_shard(&self, sid: usize) -> Self::Locked<'_> {
let inner = self.engine.shards()[sid].lock();
let index = sync::lock(&self.indexes[sid]);
(inner, index)
}
fn should_sync<'a>(locked: &Self::Locked<'a>) -> bool
where
Self: 'a,
{
locked.0.should_sync()
}
fn sync(&self, sid: usize) -> DbResult<()> {
#[cfg(feature = "ctx-gates")]
crate::ctx_gates_sync::record_sync_call();
self.engine.shards()[sid].flush()
}
fn fire(&self, ev: Self::Event<'_>) {
let (k, old, new) = ev;
self.hook
.on_write(&k, old.as_deref(), new.as_ref().map(|n| n.as_slice()));
}
fn own(ev: Self::Event<'_>) -> Self::Owned {
let (k, old, new) = ev;
(k, old, new.map(VarHookNew::into_owned))
}
fn fire_owned(&self, ev: &Self::Owned) {
let (k, old, new) = ev;
self.hook.on_write(k, old.as_deref(), new.as_deref());
}
fn publish_watch(&self, ev: &Self::Event<'_>, active: bool, wakes: &mut WatchWakeBatch) {
if !H::WATCHES {
return;
}
let Some(channel) = self.hook.watch_channel() else {
return;
};
if !active {
return;
}
let (key, old, new) = ev;
let old_existed = old.is_some();
let old = if self.hook.watch_old() {
old.as_deref()
.map(|bytes| bytes.to_vec().into_boxed_slice())
} else {
None
};
let new = new
.as_ref()
.map(|bytes| bytes.as_slice().to_vec().into_boxed_slice());
channel.publish(
std::sync::Arc::new(WatchEventInner::from_write(
*key,
old,
new,
self.hook.watch_old(),
old_existed,
)),
wakes,
);
}
const WATCHES: bool = H::WATCHES;
fn watch_enter(&self) -> crate::watch::WatchProducer<'_> {
if H::WATCHES {
self.hook
.watch_channel()
.and_then(|channel| channel.enter_producer())
.unwrap_or_else(crate::watch::WatchProducer::inactive)
} else {
crate::watch::WatchProducer::inactive()
}
}
fn apply_projection_for_event(&self, ev: &Self::Event<'_>) -> DbResult<()> {
let (k, old, new) = ev;
let new_bytes = new.as_ref().map(|n| n.as_slice());
self.apply_projection_for_bytes(k, old.as_deref(), new_bytes)
}
fn apply_projection_for_owned(&self, ev: &Self::Owned) -> DbResult<()> {
let (k, old, new) = ev;
self.apply_projection_for_bytes(k, old.as_deref(), new.as_deref())
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
#[doc(hidden)]
pub fn build_ttl_index(&self) -> DbResult<()> {
if !H::EXPIRES {
return Ok(());
}
let n = self.engine.shards().len();
let slots = self.ttl.install_and_clear(n);
let mut skipped = 0usize;
for (sid, slot) in slots.iter().enumerate() {
let inner = self.engine.shards()[sid].lock();
let index = sync::lock(&self.indexes[sid]);
for (key, disk) in index.iter() {
if let Ok(bytes) = self.read_value_locked_result(disk, sid as u8, &inner) {
let deadline = self.hook.deadline(key, &bytes);
sync::lock(slot).record(key, deadline);
} else {
tracing::warn!(
key = ?key.as_bytes(),
"ttl cold build: row omitted from index — value unreadable until rewrite or rebuild"
);
skipped += 1;
}
}
}
if skipped > 0 {
tracing::warn!(
skipped,
"ttl cold build finished with rows omitted from index — sweep will not see them until rewrite or rebuild"
);
}
Ok(())
}
fn record_ttl(&self, shard_id: usize, key: &K, expires_at: Option<u64>) {
if !H::EXPIRES {
return;
}
self.ttl.record(shard_id, key, expires_at);
}
fn deadline_for_bytes(&self, key: &K, value: &[u8]) -> Option<u64> {
if H::EXPIRES {
self.hook.deadline(key, value)
} else {
None
}
}
#[doc(hidden)]
pub fn ttl_debug_len(&self) -> usize {
self.ttl.debug_len()
}
#[doc(hidden)]
pub fn ttl_debug_deadline(&self, key: &K) -> Option<u64> {
self.ttl.debug_deadline(self.shard_for(key), key)
}
#[doc(hidden)]
pub fn ttl_debug_inject_stale(&self, key: &K, deadline: u64) {
self.ttl
.debug_inject_stale(self.shard_for(key), deadline, *key);
}
pub fn sweep_expired(&self, now_ms: u64, budget: usize) -> DbResult<SweepOutcome> {
if !H::EXPIRES {
return Ok(SweepOutcome::default());
}
let Some(slots) = self.ttl.get() else {
return Ok(SweepOutcome::default());
};
let Some(_pass) = self.ttl.begin_pass() else {
return Ok(SweepOutcome::default());
};
let shard_count = slots.len();
let mut remaining = budget;
let mut outcome = SweepOutcome::default();
while remaining > 0 && self.ttl.any_due(now_ms) {
let sid = self.ttl.next_shard(shard_count);
let mut shard_events: Vec<(K, Option<ByteView>, Option<ByteView>)> = Vec::new();
let mut wakes = WatchWakeBatch::default();
let mut locked = self.lock_shard(sid);
let producer = self.watch_enter();
let active = producer.active();
let context_active = H::WATCHES.then(|| crate::watch::ContextActiveGuard::set(active));
let mut idx = sync::lock(&slots[sid]);
while remaining > 0 {
let Some((candidate_dl, key)) = idx.pop_due(now_ms) else {
break;
};
remaining -= 1;
let action = match locked.1.get(&key) {
Some(disk) => match self.read_value_locked_result(disk, sid as u8, &locked.0) {
Ok(bytes) => {
match crate::ttl::catch_deadline(|| self.hook.deadline(&key, &bytes)) {
Some(Some(dl)) if dl <= now_ms => {
VarMapSweepAction::Delete(key, Some(bytes), candidate_dl)
}
Some(Some(dl)) => VarMapSweepAction::Reschedule(key, dl),
Some(None) => {
match crate::ttl::catch_deadline(|| {
self.hook.ttl_decode_failed(&bytes)
}) {
Some(false) => VarMapSweepAction::Drop,
_ => VarMapSweepAction::Retry(key),
}
}
None => VarMapSweepAction::Retry(key),
}
}
Err(_) => VarMapSweepAction::Retry(key),
},
None => VarMapSweepAction::Drop,
};
match action {
VarMapSweepAction::Delete(key, stage2_bytes, candidate_dl) => {
drop(idx);
match self.delete_locked(sid, &mut locked.0, &mut locked.1, &key) {
Ok(existed) => {
idx = sync::lock(&slots[sid]);
if existed {
let old = if H::NEEDS_OLD_VALUE && (H::NEEDS_WRITE || active) {
stage2_bytes
} else {
None
};
let event = (key, old.clone(), None);
self.publish_watch(&event, active, &mut wakes);
if H::NEEDS_WRITE {
shard_events.push((key, old, None));
}
outcome.deleted += 1;
}
}
Err(e) => {
sync::lock(&slots[sid]).record(&key, Some(candidate_dl));
let needs_sync = Self::should_sync(&locked);
locked.release();
crate::ttl::log_sweep_failures(None, &outcome);
drop(context_active);
let _ = crate::projection::replay_owned_events(
self,
&shard_events,
needs_sync,
sid,
);
<Self as Family>::wake_watches(wakes);
return Err(e);
}
}
}
VarMapSweepAction::Reschedule(key, dl) => {
idx.record(&key, Some(dl));
}
VarMapSweepAction::Retry(key) => {
outcome.failed += 1;
idx.record(&key, Some(now_ms + TTL_RETRY_BACKOFF_MS));
}
VarMapSweepAction::Drop => {}
}
}
drop(idx);
let needs_sync = Self::should_sync(&locked);
locked.release();
drop(context_active);
let replay =
crate::projection::replay_owned_events(self, &shard_events, needs_sync, sid);
<Self as Family>::wake_watches(wakes);
replay?;
}
outcome.more = self.ttl.any_due(now_ms);
crate::ttl::log_sweep_failures(None, &outcome);
Ok(outcome)
}
pub fn set_ttl_sweep(&self, enabled: bool) -> DbResult<()> {
if H::EXPIRES {
self.ttl.set_enabled(enabled)
} else {
Ok(())
}
}
pub fn mark_replication_follower(&self) {
if H::EXPIRES {
self.ttl.mark_replication_follower();
}
}
pub fn rebuild_ttl_index(&self) -> DbResult<()> {
if !H::EXPIRES {
return Err(DbError::Client("collection has no TTL policy"));
}
self.build_ttl_index()?;
self.ttl.clear_rebuild_requirement();
Ok(())
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
pub fn compact(&self) -> DbResult<usize> {
let mut total_compacted = 0;
for shard in self.engine.shards().iter() {
total_compacted += compact_shard(shard, self, self.compaction_threshold)?;
}
Ok(total_compacted)
}
pub fn get(&self, key: &K) -> Option<ByteView> {
metrics::counter!("armdb.ops", "op" => "get", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.get");
let result = self.try_get(key);
#[cfg(feature = "hot-path-tracing")]
if let Err(e) = &result {
tracing::error!("VarMap get error: {:?}", e);
}
result.unwrap_or_default()
}
pub fn try_get(&self, key: &K) -> DbResult<Option<ByteView>> {
#[cfg(test)]
TEST_VALUE_READS.with(|c| c.set(c.get() + 1));
let shard_id = self.shard_for(key);
for _ in 0..MAX_STALE_RETRIES {
let (disk, cached) = {
let index = sync::lock(&self.indexes[shard_id]);
let disk = match index.get(key) {
Some(d) => *d,
None => return Ok(None),
};
let cached = if disk.is_value_cache_routed() {
self.value_cache.get(&ValueKey {
shard_id: shard_id as u8,
file_id: disk.file_id,
offset: disk.offset,
})
} else {
None
};
(disk, cached)
};
if let Some(v) = cached {
metrics::counter!("armdb.value_cache.hit").increment(1);
return Ok(Some(v));
}
let is_large = disk.is_value_cache_routed();
if is_large {
metrics::counter!("armdb.value_cache.miss").increment(1);
}
match self.read_value_cached_inner(&disk, shard_id as u8) {
Ok(v) => {
if is_large {
self.value_cache.insert(
ValueKey {
shard_id: shard_id as u8,
file_id: disk.file_id,
offset: disk.offset,
},
v.clone(),
);
}
return Ok(Some(v));
}
Err(DbError::StaleDiskLoc) => {
metrics::counter!("armdb.read.stale_retry", "tree" => "var_map").increment(1);
continue;
}
Err(e) => return Err(e),
}
}
Err(DbError::StaleDiskLoc)
}
pub(crate) fn try_get_decoded<T>(
&self,
key: &K,
cache: &DecodedCache<T>,
decode: impl Fn(&ByteView) -> DbResult<T>,
) -> DbResult<Option<Arc<T>>> {
metrics::counter!("armdb.ops", "op" => "get_shared", "tree" => "var_map").increment(1);
let shard_id = self.shard_for(key);
for _ in 0..MAX_STALE_RETRIES {
let (disk, hit, cached_bytes) = {
let index = sync::lock(&self.indexes[shard_id]);
let disk = match index.get(key) {
Some(d) => *d,
None => return Ok(None),
};
let vkey = ValueKey {
shard_id: shard_id as u8,
file_id: disk.file_id,
offset: disk.offset,
};
let hit = cache.get(&vkey);
let cached_bytes = if hit.is_none() && disk.is_value_cache_routed() {
self.value_cache.get(&vkey)
} else {
None
};
(disk, hit, cached_bytes)
};
if let Some(v) = hit {
metrics::counter!("armdb.decoded_cache.hit").increment(1);
return Ok(Some(v));
}
if cache.is_enabled() {
metrics::counter!("armdb.decoded_cache.miss").increment(1);
}
let vkey = ValueKey {
shard_id: shard_id as u8,
file_id: disk.file_id,
offset: disk.offset,
};
let routed = disk.is_value_cache_routed();
let bytes = match cached_bytes {
Some(v) => {
metrics::counter!("armdb.value_cache.hit").increment(1);
v
}
None => {
if routed {
metrics::counter!("armdb.value_cache.miss").increment(1);
}
match self.read_value_cached_inner(&disk, shard_id as u8) {
Ok(v) => {
if routed {
self.value_cache.insert(vkey, v.clone());
}
v
}
Err(DbError::StaleDiskLoc) => {
metrics::counter!("armdb.read.stale_retry", "tree" => "var_map")
.increment(1);
continue;
}
Err(e) => return Err(e),
}
}
};
let value = Arc::new(decode(&bytes)?);
cache.insert(vkey, Arc::clone(&value), disk.len);
return Ok(Some(value));
}
Err(DbError::StaleDiskLoc)
}
pub fn get_or_err(&self, key: &K) -> DbResult<ByteView> {
self.try_get(key)?.ok_or(DbError::KeyNotFound)
}
pub fn put(&self, key: &K, value: &[u8]) -> DbResult<bool> {
let expires_at = self.deadline_for_bytes(key, value);
self.put_with_deadline(key, value, expires_at)
}
pub(crate) fn put_with_deadline(
&self,
key: &K,
value: &[u8],
expires_at: Option<u64>,
) -> DbResult<bool> {
metrics::counter!("armdb.ops", "op" => "put", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.put");
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let watching = H::WATCHES && crate::watch::context_active();
let old_value = if (H::NEEDS_OLD_VALUE && (H::NEEDS_WRITE || watching)) || H::CLAIMS {
if let Some(disk) = index.get(key) {
Some(self.read_value_locked_result(disk, shard_id as u8, inner)?)
} else {
None
}
} else {
None
};
self.claim_guard_bytes(key, old_value.as_deref(), Some(value))?;
let existed = self.put_locked(shard_id, inner, index, key, value, expires_at)?;
if H::NEEDS_WRITE || watching {
*ev_out = Some((*key, old_value, Some(VarHookNew::Borrowed(value))));
}
Ok(existed)
})
}
pub fn insert(&self, key: &K, value: &[u8]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "insert", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.insert");
let expires_at = self.deadline_for_bytes(key, value);
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
if (H::EMITS || (H::CLAIMS && self.hook.claim_targets().is_some()))
&& index.get(key).is_none()
{
self.claim_guard_bytes(key, None, Some(value))?;
}
self.insert_locked(shard_id, inner, index, key, value, expires_at)?;
let watching = H::WATCHES && crate::watch::context_active();
let ev = if H::NEEDS_WRITE || watching {
Some((*key, None, Some(VarHookNew::Borrowed(value))))
} else {
None
};
*ev_out = ev;
Ok(())
})
}
pub fn delete(&self, key: &K) -> DbResult<bool> {
metrics::counter!("armdb.ops", "op" => "delete", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.delete");
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let watching = H::WATCHES && crate::watch::context_active();
let old_value = if (H::NEEDS_OLD_VALUE && (H::NEEDS_WRITE || watching)) || H::CLAIMS {
if let Some(disk) = index.get(key) {
Some(self.read_value_locked_result(disk, shard_id as u8, inner)?)
} else {
None
}
} else {
None
};
self.claim_guard_bytes(key, old_value.as_deref(), None)?;
let existed = self.delete_locked(shard_id, inner, index, key)?;
let ev = if existed && (H::NEEDS_WRITE || watching) {
Some((*key, old_value, None))
} else {
None
};
*ev_out = ev;
Ok(existed)
})
}
pub fn cas(&self, key: &K, expected: &[u8], new_value: &[u8]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "cas", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.cas");
let expires_at = self.deadline_for_bytes(key, new_value);
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
let current = self.read_value_locked_result(&disk, shard_id as u8, inner)?;
if current.as_ref() != expected {
return Err(DbError::CasMismatch);
}
self.claim_guard_bytes(key, Some(current.as_ref()), Some(new_value))?;
let (new_disk_loc, _gsn) =
inner.append_entry(shard_id as u8, key.as_bytes(), new_value, false)?;
inner.add_dead_bytes(
disk.file_id,
crate::entry::entry_size(size_of::<K>(), disk.len),
);
let _old = index.upsert(*key, new_disk_loc);
debug_assert!(_old.is_some(), "cas: key must exist in index");
self.record_ttl(shard_id, key, expires_at);
let ev = if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((
*key,
if H::NEEDS_OLD_VALUE {
Some(current)
} else {
None
},
Some(VarHookNew::Borrowed(new_value)),
))
} else {
None
};
*ev_out = ev;
Ok(())
})
}
pub fn compare_delete(&self, key: &K, expected: &[u8]) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.compare_delete");
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
let current = self.read_value_locked_result(&disk, shard_id as u8, inner)?;
if current.as_ref() != expected {
return Err(DbError::CasMismatch);
}
self.claim_guard_bytes(key, Some(current.as_ref()), None)?;
let removed = self.delete_locked(shard_id, inner, index, key)?;
debug_assert!(
removed,
"compare_delete: key must exist after a successful match"
);
let ev = if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((
*key,
if H::NEEDS_OLD_VALUE {
Some(current)
} else {
None
},
None,
))
} else {
None
};
*ev_out = ev;
Ok(())
})
}
pub(crate) fn compare_delete_if_loc(&self, key: &K, expected: DiskLoc) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "var_map").increment(1);
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
if disk != expected {
return Err(DbError::CasMismatch);
}
let old_value = if (H::NEEDS_OLD_VALUE
&& (H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active())))
|| H::CLAIMS
{
Some(self.read_value_locked_result(&disk, shard_id as u8, inner)?)
} else {
None
};
self.claim_guard_bytes(key, old_value.as_deref(), None)?;
let removed = self.delete_locked(shard_id, inner, index, key)?;
debug_assert!(
removed,
"compare_delete_if_loc: key must exist after a successful loc match"
);
let ev = if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((*key, old_value, None))
} else {
None
};
*ev_out = ev;
Ok(())
})
}
pub fn update(&self, key: &K, f: impl FnOnce(&[u8]) -> ByteView) -> DbResult<Option<ByteView>> {
self.update_inner(key, f, false)
}
pub fn fetch_update(
&self,
key: &K,
f: impl FnOnce(&[u8]) -> ByteView,
) -> DbResult<Option<ByteView>> {
self.update_inner(key, f, true)
}
fn try_update_locked(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, DiskLoc>,
key: &K,
f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
) -> DbResult<UpdateOutcome> {
let disk = match index.get(key) {
Some(d) => *d,
None => return Ok(UpdateOutcome::Missing),
};
let current = self.read_value_locked_result(&disk, shard_id as u8, inner)?;
let new_value = match f(¤t)? {
Some(v) => v,
None => return Ok(UpdateOutcome::Unchanged(current)),
};
self.claim_guard_bytes(key, Some(current.as_ref()), Some(&new_value))?;
let (new_disk_loc, _gsn) =
inner.append_entry(shard_id as u8, key.as_bytes(), &new_value, false)?;
inner.add_dead_bytes(
disk.file_id,
crate::entry::entry_size(size_of::<K>(), disk.len),
);
let _old = index.upsert(*key, new_disk_loc);
debug_assert!(_old.is_some(), "update: key must exist in index");
self.record_ttl(shard_id, key, self.deadline_for_bytes(key, &new_value));
Ok(UpdateOutcome::Updated {
old: current,
new: new_value,
})
}
pub(crate) fn try_update_inner(
&self,
key: &K,
f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
return_old: bool,
) -> DbResult<Option<ByteView>> {
metrics::counter!("armdb.ops", "op" => "update", "tree" => "var_map").increment(1);
#[cfg(feature = "hot-path-tracing")]
tracing::trace!("var_map.update");
let shard_id = self.shard_for(key);
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
match self.try_update_locked(shard_id, inner, index, key, f)? {
UpdateOutcome::Missing => Ok(None),
UpdateOutcome::Unchanged(current) => Ok(Some(current)),
UpdateOutcome::Updated { old, new } => {
let (result, ev) =
if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
let event_old = if H::NEEDS_OLD_VALUE {
Some(old.clone())
} else {
None
};
let result = Some(if return_old { old } else { new.clone() });
(
result,
Some((*key, event_old, Some(VarHookNew::Owned(new)))),
)
} else {
(Some(if return_old { old } else { new }), None)
};
{
*ev_out = ev;
Ok(result)
}
}
}
})
}
fn update_inner(
&self,
key: &K,
f: impl FnOnce(&[u8]) -> ByteView,
return_old: bool,
) -> DbResult<Option<ByteView>> {
self.try_update_inner(key, |bytes| Ok(Some(f(bytes))), return_old)
}
pub fn contains(&self, key: &K) -> bool {
let index = sync::lock(&self.indexes[self.shard_for(key)]);
index.contains_key(key)
}
pub fn entry_len(&self, key: &K) -> Option<u32> {
let index = sync::lock(&self.indexes[self.shard_for(key)]);
index.get(key).map(|disk| disk.len)
}
pub fn len(&self) -> usize {
self.entry_count.load(AtomicOrdering::Relaxed)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn sync_hints(&self) -> DbResult<()> {
for shard in self.engine.shards().iter() {
shard.write_active_hint(size_of::<K>())?;
}
Ok(())
}
pub fn iter_view(&self) -> Option<VarMapIterView<'_, K, H>> {
self.iterable.then(|| VarMapIterView {
map: self,
inner: MapIterView::new(self, dir_from(self.reversed)),
})
}
pub fn scope<const N: usize>(
&self,
key: &K,
) -> Option<crate::scope::Scope<'_, crate::scope::MapScan<&Self>, N>> {
crate::scope::map_scope::<_, _, N>(self, key, dir_from(self.reversed), self.iterable)
}
pub fn warmup(&self) -> DbResult<()> {
use std::collections::BTreeSet;
let mut blocks: BTreeSet<(u8, u32, u64)> = BTreeSet::new();
for (i, index) in self.indexes.iter().enumerate() {
let map = sync::lock(index);
for disk in map.values() {
if disk.is_value_cache_routed() {
continue;
}
let block_offset = disk.offset as u64 & !4095;
blocks.insert((i as u8, disk.file_id, block_offset));
if (disk.offset & 4095) as usize + disk.len as usize > 4096 {
blocks.insert((i as u8, disk.file_id, block_offset + 4096));
}
}
}
for (shard_id, file_id, block_offset) in &blocks {
let key = BlockKey {
shard_id: *shard_id,
file_id: *file_id,
block_offset: *block_offset,
};
if self.block_cache.get(&key).is_some() {
continue;
}
let shard = &self.engine.shards()[*shard_id as usize];
let (buf, is_full_block) = shard.read_block(*file_id, *block_offset)?;
if is_full_block {
self.block_cache.insert(key, Arc::new(buf));
}
}
Ok(())
}
pub fn for_each(&self, mut f: impl FnMut(K, ByteView)) {
for i in 0..self.engine.shards().len() {
let keys: Vec<K> = {
let index = sync::lock(&self.indexes[i]);
index.keys().copied().collect()
};
for key in keys {
match self.try_get(&key) {
Ok(Some(value)) => f(key, value),
Ok(None) => {}
Err(e) => {
tracing::warn!(
key = ?key.as_bytes(),
error = ?e,
"var_map for_each: skipping entry — value read failed"
);
}
}
}
}
}
pub(crate) fn for_each_while(
&self,
mut f: impl FnMut(DbResult<(K, ByteView)>) -> ControlFlow<()>,
) {
for i in 0..self.engine.shards().len() {
let keys: Vec<K> = {
let index = sync::lock(&self.indexes[i]);
index.keys().copied().collect()
};
for key in keys {
let step = match self.try_get(&key) {
Ok(Some(value)) => f(Ok((key, value))),
Ok(None) => ControlFlow::Continue(()),
Err(e) => f(Err(e)),
};
if step.is_break() {
return;
}
}
}
}
pub fn migrate(
&self,
f: impl Fn(&K, &[u8]) -> crate::MigrateAction<ByteView>,
) -> DbResult<usize> {
self.migrate_inner(f, true)
}
pub(crate) fn migrate_inner(
&self,
f: impl Fn(&K, &[u8]) -> crate::MigrateAction<ByteView>,
fire_init: bool,
) -> DbResult<usize> {
use crate::MigrateAction;
let mut count = 0;
for i in 0..self.engine.shards().len() {
let keys: Vec<K> = {
let index = sync::lock(&self.indexes[i]);
index.keys().copied().collect()
};
for key in keys {
let value = match self.read_value_cached_result(&key) {
Ok(Some(v)) => v,
Ok(None) => continue,
Err(e) => {
tracing::error!(
key = ?key.as_bytes(),
error = %e,
"var_map migrate: value read failed — aborting migration step"
);
return Err(e);
}
};
match f(&key, &value) {
MigrateAction::Keep => {
if fire_init && H::NEEDS_INIT {
self.hook.on_init(&key, &value);
}
}
MigrateAction::Update(new_value) => {
{
let mut inner = self.engine.shards()[i].lock();
let mut index = sync::lock(&self.indexes[i]);
self.claim_guard_bytes(&key, Some(&value), Some(&new_value))?;
let expires_at = self.deadline_for_bytes(&key, &new_value);
self.put_locked(
i, &mut inner, &mut index, &key, &new_value, expires_at,
)?;
}
if fire_init && H::NEEDS_INIT {
self.hook.on_init(&key, &new_value);
}
count += 1;
}
MigrateAction::Delete => {
let mut inner = self.engine.shards()[i].lock();
let mut index = sync::lock(&self.indexes[i]);
self.claim_guard_bytes(&key, Some(&value), None)?;
self.delete_locked(i, &mut inner, &mut index, &key)?;
count += 1;
}
}
}
}
tracing::info!(mutations = count, "var_map migration complete");
Ok(count)
}
pub(crate) fn replay_init(&self) {
if !H::NEEDS_INIT {
return;
}
let mut count = 0usize;
for i in 0..self.engine.shards().len() {
let keys: Vec<K> = {
let index = sync::lock(&self.indexes[i]);
index.keys().copied().collect()
};
for key in keys {
let value = match self.get(&key) {
Some(v) => v,
None => {
tracing::warn!(
key = ?key.as_bytes(),
"var_map replay_init: skipping entry — value read failed"
);
continue;
}
};
self.hook.on_init(&key, &value);
count += 1;
}
}
tracing::debug!(replayed = count, "var_map replay_init complete");
}
pub(crate) fn effective_projection_limits(&self) -> (usize, usize) {
let cfg = self.config();
let (hook_rows, hook_bytes) = self.hook.projection_limits();
(
hook_rows.min(cfg.max_projected_rows),
hook_bytes.min(cfg.max_projected_bytes),
)
}
#[allow(clippy::collapsible_if)]
#[inline]
fn claim_guard_bytes(&self, key: &K, old: Option<&[u8]>, new: Option<&[u8]>) -> DbResult<()> {
#[cfg(feature = "armour")]
if H::EMITS && !crate::txn_depth::in_stream_apply() {
if let Some(targets) = self.hook.stream_targets() {
let collection = self
.counter_dirty
.get()
.map(crate::counter_dirty::CounterDirty::name)
.unwrap_or("");
return Err(DbError::ChangeStreamRequiresTx {
collection,
queue: targets.first_name(),
});
}
}
if !H::CLAIMS {
return Ok(());
}
if self.hook.claim_targets().is_none() {
return Ok(());
}
let collection = self
.counter_dirty
.get()
.map(crate::counter_dirty::CounterDirty::name)
.unwrap_or("");
crate::unique::guard_nontx_write_bytes(
collection,
&self.hook,
self.effective_projection_limits(),
key,
old,
new,
)
}
pub(crate) fn reconcile_projections(&self) -> DbResult<()> {
if !H::PROJECTS {
return Ok(());
}
let Some(targets) = self.hook.projection_targets() else {
return Ok(());
};
let (max_rows, max_bytes) = self.effective_projection_limits();
crate::projection::reconcile_carrier_projections(
targets,
|emit| {
let mut err: Option<DbError> = None;
self.for_each_while(|row| {
if err.is_some() {
return ControlFlow::Break(());
}
let row = row
.and_then(|r| crate::projection::reconcile_value_read_fault().map(|()| r));
match row {
Ok((k, bytes)) => {
let mut sink = crate::projection::ProjectionSink::new(
targets, max_rows, max_bytes,
);
match self
.hook
.project(&k, None, Some(bytes.as_ref()), &mut sink)
.and_then(|()| {
for row in sink.rows() {
emit(row.target, &row.key, &row.value)?;
}
Ok(())
}) {
Ok(()) => ControlFlow::Continue(()),
Err(e) => {
err = Some(e);
ControlFlow::Break(())
}
}
}
Err(e) => {
err = Some(e);
ControlFlow::Break(())
}
}
});
err.map_or(Ok(()), Err)
},
true,
)
}
pub(crate) fn reconcile_counters(&self) -> DbResult<u64> {
if !H::COUNTS {
return Ok(0);
}
let Some(targets) = self.hook.counter_targets() else {
return Ok(0);
};
let (max_rows, max_bytes) = self.effective_projection_limits();
crate::counter::recount(
|k: &K, v: &[u8], sink| self.hook.counter_project(k, None, Some(v), sink),
targets,
(max_rows, max_bytes),
&mut |visit| {
let mut err: Option<DbError> = None;
self.for_each_while(|row| {
if err.is_some() {
return ControlFlow::Break(());
}
let row = row
.and_then(|r| crate::projection::reconcile_value_read_fault().map(|()| r));
match row {
Ok((k, bytes)) => match visit(&k, bytes.as_ref()) {
Ok(()) => ControlFlow::Continue(()),
Err(e) => {
err = Some(e);
ControlFlow::Break(())
}
},
Err(e) => {
err = Some(e);
ControlFlow::Break(())
}
}
});
err.map_or(Ok(()), Err)
},
)
}
pub(crate) fn counter_target_names(&self) -> Vec<&'static str> {
self.hook
.counter_targets()
.map(|t| t.all().map(|m| m.collection_name()).collect())
.unwrap_or_default()
}
pub fn atomic<R>(
&self,
shard_key: &K,
f: impl FnOnce(&mut VarMapShard<'_, K, H>) -> DbResult<R>,
) -> DbResult<R> {
let _atomic_depth = crate::AtomicDepthGuard::try_enter()?;
let shard_id = self.shard_for(shard_key);
let mut shard = VarMapShard {
tree: self,
held: Held {
locked: self.lock_shard(shard_id),
sid: shard_id,
events: Vec::new(),
watch_wakes: WatchWakeBatch::default(),
publish_watches: true,
},
};
let result = f(&mut shard);
let VarMapShard { held, .. } = shard;
let Held {
locked,
events,
sid,
watch_wakes,
..
} = held;
let needs_sync = Self::should_sync(&locked);
locked.release();
let result =
crate::projection::finish_atomic_teardown(self, sid, result, needs_sync, &events);
<Self as Family>::wake_watches(watch_wakes);
result
}
pub(crate) fn mark_counter_dirty(&self) -> DbResult<()> {
match self.counter_dirty.get() {
Some(d) => d.mark(),
None => Ok(()),
}
}
fn apply_projection_for_bytes(
&self,
key: &K,
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<()> {
if !H::PROJECTS && !H::COUNTS {
return Ok(());
}
let (max_rows, max_bytes) = self.effective_projection_limits();
let mut plan = match (H::PROJECTS, self.hook.projection_targets()) {
(true, Some(targets)) => crate::projection::diff_projection_bytes(
&self.hook, targets, max_rows, max_bytes, key, old, new,
)?,
_ => crate::projection::PendingProjection::default(),
};
if H::COUNTS
&& let Some(counter_targets) = self.hook.counter_targets()
{
let deltas = crate::counter::diff_counters_bytes(
&self.hook,
counter_targets,
max_rows,
max_bytes,
key,
old,
new,
)?;
if !deltas.is_empty() {
plan.set_deltas(deltas, counter_targets.clone());
}
}
plan.apply_or_incomplete()
}
pub(crate) fn indexes(&self) -> &[Mutex<ShardIndex<K, DiskLoc>>] {
&self.indexes
}
fn put_locked(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, DiskLoc>,
key: &K,
value: &[u8],
expires_at: Option<u64>,
) -> DbResult<bool> {
let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;
if let Some(old_disk) = index.upsert(*key, disk_loc) {
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
self.record_ttl(shard_id, key, expires_at);
Ok(true)
} else {
self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
self.record_ttl(shard_id, key, expires_at);
Ok(false)
}
}
fn insert_locked(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, DiskLoc>,
key: &K,
value: &[u8],
expires_at: Option<u64>,
) -> DbResult<()> {
if index.contains_key(key) {
return Err(DbError::KeyExists);
}
let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;
index.upsert(*key, disk_loc);
self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
self.record_ttl(shard_id, key, expires_at);
Ok(())
}
fn delete_locked(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, DiskLoc>,
key: &K,
) -> DbResult<bool> {
#[cfg(test)]
if crate::durability::ttl_delete_should_fail() {
return Err(DbError::Internal("test: TTL delete injected failure"));
}
if !index.contains_key(key) {
return Ok(false);
}
inner.append_entry(shard_id as u8, key.as_bytes(), &[], true)?;
if let Some(old_disk) = index.remove(key) {
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
self.entry_count.fetch_sub(1, AtomicOrdering::Relaxed);
self.record_ttl(shard_id, key, None);
Ok(true)
} else {
Ok(false)
}
}
fn read_value_cached_result(&self, key: &K) -> DbResult<Option<ByteView>> {
self.try_get(key)
}
fn read_value_cached_inner(&self, disk: &DiskLoc, shard_id: u8) -> DbResult<ByteView> {
let len = disk.len as usize;
let start = (disk.offset & 4095) as usize;
if start + len > 8192 {
{
let shard = &self.engine.shards()[shard_id as usize];
let inner = shard.lock();
if inner.active.file_id == disk.file_id {
if let Some(bytes) = inner.write_buf.read(disk.offset as u64, len) {
return Ok(ByteView::new(bytes));
}
if inner.straddles_write_buf(disk) {
return Ok(ByteView::from_vec(
inner.read_straddling_value_locked(disk)?,
));
}
}
}
let shard = &self.engine.shards()[shard_id as usize];
return Ok(ByteView::from_vec(shard.read_value(
disk.file_id,
disk.offset,
len,
)?));
}
let block_offset = disk.offset as u64 & !4095;
let cache_key = BlockKey {
shard_id,
file_id: disk.file_id,
block_offset,
};
if let Some(block) = self.block_cache.get(&cache_key) {
metrics::counter!("armdb.block_cache.hit").increment(1);
if start + len > 4096 {
let shard = &self.engine.shards()[shard_id as usize];
let inner = shard.lock();
if inner.active.file_id == disk.file_id && inner.straddles_write_buf(disk) {
return Ok(ByteView::from_vec(
inner.read_straddling_value_locked(disk)?,
));
}
}
return Self::extract_from_block(&block, start, len, || {
self.get_or_read_block(shard_id, disk.file_id, block_offset + 4096)
});
}
{
let shard = &self.engine.shards()[shard_id as usize];
let inner = shard.lock();
if inner.active.file_id == disk.file_id {
if let Some(bytes) = inner.write_buf.read(disk.offset as u64, len) {
return Ok(ByteView::new(bytes));
}
if inner.straddles_write_buf(disk) {
return Ok(ByteView::from_vec(
inner.read_straddling_value_locked(disk)?,
));
}
}
}
metrics::counter!("armdb.block_cache.miss").increment(1);
let block = self.get_or_read_block(shard_id, disk.file_id, block_offset)?;
Self::extract_from_block(&block, start, len, || {
self.get_or_read_block(shard_id, disk.file_id, block_offset + 4096)
})
}
fn extract_from_block(
block: &AlignedBuf,
start: usize,
len: usize,
next_block: impl FnOnce() -> DbResult<Arc<AlignedBuf>>,
) -> DbResult<ByteView> {
debug_assert!(
start + len <= 8192,
"extract_from_block supports at most 2 blocks (8192 bytes)"
);
if start + len <= 4096 {
Ok(ByteView::new(&block[start..start + len]))
} else {
let next = next_block()?;
let first_part = &block[start..];
let second_len = len - first_part.len();
let mut combined = Vec::with_capacity(len);
combined.extend_from_slice(first_part);
combined.extend_from_slice(&next[..second_len]);
Ok(ByteView::from_vec(combined))
}
}
fn get_or_read_block(
&self,
shard_id: u8,
file_id: u32,
block_offset: u64,
) -> DbResult<Arc<AlignedBuf>> {
let key = BlockKey {
shard_id,
file_id,
block_offset,
};
if let Some(cached) = self.block_cache.get(&key) {
return Ok(cached);
}
let shard = &self.engine.shards()[shard_id as usize];
let (buf, is_full_block) = shard.read_block(file_id, block_offset)?;
let arc = Arc::new(buf);
if is_full_block {
self.block_cache.insert(key, arc.clone());
}
Ok(arc)
}
fn read_value_locked_result(
&self,
disk: &DiskLoc,
shard_id: u8,
inner: &ShardInner,
) -> DbResult<ByteView> {
let len = disk.len as usize;
if inner.active.file_id == disk.file_id {
if let Some(bytes) = inner.write_buf.read(disk.offset as u64, len) {
return Ok(ByteView::new(bytes));
}
if inner.straddles_write_buf(disk) {
return Ok(ByteView::from_vec(
inner.read_straddling_value_locked(disk)?,
));
}
}
let block_offset = disk.offset as u64 & !4095;
let start = (disk.offset & 4095) as usize;
if start + len <= 4096 {
let cache_key = BlockKey {
shard_id,
file_id: disk.file_id,
block_offset,
};
if let Some(block) = self.block_cache.get(&cache_key) {
return Ok(ByteView::new(&block[start..start + len]));
}
let (buf, is_full_block) = inner.read_block_locked(disk.file_id, block_offset)?;
let arc = Arc::new(buf);
if is_full_block {
self.block_cache.insert(cache_key, arc.clone());
}
return Ok(ByteView::new(&arc[start..start + len]));
}
let bytes = inner.read_value_from_disk_locked(disk)?;
Ok(ByteView::new(&bytes))
}
pub fn shard_for(&self, key: &K) -> usize {
crate::shard_for_key(key, self.shard_prefix_bits, self.engine.shards().len())
}
pub(crate) fn name(&self) -> &'static str {
""
}
pub(crate) fn typ_hash(&self) -> u64 {
0
}
pub(crate) fn shard_count(&self) -> usize {
self.engine.shards().len()
}
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) fn bind_tx_self(arc: &Arc<Self>) {
let _ = arc.tx_self.set(Arc::downgrade(arc));
}
pub(crate) fn tx_arc(&self) -> DbResult<Arc<Self>> {
self.tx_self
.get()
.and_then(Weak::upgrade)
.ok_or(DbError::Internal(
"VarMap::bind_tx_self must be called before transaction use",
))
}
pub(crate) fn tx_install_projection_metadata(&self, db: &crate::armour::Db) -> DbResult<()> {
if H::PROJECTS
&& let Some(targets) = self.hook.projection_targets()
{
db.tx_register_projection_targets(targets)?;
}
if H::COUNTS
&& let Some(targets) = self.hook.counter_targets()
{
db.tx_register_counter_targets(targets)?;
}
if H::CLAIMS
&& let Some(targets) = self.hook.claim_targets()
{
db.tx_register_claim_targets(targets)?;
}
if H::EMITS
&& let Some(targets) = self.hook.stream_targets()
{
db.tx_register_stream_targets(targets)?;
}
Ok(())
}
pub(crate) fn tx_preplan_projection_targets(
&self,
) -> DbResult<Vec<Arc<dyn crate::txn::participant::TxParticipant>>> {
let mut out = Vec::new();
if H::PROJECTS
&& let Some(targets) = self.hook.projection_targets()
{
out.extend(
targets
.erased_targets()
.iter()
.filter_map(|h| targets.resolve_tx_participant(h.name)),
);
}
if H::COUNTS
&& let Some(targets) = self.hook.counter_targets()
{
out.extend(targets.all().map(|m| m.as_tx_participant()));
}
if H::CLAIMS
&& let Some(targets) = self.hook.claim_targets()
{
out.extend(targets.all().map(|t| t.as_tx_participant()));
}
if H::EMITS
&& let Some(targets) = self.hook.stream_targets()
{
out.extend(targets.tx_participants());
}
Ok(out)
}
#[allow(dead_code)]
pub(crate) fn tx_preplan_stream_targets(
&self,
) -> Vec<Arc<dyn crate::txn::participant::TxParticipant>> {
self.hook
.stream_targets()
.map(|t| t.tx_participants())
.unwrap_or_default()
}
#[allow(dead_code)]
pub(crate) fn tx_prepare_streams(
&self,
) -> DbResult<Vec<crate::queue::retention::ProducerPermit>> {
self.hook
.stream_targets()
.map_or(Ok(Vec::new()), crate::QueueTargets::prepare_streams)
}
pub(crate) fn tx_stream_writes(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<crate::txn::participant::TxStreamWrites> {
if !H::EMITS {
return Ok(crate::txn::participant::TxStreamWrites::empty());
}
let Some(targets) = self.hook.stream_targets() else {
return Ok(crate::txn::participant::TxStreamWrites::empty());
};
let key = K::from_bytes(key);
let (max_rows, max_bytes) = self.effective_projection_limits();
let mut sink = crate::EmitSink::new(targets, max_rows, max_bytes);
self.hook.stream_project(&key, old, new, &mut sink)?;
targets.lower(&sink)
}
pub(crate) fn tx_project_writes(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::TxDerivedWrite>> {
if !H::PROJECTS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.projection_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let (max_rows, max_bytes) = self.effective_projection_limits();
let plan = crate::projection::diff_projection_bytes(
&self.hook, targets, max_rows, max_bytes, &key_typed, old, new,
)?;
plan.into_tx_derived(targets)
}
pub(crate) fn tx_counter_deltas_for(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::TxCounterDelta>> {
if !H::COUNTS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.counter_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let (max_rows, max_bytes) = self.effective_projection_limits();
let deltas = crate::counter::diff_counters_bytes(
&self.hook, targets, max_rows, max_bytes, &key_typed, old, new,
)?;
crate::counter::into_tx_counter_deltas(targets, deltas)
}
pub(crate) fn tx_claims_for(
&self,
name: &'static str,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::ClaimOp>> {
if !H::CLAIMS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.claim_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let (max_rows, max_bytes) = self.effective_projection_limits();
let delta = crate::unique::claim_delta_bytes(
&self.hook, targets, max_rows, max_bytes, &key_typed, old, new,
)?;
crate::unique::into_tx_claim_ops(targets, name, key, delta)
}
pub(crate) fn tx_shard(&self, sid: usize) -> VarMapShard<'_, K, H> {
VarMapShard {
tree: self,
held: Held {
locked: self.lock_shard(sid),
sid,
events: Vec::new(),
watch_wakes: WatchWakeBatch::default(),
publish_watches: true,
},
}
}
pub(crate) fn tx_release_shard(
shard: VarMapShard<'_, K, H>,
) -> (Vec<<Self as Family>::Owned>, WatchWakeBatch) {
let VarMapShard { held, .. } = shard;
let Held {
locked,
events,
watch_wakes,
..
} = held;
locked.release();
(events, watch_wakes)
}
pub(crate) fn tx_reserve_gsns(&self, count: usize) -> u64 {
self.engine
.gsn()
.fetch_add(count as u64, std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn tx_encrypted(&self) -> bool {
#[cfg(feature = "encryption")]
{
self.engine.config().encryption_key.is_some()
}
#[cfg(not(feature = "encryption"))]
{
false
}
}
pub(crate) fn tx_durable_recovered_gsn(&self, sid: usize) -> u64 {
self.engine.shards()[sid].durable_recovered_gsn()
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> ShardSource for VarMap<K, H> {
type Key = K;
type Val = VarMapIterVal;
fn shard_count(&self) -> usize {
self.engine.shards().len()
}
fn next_batch(
&self,
shard: usize,
after: Option<&K>,
dir: Dir,
n: usize,
range_start: Bound<OrdKey<K>>,
range_end: Bound<OrdKey<K>>,
) -> Vec<(K, VarMapIterVal)> {
let mut out = Vec::with_capacity(n);
let mut cursor: Option<K> = after.copied();
loop {
let snapshots: Vec<(K, DiskLoc)> = {
let index = sync::lock(&self.indexes[shard]);
let Some(ki) = index.key_index() else {
return out;
};
let keys = key_index_batch(ki, cursor.as_ref(), dir, n, range_start, range_end);
keys.into_iter()
.filter_map(|k| index.get(&k).map(|disk| (k, *disk)))
.collect()
};
if snapshots.is_empty() {
return out; }
let scanned = snapshots.len();
let last_scanned = snapshots.last().map(|(k, _)| *k);
for (key, loc) in snapshots {
match self.try_get(&key) {
Ok(Some(value)) => out.push((key, VarMapIterVal { loc, value })),
Ok(None) => {}
Err(e) => {
tracing::warn!(
key = ?key.as_bytes(),
error = ?e,
"var_map iter: skipping entry — value read failed"
);
}
}
}
if out.len() >= n || scanned < n {
return out;
}
cursor = last_scanned;
}
}
fn try_next_batch(
&self,
shard: usize,
after: Option<&K>,
dir: Dir,
n: usize,
range_start: Bound<OrdKey<K>>,
range_end: Bound<OrdKey<K>>,
) -> DbResult<Vec<(K, VarMapIterVal)>> {
let mut out = Vec::with_capacity(inspection_batch_initial_capacity(n));
let mut cursor: Option<K> = after.copied();
loop {
let snapshots: Vec<(K, DiskLoc)> = {
let index = sync::lock(&self.indexes[shard]);
let Some(ki) = index.key_index() else {
return Ok(out);
};
let keys = key_index_batch(ki, cursor.as_ref(), dir, n, range_start, range_end);
keys.into_iter()
.filter_map(|k| index.get(&k).map(|disk| (k, *disk)))
.collect()
};
if snapshots.is_empty() {
return Ok(out);
}
let scanned = snapshots.len();
let last_scanned = snapshots.last().map(|(k, _)| *k);
for (key, loc) in snapshots {
match self.try_get(&key) {
Ok(Some(value)) => out.push((key, VarMapIterVal { loc, value })),
Ok(None) => {}
Err(error) => return Err(error),
}
}
if out.len() >= n || scanned < n {
return Ok(out);
}
cursor = last_scanned;
}
}
fn next_keys_batch(
&self,
shard: usize,
after: Option<&K>,
dir: Dir,
n: usize,
range_start: Bound<OrdKey<K>>,
range_end: Bound<OrdKey<K>>,
) -> Vec<K> {
let index = sync::lock(&self.indexes[shard]);
let Some(ki) = index.key_index() else {
return Vec::new();
};
key_index_batch(ki, after, dir, n, range_start, range_end)
}
}
fn inspection_batch_initial_capacity(requested: usize) -> usize {
requested.min(1024)
}
impl<'a, K, H> VarMapIterView<'a, K, H>
where
K: Key + Send + Sync + Hash + Eq,
H: WriteHook<K>,
{
pub fn iter(&self) -> VarMapValueIter<'a, K, H> {
VarMapValueIter {
inner: self.inner.iter(),
}
}
pub fn range(&self, start: &K, end: &K) -> VarMapValueIter<'a, K, H> {
VarMapValueIter {
inner: self.inner.range(start, end),
}
}
pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> VarMapValueIter<'a, K, H> {
VarMapValueIter {
inner: self.inner.range_bounds(start, end),
}
}
pub fn prefix_iter(&self, prefix: &[u8]) -> VarMapValueIter<'a, K, H> {
VarMapValueIter {
inner: self.inner.prefix_iter(prefix),
}
}
pub fn paginate(&self, after: Option<&K>, limit: usize) -> Page<K, ByteView> {
let iter = match (after, dir_from(self.map.reversed)) {
(None, _) => self.iter(),
(Some(k), Dir::Asc) => self.range_bounds(Bound::Excluded(k), Bound::Unbounded),
(Some(k), Dir::Desc) => self.range_bounds(Bound::Unbounded, Bound::Excluded(k)),
};
let items: Vec<_> = iter.take(limit).collect();
let next = if items.len() == limit {
items.last().map(|(k, _)| *k)
} else {
None
};
Page { items, next }
}
pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
self.inner.keys()
}
pub fn keys_range(&self, start: &K, end: &K) -> impl Iterator<Item = K> + '_ {
self.inner.keys_range(start, end)
}
pub fn keys_range_bounds(
&self,
start: Bound<&K>,
end: Bound<&K>,
) -> impl Iterator<Item = K> + '_ {
self.inner.keys_range_bounds(start, end)
}
pub fn keys_prefix(&self, prefix: &[u8]) -> impl Iterator<Item = K> + '_ {
self.inner.keys_prefix(prefix)
}
pub fn keys_paginate(&self, after: Option<&K>, limit: usize) -> KeyPage<K> {
self.inner.keys_paginate(after, limit)
}
pub fn retain(&self, mut f: impl FnMut(&K, &ByteView) -> bool) -> DbResult<()> {
for (key, entry) in self.inner.iter() {
if f(&key, &entry.value) {
continue;
}
match self.map.compare_delete_if_loc(&key, entry.loc) {
Ok(()) => {}
Err(DbError::CasMismatch) | Err(DbError::KeyNotFound) => {
tracing::debug!("retain: compare_delete skipped (concurrent change)");
}
Err(e) => return Err(e),
}
}
Ok(())
}
}
pub struct VarMapIterView<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
map: &'a VarMap<K, H>,
inner: MapIterView<'a, VarMap<K, H>>,
}
pub struct VarMapValueIter<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> {
inner: MapIter<'a, VarMap<K, H>>,
}
#[doc(hidden)]
pub struct VarMapTryIter<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> {
inner: TryMapIter<'a, VarMap<K, H>>,
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> Iterator for VarMapTryIter<'_, K, H> {
type Item = DbResult<(K, ByteView)>;
fn next(&mut self) -> Option<Self::Item> {
self.inner
.next()
.map(|row| row.map(|(key, entry)| (key, entry.value)))
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> Iterator for VarMapValueIter<'_, K, H> {
type Item = (K, ByteView);
fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|(k, entry)| (k, entry.value))
}
}
#[cfg(feature = "replication")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> crate::replication::ReplicationTarget
for VarMap<K, H>
{
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::from_bytes(key);
let disk = DiskLoc::new(
file_id,
(entry_offset + size_of::<crate::entry::EntryHeader>() as u64 + size_of::<K>() as u64)
as u32,
header.value_len,
);
if header.is_tombstone() {
let old = sync::lock(&self.indexes[self.shard_for(&key)]).remove(&key);
match old {
Some(old_disk) => {
self.entry_count.fetch_sub(1, AtomicOrdering::Relaxed);
Ok(ApplyOutcome::TombstoneRemoved(old_disk))
}
None => Ok(ApplyOutcome::Inserted), }
} else {
let old = sync::lock(&self.indexes[self.shard_for(&key)]).upsert(key, disk);
match old {
Some(old_disk) => Ok(ApplyOutcome::Replaced(old_disk)),
None => {
self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
Ok(ApplyOutcome::Inserted)
}
}
}
}
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>()
}
fn mark_replication_follower(&self) {
VarMap::mark_replication_follower(self);
}
}
#[cfg(feature = "replication")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
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.engine.shards().clone(),
consumers,
self.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.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.engine.shards().clone(),
registry,
size_of::<K>() as u16,
signal,
)
}
}
#[cfg(feature = "replication")]
impl<K, H> VarMap<K, H>
where
K: Key + Send + Sync + Hash + Eq + '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))
}
}
pub struct VarMapShard<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
tree: &'a VarMap<K, H>,
pub(crate) held: Held<'a, VarMap<K, H>>,
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMapShard<'_, K, H> {
pub fn put(&mut self, key: &K, value: &[u8]) -> DbResult<bool> {
self.check_shard(key)?;
let tree = self.tree;
let sid = self.held.sid;
Ctx::with_shard(&mut self.held, tree, sid, |locked, ev_out| {
let (inner, index) = locked;
let old = if (H::NEEDS_OLD_VALUE
&& (H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active())))
|| H::CLAIMS
{
if let Some(disk) = index.get(key) {
Some(tree.read_value_locked_result(disk, sid as u8, inner)?)
} else {
None
}
} else {
None
};
tree.claim_guard_bytes(key, old.as_deref(), Some(value))?;
let expires_at = tree.deadline_for_bytes(key, value);
let existed = tree.put_locked(sid, inner, index, key, value, expires_at)?;
let ev = if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((*key, old, Some(VarHookNew::Borrowed(value))))
} else {
None
};
*ev_out = ev;
Ok(existed)
})
}
pub fn insert(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
self.check_shard(key)?;
let tree = self.tree;
let sid = self.held.sid;
Ctx::with_shard(&mut self.held, tree, sid, |locked, ev_out| {
let (inner, index) = locked;
if (H::EMITS || (H::CLAIMS && tree.hook.claim_targets().is_some()))
&& index.get(key).is_none()
{
tree.claim_guard_bytes(key, None, Some(value))?;
}
let expires_at = tree.deadline_for_bytes(key, value);
tree.insert_locked(sid, inner, index, key, value, expires_at)?;
let ev = if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((*key, None, Some(VarHookNew::Borrowed(value))))
} else {
None
};
*ev_out = ev;
Ok(())
})
}
pub fn delete(&mut self, key: &K) -> DbResult<bool> {
self.check_shard(key)?;
let tree = self.tree;
let sid = self.held.sid;
Ctx::with_shard(&mut self.held, tree, sid, |locked, ev_out| {
let (inner, index) = locked;
let old = if (H::NEEDS_OLD_VALUE
&& (H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active())))
|| H::CLAIMS
{
if let Some(disk) = index.get(key) {
Some(tree.read_value_locked_result(disk, sid as u8, inner)?)
} else {
None
}
} else {
None
};
tree.claim_guard_bytes(key, old.as_deref(), None)?;
let existed = tree.delete_locked(sid, inner, index, key)?;
let ev =
if existed && (H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active())) {
Some((*key, old, None))
} else {
None
};
*ev_out = ev;
Ok(existed)
})
}
pub fn try_get(&self, key: &K) -> DbResult<Option<ByteView>> {
self.check_shard(key)?;
match self.held.locked.1.get(key) {
Some(d) => {
let disk = *d;
let view = self.tree.read_value_locked_result(
&disk,
self.held.sid as u8,
&self.held.locked.0,
)?;
Ok(Some(view))
}
None => Ok(None),
}
}
pub fn get_or_err(&self, key: &K) -> DbResult<ByteView> {
self.check_shard(key)?;
let disk = *self.held.locked.1.get(key).ok_or(DbError::KeyNotFound)?;
self.tree
.read_value_locked_result(&disk, self.held.sid as u8, &self.held.locked.0)
}
pub fn try_contains(&self, key: &K) -> DbResult<bool> {
self.check_shard(key)?;
Ok(self.held.locked.1.contains_key(key))
}
pub fn update(
&mut self,
key: &K,
f: impl FnOnce(&[u8]) -> ByteView,
) -> DbResult<Option<ByteView>> {
self.update_inner(key, |b| Ok(Some(f(b))), false)
}
pub fn fetch_update(
&mut self,
key: &K,
f: impl FnOnce(&[u8]) -> ByteView,
) -> DbResult<Option<ByteView>> {
self.update_inner(key, |b| Ok(Some(f(b))), true)
}
pub(crate) fn update_inner(
&mut self,
key: &K,
f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
return_old: bool,
) -> DbResult<Option<ByteView>> {
self.check_shard(key)?;
let tree = self.tree;
let sid = self.held.sid;
Ctx::with_shard(&mut self.held, tree, sid, |locked, ev_out| {
let (inner, index) = locked;
match tree.try_update_locked(sid, inner, index, key, f)? {
UpdateOutcome::Missing => Ok(None),
UpdateOutcome::Unchanged(current) => Ok(Some(current)),
UpdateOutcome::Updated { old, new } => {
let (result, ev) =
if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
let event_old = if H::NEEDS_OLD_VALUE {
Some(old.clone())
} else {
None
};
let result = Some(if return_old { old } else { new.clone() });
(
result,
Some((*key, event_old, Some(VarHookNew::Owned(new)))),
)
} else {
(Some(if return_old { old } else { new }), None)
};
{
*ev_out = ev;
Ok(result)
}
}
}
})
}
fn check_shard(&self, key: &K) -> DbResult<()> {
if self.tree.shard_for(key) != self.held.sid {
return Err(DbError::ShardMismatch);
}
Ok(())
}
}
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMapShard<'_, K, H> {
pub(crate) fn tx_arm_gsn(&mut self, gsn: u64) {
self.held.locked.0.arm_gsn(gsn);
}
pub(crate) fn tx_clear_armed_gsn(&mut self) {
self.held.locked.0.clear_armed_gsn();
}
pub(crate) fn tx_fence(&mut self, gsn: u64) -> u64 {
self.held.locked.0.fence(gsn)
}
pub(crate) fn tx_buffered_events(&self) -> usize {
self.held.events.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Config;
use crate::compaction::compact_shard;
use crate::config::CacheConfig;
use crate::decoded_cache::DecodedCache;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use tempfile::tempdir;
fn decoded_cache_on() -> DecodedCache<Vec<u8>> {
DecodedCache::new(&CacheConfig {
max_size: 1 << 20,
estimated_items: 64,
})
}
fn counting_decode(counter: &AtomicUsize) -> impl Fn(&ByteView) -> DbResult<Vec<u8>> + '_ {
move |b| {
counter.fetch_add(1, AtomicOrdering::Relaxed);
Ok(b.as_bytes().to_vec())
}
}
#[derive(Default)]
struct CountingHook<const NEEDS_INIT: bool, const NEEDS_OLD: bool> {
writes: AtomicUsize,
writes_with_old: AtomicUsize,
inits: AtomicUsize,
last_write_new: crate::sync::Mutex<Option<Vec<u8>>>,
last_init_value: crate::sync::Mutex<Option<Vec<u8>>>,
}
impl<const NEEDS_INIT: bool, const NEEDS_OLD: bool> WriteHook<[u8; 8]>
for CountingHook<NEEDS_INIT, NEEDS_OLD>
{
const NEEDS_OLD_VALUE: bool = NEEDS_OLD;
const NEEDS_INIT: bool = NEEDS_INIT;
fn on_write(&self, _key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
self.writes.fetch_add(1, AtomicOrdering::Relaxed);
if old.is_some() {
self.writes_with_old.fetch_add(1, AtomicOrdering::Relaxed);
}
*crate::sync::lock(&self.last_write_new) = new.map(<[u8]>::to_vec);
}
fn on_init(&self, _key: &[u8; 8], value: &[u8]) {
self.inits.fetch_add(1, AtomicOrdering::Relaxed);
*crate::sync::lock(&self.last_init_value) = Some(value.to_vec());
}
}
fn open_test_map_hooked<const NEEDS_INIT: bool, const NEEDS_OLD: bool>(
dir: &std::path::Path,
hook: CountingHook<NEEDS_INIT, NEEDS_OLD>,
) -> VarMap<[u8; 8], CountingHook<NEEDS_INIT, NEEDS_OLD>> {
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 8192;
cfg.write_buffer_size = 8192;
cfg.compaction_threshold = 0.0;
VarMap::open_hooked(dir, cfg, hook).expect("open hooked test map")
}
fn open_test_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 8192;
cfg.write_buffer_size = 8192;
cfg.compaction_threshold = 0.0;
VarMap::open(dir, cfg).expect("open test map")
}
#[test]
fn inspection_batch_initial_capacity_is_bounded() {
assert_eq!(inspection_batch_initial_capacity(32), 32);
assert_eq!(inspection_batch_initial_capacity(usize::MAX), 1024);
}
#[test]
fn strict_range_bounds_surfaces_stale_and_terminates() {
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.iterable = true;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).unwrap();
let key = 1u64.to_be_bytes();
map.put(&key, b"payload").unwrap();
let shard = &map.engine.shards()[map.shard_for(&key)];
let mut inner = shard.lock();
inner.immutable = Vec::new();
inner.active.file_id = u32::MAX;
drop(inner);
let mut rows =
map.inspection_range_bounds_dir(Bound::Unbounded, Bound::Unbounded, Dir::Asc, 1);
let first = rows.next();
assert!(
matches!(first, Some(Err(DbError::StaleDiskLoc))),
"expected stale strict row, got {first:?}"
);
assert!(
rows.next().is_none(),
"strict iterator must terminate after an error"
);
}
#[test]
fn open_resolves_hints_true_for_var() {
let dir = tempfile::tempdir().unwrap();
let cfg = crate::Config::test_no_hints();
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).unwrap();
assert_eq!(map.config().hints, Some(true));
}
#[test]
fn put_reports_whether_key_existed() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 42u64.to_be_bytes();
assert!(!map.put(&key, b"v1").expect("put"));
assert!(map.put(&key, b"v2").expect("put"));
map.flush_buffers().expect("flush");
assert!(map.put(&key, b"v3").expect("put"));
assert!(!map.put(&7u64.to_be_bytes(), b"x").expect("put"));
}
#[test]
fn for_each_visits_all_entries() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
map.put(&1u64.to_be_bytes(), b"one").unwrap();
map.put(&2u64.to_be_bytes(), b"two").unwrap();
map.put(&3u64.to_be_bytes(), b"three").unwrap();
let mut seen: std::collections::HashMap<[u8; 8], Vec<u8>> =
std::collections::HashMap::new();
map.for_each(|key, value| {
seen.insert(key, value.to_vec());
});
assert_eq!(seen.len(), 3);
assert_eq!(seen[&1u64.to_be_bytes()], b"one".to_vec());
assert_eq!(seen[&2u64.to_be_bytes()], b"two".to_vec());
assert_eq!(seen[&3u64.to_be_bytes()], b"three".to_vec());
}
#[test]
fn for_each_while_stops_reading_on_break() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
for id in 1..=6u64 {
map.put(&id.to_be_bytes(), b"v").unwrap();
}
let before = test_value_reads();
let mut visited = 0;
map.for_each_while(|row| {
row.expect("readable row");
visited += 1;
ControlFlow::Break(())
});
assert_eq!(visited, 1);
assert_eq!(test_value_reads() - before, 1);
let before = test_value_reads();
let mut visited = 0;
map.for_each_while(|_| {
visited += 1;
ControlFlow::Continue(())
});
assert_eq!(visited, 6);
assert_eq!(test_value_reads() - before, 6);
}
fn put_until_compactable(map: &VarMap<[u8; 8]>, key: [u8; 8]) -> DiskLoc {
map.put(&key, &[0u8; 256]).expect("first put");
let snap = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
for i in 1..65u8 {
map.put(&key, &[i; 256]).expect("overwrite");
}
map.put(&key, b"final-value-payload-XX").expect("final put");
snap
}
#[test]
fn var_map_len_tracks_mutations() {
let dir = tempdir().unwrap();
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), Config::balanced().build()).unwrap();
assert_eq!(map.len(), 0);
assert!(map.is_empty());
let k1 = [1u8; 8];
let k2 = [2u8; 8];
map.put(&k1, b"a").unwrap();
assert_eq!(map.len(), 1);
map.put(&k2, b"b").unwrap();
assert_eq!(map.len(), 2);
map.put(&k1, b"c").unwrap();
assert_eq!(map.len(), 2);
map.delete(&k1).unwrap();
assert_eq!(map.len(), 1);
map.delete(&k2).unwrap();
assert_eq!(map.len(), 0);
assert!(map.is_empty());
map.delete(&k1).unwrap();
assert_eq!(map.len(), 0);
}
#[test]
fn var_map_len_survives_reopen() {
let dir = tempdir().unwrap();
{
let map: VarMap<[u8; 8]> =
VarMap::open(dir.path(), Config::balanced().build()).unwrap();
for i in 0u64..50 {
map.put(&i.to_le_bytes(), b"val").unwrap();
}
assert_eq!(map.len(), 50);
}
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), Config::balanced().build()).unwrap();
assert_eq!(map.len(), 50);
}
#[test]
fn read_value_cached_inner_returns_stale_after_compaction() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 7u64.to_be_bytes();
let snap = put_until_compactable(&map, key);
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let _ = compact_shard(shard, &map, 0.0).expect("compaction");
match map.read_value_cached_inner(&snap, shard_id) {
Err(DbError::StaleDiskLoc) => {}
Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
}
}
#[test]
fn get_during_compaction_returns_some() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 11u64.to_be_bytes();
let _snap = put_until_compactable(&map, key);
let shard = &map.engine.shards()[0];
let _ = compact_shard(shard, &map, 0.0).expect("compaction");
let v = map.get(&key).expect("post-compaction get");
assert_eq!(v.as_bytes(), b"final-value-payload-XX");
}
#[test]
fn get_or_read_block_returns_stale_for_unknown_file_id() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
match map.get_or_read_block(0, 9999, 0) {
Err(DbError::StaleDiskLoc) => {}
Ok(_) => panic!("expected StaleDiskLoc, got Ok"),
Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
}
}
#[test]
fn retry_limit_returns_none_on_persistent_stale() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 99u64.to_be_bytes();
map.put(&key, b"payload").expect("put");
{
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let mut inner = shard.lock();
inner.immutable = Vec::new();
inner.active.file_id = u32::MAX;
}
assert!(
map.get(&key).is_none(),
"MAX_STALE_RETRIES must terminate the retry loop and return None"
);
}
#[test]
fn var_map_replay_init_fires_on_init_per_live_key_raw() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());
for i in 0u64..5 {
map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
}
map.delete(&3u64.to_be_bytes()).expect("delete");
map.hook.writes.store(0, AtomicOrdering::Relaxed);
map.hook.inits.store(0, AtomicOrdering::Relaxed);
map.replay_init();
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 4);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
}
#[test]
fn var_map_replay_init_no_hook_is_noop() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
for i in 0u64..3 {
map.put(&i.to_be_bytes(), &[i as u8; 8]).expect("put");
}
map.replay_init();
assert!(map.get(&0u64.to_be_bytes()).is_some());
}
#[test]
fn var_map_migrate_keep_fires_on_init_not_on_write_raw() {
use crate::MigrateAction;
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());
for i in 0u64..4 {
map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
}
map.hook.writes.store(0, AtomicOrdering::Relaxed);
map.hook.inits.store(0, AtomicOrdering::Relaxed);
let mutated = map.migrate(|_, _| MigrateAction::Keep).expect("migrate");
assert_eq!(mutated, 0);
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 4);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
}
#[test]
fn var_map_migrate_update_fires_on_init_with_new_value_raw() {
use crate::MigrateAction;
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());
let key = 42u64.to_be_bytes();
map.put(&key, b"old-value").expect("put");
map.hook.writes.store(0, AtomicOrdering::Relaxed);
map.hook.inits.store(0, AtomicOrdering::Relaxed);
let new = ByteView::new(b"new-value");
let mutated = map
.migrate(move |_, _| MigrateAction::Update(new.clone()))
.expect("migrate");
assert_eq!(mutated, 1);
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 1);
let g = crate::sync::lock(&map.hook.last_init_value);
assert_eq!(g.as_deref(), Some(b"new-value".as_ref()));
drop(g);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
assert_eq!(map.get(&key).unwrap().as_bytes(), b"new-value");
}
#[test]
fn var_map_migrate_delete_fires_no_hooks_raw() {
use crate::MigrateAction;
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());
let key = 7u64.to_be_bytes();
map.put(&key, b"x").expect("put");
map.hook.writes.store(0, AtomicOrdering::Relaxed);
map.hook.inits.store(0, AtomicOrdering::Relaxed);
let mutated = map.migrate(|_, _| MigrateAction::Delete).expect("migrate");
assert_eq!(mutated, 1);
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
assert!(map.get(&key).is_none());
}
#[test]
fn var_map_public_put_still_fires_on_write_once() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());
map.put(&1u64.to_be_bytes(), b"v").expect("put");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
}
#[test]
fn var_map_migrate_no_init_hook_is_silent_for_keep_and_update() {
use crate::MigrateAction;
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<false, false>(dir.path(), CountingHook::default());
for i in 0u64..3 {
map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
}
map.hook.writes.store(0, AtomicOrdering::Relaxed);
map.hook.inits.store(0, AtomicOrdering::Relaxed);
map.migrate(|_, _| MigrateAction::Keep)
.expect("migrate keep");
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
let new = ByteView::new(b"new");
map.migrate(move |_, _| MigrateAction::Update(new.clone()))
.expect("migrate update");
assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
}
#[test]
fn var_map_atomic_fires_hooks() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
let key = 1u64.to_be_bytes();
map.atomic(&key, |shard| {
shard.put(&key, b"a")?; shard.delete(&key)?; Ok(())
})
.expect("atomic");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1); }
#[test]
fn var_map_atomic_fires_for_applied_on_err() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
let key = 1u64.to_be_bytes();
let r: DbResult<()> = map.atomic(&key, |shard| {
shard.put(&key, b"x")?;
Err(DbError::KeyNotFound)
});
assert!(r.is_err());
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
}
#[test]
fn var_map_atomic_same_key_old_from_disk() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<false, true>(dir.path(), CountingHook::default());
let key = 1u64.to_be_bytes();
map.atomic(&key, |shard| {
shard.put(&key, b"first")?;
shard.put(&key, b"second")?;
Ok(())
})
.expect("atomic");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1); assert_eq!(
crate::sync::lock(&map.hook.last_write_new).clone(),
Some(b"second".to_vec())
);
}
#[test]
fn var_map_shard_put_returns_existed() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
let key = 1u64.to_be_bytes();
let (fresh, overwrite) = map
.atomic(&key, |shard| {
let fresh = shard.put(&key, b"a")?;
let overwrite = shard.put(&key, b"b")?;
Ok((fresh, overwrite))
})
.expect("atomic");
assert!(!fresh);
assert!(overwrite);
}
#[test]
fn compare_delete_match_mismatch_absent() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let k = [1u8; 8];
map.put(&k, b"hello").unwrap();
assert!(matches!(
map.compare_delete(&k, b"WRONG"),
Err(DbError::CasMismatch)
));
assert_eq!(map.get(&k).as_deref(), Some(&b"hello"[..]));
assert!(map.compare_delete(&k, b"hello").is_ok());
assert!(map.get(&k).is_none());
assert!(matches!(
map.compare_delete(&k, b"hello"),
Err(DbError::KeyNotFound)
));
}
#[test]
fn compare_delete_hook_needs_old_true_gets_old() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<false, true>(dir.path(), CountingHook::default());
let k = [2u8; 8];
map.put(&k, b"payload").unwrap();
let with_old_before = map.hook.writes_with_old.load(AtomicOrdering::Relaxed);
map.compare_delete(&k, b"payload").unwrap();
assert_eq!(
map.hook.writes_with_old.load(AtomicOrdering::Relaxed) - with_old_before,
1
);
}
#[test]
fn compare_delete_hook_needs_old_false_gets_none() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<false, false>(dir.path(), CountingHook::default());
let k = [3u8; 8];
map.put(&k, b"payload").unwrap();
let writes_before = map.hook.writes.load(AtomicOrdering::Relaxed);
let with_old_before = map.hook.writes_with_old.load(AtomicOrdering::Relaxed);
map.compare_delete(&k, b"payload").unwrap();
assert_eq!(
map.hook.writes.load(AtomicOrdering::Relaxed) - writes_before,
1
);
assert_eq!(
map.hook.writes_with_old.load(AtomicOrdering::Relaxed) - with_old_before,
0
);
}
#[test]
fn read_value_locked_result_ok_from_write_buf() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 1u64.to_be_bytes();
let payload = b"in-write-buffer-value";
map.put(&key, payload).expect("put");
let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let inner = shard.lock();
let v = map
.read_value_locked_result(&disk, shard_id, &inner)
.expect("write-buf read must succeed");
assert_eq!(v.as_bytes(), payload);
}
#[test]
fn read_value_locked_result_ok_from_disk_immutable() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 2u64.to_be_bytes();
let payload = b"single-block-immutable";
map.put(&key, payload).expect("put");
for i in 100u64..135 {
map.put(&i.to_be_bytes(), &[i as u8; 256]).expect("rotator");
}
let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let inner = shard.lock();
assert_ne!(
disk.file_id, inner.active.file_id,
"test setup failed: key=2 entry is still in the active file's write buffer",
);
let v = map
.read_value_locked_result(&disk, shard_id, &inner)
.expect("disk read must succeed");
assert_eq!(v.as_bytes(), payload);
}
#[test]
fn read_value_locked_result_propagates_stale_disk_loc() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 3u64.to_be_bytes();
let snap = put_until_compactable(&map, key);
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let _ = compact_shard(shard, &map, 0.0).expect("compaction");
let inner = shard.lock();
match map.read_value_locked_result(&snap, shard_id, &inner) {
Err(DbError::StaleDiskLoc) => {}
Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
}
}
#[test]
fn warmup_loads_both_blocks_of_a_straddling_value() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 1 << 20;
cfg.write_buffer_size = 1 << 20;
cfg.compaction_threshold = 0.0;
cfg.block_cache.max_size = 1 << 20;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");
map.put(&0u64.to_be_bytes(), &vec![0x11u8; 3000]).unwrap();
map.put(&1u64.to_be_bytes(), &vec![0x5Au8; 1500]).unwrap();
map.put(&2u64.to_be_bytes(), &vec![0x77u8; 30_000]).unwrap();
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
let straddler = *sync::lock(&map.indexes[0])
.get(&1u64.to_be_bytes())
.expect("indexed");
let start = (straddler.offset & 4095) as usize;
assert!(
!straddler.is_value_cache_routed() && start + straddler.len as usize > 4096,
"test precondition: value must straddle a block boundary and stay block-cache routed \
(start={start}, len={})",
straddler.len
);
let tail = *sync::lock(&map.indexes[0])
.get(&2u64.to_be_bytes())
.expect("indexed");
assert!(
tail.is_value_cache_routed(),
"test precondition: the tail value must be value-cache routed so warmup skips it"
);
map.warmup().unwrap();
let block_offset = straddler.offset as u64 & !4095;
let bkey = |off| BlockKey {
shard_id: 0,
file_id: straddler.file_id,
block_offset: off,
};
assert!(
map.block_cache.get(&bkey(block_offset)).is_some(),
"warmup must cache the value's first block"
);
assert!(
map.block_cache.get(&bkey(block_offset + 4096)).is_some(),
"warmup must also cache the block the value spills into"
);
}
#[test]
fn large_value_with_first_block_cached_uses_fallback() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.block_cache.max_size = 1 << 20;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");
let key = 42u64.to_be_bytes();
let payload: Vec<u8> = (0..20_000u32).map(|i| i as u8).collect();
map.put(&key, &payload).expect("put large");
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
let start = (disk.offset & 4095) as usize;
assert!(
start + disk.len as usize > 8192,
"test precondition: large value must span >2 blocks",
);
let block_offset = disk.offset as u64 & !4095;
let shard_id = map.shard_for(&key) as u8;
let cache_key = BlockKey {
shard_id,
file_id: disk.file_id,
block_offset,
};
map.block_cache
.insert(cache_key, Arc::new(AlignedBuf::zeroed(4096)));
assert!(
map.block_cache.get(&cache_key).is_some(),
"cache must be enabled for warm-cache scenario",
);
let v = map
.read_value_cached_inner(&disk, shard_id)
.expect("large value must read via locked fallback");
assert_eq!(v.as_bytes(), payload.as_slice());
}
#[test]
fn extract_from_block_single_block() {
let mut block = AlignedBuf::zeroed(4096);
for (i, byte) in block.iter_mut().enumerate() {
*byte = i as u8;
}
let v = VarMap::<[u8; 8]>::extract_from_block(&block, 100, 50, || {
panic!("next_block must not be called for single-block reads")
})
.expect("ok");
let expected: Vec<u8> = (100u8..150u8).collect();
assert_eq!(v.as_bytes(), expected.as_slice());
}
#[test]
fn extract_from_block_two_blocks_exact() {
let mut first = AlignedBuf::zeroed(4096);
for byte in first.iter_mut() {
*byte = 0xAA;
}
let mut second = AlignedBuf::zeroed(4096);
for byte in second.iter_mut() {
*byte = 0xBB;
}
let v = VarMap::<[u8; 8]>::extract_from_block(&first, 4095, 4097, || Ok(Arc::new(second)))
.expect("ok");
let bytes = v.as_bytes();
assert_eq!(bytes.len(), 4097);
assert_eq!(bytes[0], 0xAA);
assert_eq!(bytes[1], 0xBB);
assert_eq!(bytes[4096], 0xBB);
}
#[test]
fn extract_from_block_two_blocks_partial() {
let mut first = AlignedBuf::zeroed(4096);
for byte in first.iter_mut() {
*byte = 0x11;
}
let mut second = AlignedBuf::zeroed(4096);
for byte in second.iter_mut() {
*byte = 0x22;
}
let v = VarMap::<[u8; 8]>::extract_from_block(&first, 4000, 200, || Ok(Arc::new(second)))
.expect("ok");
let bytes = v.as_bytes();
assert_eq!(bytes.len(), 200);
assert!(bytes[..96].iter().all(|b| *b == 0x11));
assert!(bytes[96..].iter().all(|b| *b == 0x22));
}
fn open_large_value_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
VarMap::open(dir, cfg).expect("open large-value test map")
}
fn build_large_payload(seed: u8) -> Vec<u8> {
(0..50_000u32)
.map(|i| (i as u8).wrapping_add(seed))
.collect()
}
#[test]
fn large_value_read_via_locked_fallback() {
let dir = tempdir().unwrap();
let map = open_large_value_map(dir.path());
let key = 100u64.to_be_bytes();
let payload = build_large_payload(0);
map.put(&key, &payload).expect("put large");
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
let v = map
.get(&key)
.expect("read must succeed via locked fallback");
assert_eq!(v.as_bytes(), payload.as_slice());
}
#[cfg(feature = "encryption")]
#[test]
fn varmap_strict_reads_and_mutators_surface_read_fault() {
use crate::test_faults::{big_value, corrupt_tags};
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.encryption_key = Some([7u8; 32]);
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");
let key = 1u64.to_be_bytes();
map.put(&key, &big_value(0xCD)).unwrap();
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
corrupt_tags(&map.engine.shard_dirs()[0]);
assert!(map.get(&key).is_none());
assert!(matches!(
map.get_or_err(&key),
Err(DbError::EncryptionError(_))
));
assert!(matches!(
map.try_get(&key),
Err(DbError::EncryptionError(_))
));
assert!(matches!(
map.cas(&key, b"x", b"y"),
Err(DbError::EncryptionError(_))
));
assert!(matches!(
map.compare_delete(&key, b"x"),
Err(DbError::EncryptionError(_))
));
assert!(matches!(
map.update(&key, |_| ByteView::new(b"z")),
Err(DbError::EncryptionError(_))
));
let absent = 999u64.to_be_bytes();
assert!(matches!(map.try_get(&absent), Ok(None)));
assert!(matches!(map.get_or_err(&absent), Err(DbError::KeyNotFound)));
}
#[cfg(feature = "encryption")]
#[test]
fn large_value_read_encrypted() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.encryption_key = Some([7u8; 32]);
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");
let key = 101u64.to_be_bytes();
let payload = build_large_payload(0xAB);
map.put(&key, &payload).expect("put encrypted large");
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
let v = map
.get(&key)
.expect("encrypted large read must succeed via pread_value_encrypted");
assert_eq!(v.as_bytes(), payload.as_slice());
}
#[test]
fn large_value_stale_disk_loc_deterministic() {
let dir = tempdir().unwrap();
let map = open_large_value_map(dir.path());
let key = 102u64.to_be_bytes();
let payload = build_large_payload(0x42);
map.put(&key, &payload).expect("first large put");
let snap = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate after large put");
for i in 1..20u8 {
map.put(&key, &[i; 256]).expect("overwrite");
}
map.put(&key, b"live-after-compaction").expect("final put");
let shard_id = map.shard_for(&key) as u8;
let shard = &map.engine.shards()[shard_id as usize];
let _ = compact_shard(shard, &map, 0.0).expect("compaction");
match map.read_value_cached_inner(&snap, shard_id) {
Err(DbError::StaleDiskLoc) => {}
Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
}
let v = map
.get(&key)
.expect("public path must retry and return live value");
assert_eq!(v.as_bytes(), b"live-after-compaction");
}
#[test]
fn read_value_locked_result_ok_from_cache_single_block() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 8192;
cfg.write_buffer_size = 8192;
cfg.compaction_threshold = 0.0;
cfg.block_cache.max_size = 1 << 20;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");
let key = 9u64.to_be_bytes();
let payload = b"small-single-block-value";
map.put(&key, payload).expect("put");
for i in 100u64..135 {
map.put(&i.to_be_bytes(), &[i as u8; 256]).expect("rotator");
}
let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
let start = (disk.offset & 4095) as usize;
let len = disk.len as usize;
assert!(
start + len <= 4096,
"test precondition: value must fit in a single block",
);
let block_offset = disk.offset as u64 & !4095;
let shard_id = map.shard_for(&key) as u8;
let cache_key = BlockKey {
shard_id,
file_id: disk.file_id,
block_offset,
};
let block = map
.get_or_read_block(shard_id, disk.file_id, block_offset)
.expect("read block");
map.block_cache.insert(cache_key, block);
assert!(
map.block_cache.get(&cache_key).is_some(),
"cache must contain the block for step 2 to fire",
);
let shard = &map.engine.shards()[shard_id as usize];
let inner = shard.lock();
assert_ne!(
disk.file_id, inner.active.file_id,
"test setup failed: key entry is still in the active write buffer",
);
let v = map
.read_value_locked_result(&disk, shard_id, &inner)
.expect("cache-hit read must succeed");
assert_eq!(v.as_bytes(), payload);
}
#[test]
fn var_map_get_with_file_id_above_u16() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.block_cache.max_size = 1 << 20;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");
let shard = &map.engine.shards()[0];
shard.set_next_file_id(70_000);
shard.rotate_active_for_test(8).expect("first rotate");
assert!(
shard.active_file_id() >= 70_000,
"active_file_id should be >= 70_000 after rotation"
);
let key = 42u64.to_be_bytes();
let value = vec![0xC3u8; 512];
map.put(&key, &value).expect("put");
{
let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
assert!(
disk.file_id > u16::MAX as u32,
"DiskLoc.file_id must be above u16::MAX, got {}",
disk.file_id,
);
}
shard.flush().expect("flush");
shard.rotate_active_for_test(8).expect("second rotate");
let got = map.get(&key).expect("get must return Some");
assert_eq!(got.as_bytes(), value.as_slice());
}
#[test]
fn var_map_iter_reads_values_after_unlock() {
let dir = tempfile::tempdir().unwrap();
let cfg = Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.build();
let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
m.put(&[1, 0], b"alpha").unwrap();
m.put(&[2, 0], b"beta").unwrap();
let got: Vec<(u8, Vec<u8>)> = m
.iter_view()
.unwrap()
.iter()
.map(|(k, v)| (k[0], v.as_ref().to_vec()))
.collect();
assert_eq!(got, [(2, b"beta".to_vec()), (1, b"alpha".to_vec())]);
}
#[test]
fn var_map_iter_view_none_unless_iterable() {
let dir = tempfile::tempdir().unwrap();
let m = open_test_map(dir.path());
assert!(m.iter_view().is_none());
}
#[test]
fn var_map_keys_only_skips_value_read() {
let dir = tempfile::tempdir().unwrap();
let cfg = Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.build();
let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
m.put(&[1, 0], b"alpha").unwrap();
m.put(&[2, 0], b"beta").unwrap();
let keys: Vec<u8> = m.iter_view().unwrap().keys().map(|k| k[0]).collect();
assert_eq!(keys, vec![2, 1], "DESC keys");
let page = m.iter_view().unwrap().keys_paginate(None, 1);
assert_eq!(page.keys, vec![[2, 0]]);
assert_eq!(page.next, Some([2, 0]));
}
#[test]
fn var_map_retain_deletes_rejected_entries() {
let dir = tempfile::tempdir().unwrap();
let cfg = Config::balanced().shard_count(1).iterable(true).build();
let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
m.put(&[1, 0], b"odd").unwrap();
m.put(&[2, 0], b"even").unwrap();
m.put(&[3, 0], b"odd2").unwrap();
m.iter_view()
.unwrap()
.retain(|_k, v| v.as_ref() == b"even")
.unwrap();
let got: Vec<u8> = m.iter_view().unwrap().iter().map(|(k, _)| k[0]).collect();
assert_eq!(got, vec![2]);
assert_eq!(m.len(), 1);
}
#[cfg(feature = "encryption")]
#[test]
fn var_map_iter_skips_unreadable_key() {
use crate::test_faults::{big_value, corrupt_tags};
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.encryption_key = Some([7u8; 32]);
cfg.iterable = true;
let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
m.put(&k2, &big_value(0x22)).unwrap();
m.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
corrupt_tags(&m.engine.shard_dirs()[0]);
m.put(&k1, b"readable").unwrap();
let got: Vec<u8> = m.iter_view().unwrap().iter().map(|(k, _)| k[7]).collect();
assert_eq!(got, vec![1]);
assert_eq!(m.len(), 2, "unreadable entry must not be deleted by iter");
assert!(matches!(m.try_get(&k2), Err(DbError::EncryptionError(_))));
let mut keys: Vec<u8> = m.iter_view().unwrap().keys().map(|k| k[7]).collect();
keys.sort_unstable();
assert_eq!(keys, vec![1, 2], "keys-only returns the unreadable key too");
}
#[cfg(feature = "encryption")]
#[test]
fn var_map_retain_skips_unreadable_entry() {
use crate::test_faults::{big_value, corrupt_tags};
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.encryption_key = Some([7u8; 32]);
cfg.iterable = true;
let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
m.put(&k2, &big_value(0x22)).unwrap();
m.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
corrupt_tags(&m.engine.shard_dirs()[0]);
m.put(&k1, b"keep").unwrap();
m.iter_view().unwrap().retain(|_k, _v| false).unwrap();
assert_eq!(m.len(), 1);
assert!(m.get(&k1).is_none());
assert!(matches!(m.try_get(&k2), Err(DbError::EncryptionError(_))));
}
#[cfg(feature = "encryption")]
#[test]
fn var_map_iter_advances_past_full_unreadable_batch() {
use crate::test_faults::{big_value, corrupt_tags};
let dir = tempfile::tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 4 * 1024 * 1024;
cfg.write_buffer_size = 4 * 1024 * 1024;
cfg.compaction_threshold = 0.0;
cfg.encryption_key = Some([9u8; 32]);
cfg.iterable = true;
let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");
for i in 2u64..=66 {
m.put(&i.to_be_bytes(), &big_value(i as u8)).unwrap();
}
m.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
corrupt_tags(&m.engine.shard_dirs()[0]);
let readable = 1u64.to_be_bytes();
m.put(&readable, b"alive").unwrap();
let got: Vec<[u8; 8]> = m.iter_view().unwrap().iter().map(|(k, _)| k).collect();
assert_eq!(
got,
vec![readable],
"iter must scan past a full unreadable batch"
);
assert_eq!(m.iter_view().unwrap().keys().count(), 66);
}
fn open_value_cache_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
cfg.value_cache.max_size = 4 << 20;
VarMap::open(dir, cfg).expect("open value-cache test map")
}
fn large_payload(seed: u8) -> Vec<u8> {
(0..50_000u32)
.map(|i| (i as u8).wrapping_add(seed))
.collect()
}
#[test]
fn map_large_value_miss_then_cached() {
let dir = tempdir().unwrap();
let map = open_value_cache_map(dir.path());
let key = 1u64.to_be_bytes();
let payload = large_payload(0x33);
map.put(&key, &payload).unwrap();
assert_eq!(map.get(&key).expect("first get").as_bytes(), &payload[..]);
let disk = {
let idx = sync::lock(&map.indexes[0]);
*idx.get(&key).unwrap()
};
let vkey = crate::value_cache::ValueKey {
shard_id: 0,
file_id: disk.file_id,
offset: disk.offset,
};
assert!(
map.value_cache.get(&vkey).is_some(),
"value cached after read"
);
assert_eq!(map.get(&key).expect("second get").as_bytes(), &payload[..]);
}
#[test]
fn map_large_value_read_then_edit_returns_new() {
let dir = tempdir().unwrap();
let map = open_value_cache_map(dir.path());
let key = 7u64.to_be_bytes();
let v1 = large_payload(0x11);
let v2 = large_payload(0x22);
map.put(&key, &v1).unwrap();
assert_eq!(map.get(&key).expect("get v1").as_bytes(), &v1[..]);
map.put(&key, &v2).unwrap();
assert_eq!(map.get(&key).expect("get v2").as_bytes(), &v2[..]);
}
#[test]
fn map_large_value_read_then_compact_returns_live() {
let dir = tempdir().unwrap();
let map = open_value_cache_map(dir.path());
let key = 102u64.to_be_bytes();
let payload = large_payload(0x42);
map.put(&key, &payload).unwrap();
assert_eq!(map.get(&key).expect("initial get").as_bytes(), &payload[..]);
map.engine.shards()[0]
.rotate_active_for_test(8)
.expect("rotate");
for i in 1..20u8 {
map.put(&key, &[i; 256]).unwrap();
}
map.put(&key, b"live-after-compaction").unwrap();
let _ = compact_shard(&map.engine.shards()[0], &map, 0.0).expect("compaction");
assert_eq!(
map.get(&key).expect("post-compaction get").as_bytes(),
b"live-after-compaction"
);
}
#[test]
fn var_map_shard_update_and_fetch_update() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 1u64.to_be_bytes();
map.put(&key, b"a").unwrap();
let (upd, fetched, missing) = map
.atomic(&key, |shard| {
let upd = shard.update(&key, |old| {
let mut v = old.to_vec();
v.push(b'b');
ByteView::from(v.as_slice())
})?;
let fetched = shard.fetch_update(&key, |old| {
let mut v = old.to_vec();
v.push(b'c');
ByteView::from(v.as_slice())
})?;
let missing = shard.update(&2u64.to_be_bytes(), |old| ByteView::from(old))?;
Ok((upd, fetched, missing))
})
.expect("atomic");
assert_eq!(upd.as_deref(), Some(&b"ab"[..]));
assert_eq!(fetched.as_deref(), Some(&b"ab"[..]));
assert!(missing.is_none());
assert_eq!(map.get(&key).as_deref(), Some(&b"abc"[..]));
}
#[test]
fn var_map_atomic_update_fires_hook() {
let dir = tempdir().unwrap();
let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
let key = 1u64.to_be_bytes();
map.atomic(&key, |shard| {
shard.put(&key, b"hello")?;
shard.update(&key, |old| {
let mut v = old.to_vec();
v.push(b'!');
ByteView::from(v.as_slice())
})?;
Ok(())
})
.expect("atomic");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1);
assert_eq!(map.get(&key).as_deref(), Some(&b"hello!"[..]));
}
#[test]
fn straddle_read_2block_direct_io_returns_full_value() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 1 << 20;
cfg.write_buffer_size = 8192;
cfg.direct_io = true;
cfg.compaction_threshold = 0.0;
let map = VarMap::<[u8; 8]>::open(dir.path(), cfg).expect("open");
let v1: Vec<u8> = (0..3000).map(|i| ((i % 251) + 1) as u8).collect();
let v2: Vec<u8> = (0..3000).map(|i| ((i % 241) + 2) as u8).collect();
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
map.put(&k1, &v1).expect("put v1");
map.put(&k2, &v2).expect("put v2");
map.flush_buffers().expect("flush");
assert_eq!(
map.get(&k2).as_deref(),
Some(&v2[..]),
"straddling value must read back byte-for-byte (tail was zeroed before the fix)"
);
assert_eq!(map.get(&k1).as_deref(), Some(&v1[..]));
}
#[test]
fn straddle_read_large_value_direct_io_try_get_ok() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 1 << 20;
cfg.write_buffer_size = 16 * 1024;
cfg.direct_io = true;
cfg.compaction_threshold = 0.0;
let map = VarMap::<[u8; 8]>::open(dir.path(), cfg).expect("open");
let v1: Vec<u8> = (0..3000).map(|i| ((i % 251) + 1) as u8).collect();
let v2: Vec<u8> = (0..9000).map(|i| ((i % 239) + 3) as u8).collect();
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
map.put(&k1, &v1).expect("put v1");
map.put(&k2, &v2).expect("put v2");
map.flush_buffers().expect("flush");
let got = map
.try_get(&k2)
.expect("try_get must not error for a live key");
assert_eq!(
got.as_deref(),
Some(&v2[..]),
"straddling large value must read back byte-for-byte"
);
}
#[cfg(feature = "encryption")]
#[test]
fn straddle_read_2block_encrypted_returns_full_value() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 1 << 20;
cfg.write_buffer_size = 8192;
cfg.encryption_key = Some([7u8; 32]);
cfg.compaction_threshold = 0.0;
let map = VarMap::<[u8; 8]>::open(dir.path(), cfg).expect("open enc");
let v1: Vec<u8> = (0..3000).map(|i| ((i % 251) + 1) as u8).collect();
let v2: Vec<u8> = (0..3000).map(|i| ((i % 241) + 2) as u8).collect();
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
map.put(&k1, &v1).expect("put v1");
map.put(&k2, &v2).expect("put v2");
map.flush_buffers().expect("flush");
let got = map
.try_get(&k2)
.expect("encrypted straddle read must not error");
assert_eq!(got.as_deref(), Some(&v2[..]));
assert_eq!(map.get(&k1).as_deref(), Some(&v1[..]));
}
#[test]
fn straddle_read_2block_cache_hit_returns_full_value() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 1 << 20;
cfg.write_buffer_size = 8192;
cfg.direct_io = true;
cfg.block_cache.max_size = 1 << 20;
cfg.compaction_threshold = 0.0;
let map = VarMap::<[u8; 8]>::open(dir.path(), cfg).expect("open");
let v1: Vec<u8> = (0..3000).map(|i| ((i % 251) + 1) as u8).collect();
let v2: Vec<u8> = (0..3000).map(|i| ((i % 241) + 2) as u8).collect();
let k1 = 1u64.to_be_bytes();
let k2 = 2u64.to_be_bytes();
map.put(&k1, &v1).expect("put v1");
map.put(&k2, &v2).expect("put v2");
map.flush_buffers().expect("flush");
let shard_id = map.shard_for(&k2) as u8;
let key0 = BlockKey {
shard_id,
file_id: 1,
block_offset: 0,
};
let (blk0, _is_full) = map.engine.shards()[shard_id as usize]
.read_block(1, 0)
.expect("read block 0");
map.block_cache.insert(key0, Arc::new(blk0));
assert_eq!(
map.get(&k2).as_deref(),
Some(&v2[..]),
"straddling value via cache-hit path must read back byte-for-byte (tail was zeroed before the fix)"
);
}
#[test]
fn map_try_get_decoded_hit_does_not_decode() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 3u64.to_be_bytes();
map.put(&key, b"payload").unwrap();
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
let first = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.expect("read")
.expect("present");
let second = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.expect("read")
.expect("present");
assert_eq!(first.as_slice(), b"payload");
assert_eq!(second.as_slice(), b"payload");
assert_eq!(decodes.load(AtomicOrdering::Relaxed), 1);
assert!(Arc::ptr_eq(&first, &second)); }
#[test]
fn map_try_get_decoded_covers_both_byte_routes() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
cfg.max_file_size = 128 * 1024;
cfg.write_buffer_size = 128 * 1024;
cfg.compaction_threshold = 0.0;
let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open test map");
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
let small_key = 21u64.to_be_bytes();
map.put(&small_key, &[7u8; 64]).unwrap();
let big_key = 22u64.to_be_bytes();
map.put(&big_key, &vec![9u8; 12 * 1024]).unwrap();
let routed = |k: &[u8; 8]| {
let shard = map.shard_for(k);
let index = crate::sync::lock(&map.indexes[shard]);
index.get(k).expect("indexed").is_value_cache_routed()
};
assert!(
!routed(&small_key),
"test setup: the small value must take the block-cache route"
);
assert!(
routed(&big_key),
"test setup: the large value must take the value-cache route"
);
let small = map
.try_get_decoded(&small_key, &cache, counting_decode(&decodes))
.expect("small read")
.expect("present");
assert_eq!(small.as_slice(), &vec![7u8; 64][..]);
let big = map
.try_get_decoded(&big_key, &cache, counting_decode(&decodes))
.expect("big read")
.expect("present");
assert_eq!(big.as_slice(), &vec![9u8; 12 * 1024][..]);
let before = decodes.load(AtomicOrdering::Relaxed);
map.try_get_decoded(&small_key, &cache, counting_decode(&decodes))
.expect("small re-read")
.expect("present");
map.try_get_decoded(&big_key, &cache, counting_decode(&decodes))
.expect("big re-read")
.expect("present");
assert_eq!(decodes.load(AtomicOrdering::Relaxed), before);
}
#[test]
fn map_try_get_decoded_absent_is_none() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
assert!(
map.try_get_decoded(&77u64.to_be_bytes(), &cache, counting_decode(&decodes))
.expect("absent is not an error")
.is_none()
);
assert_eq!(decodes.load(AtomicOrdering::Relaxed), 0);
}
#[test]
fn map_try_get_decoded_sees_a_write() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 5u64.to_be_bytes();
map.put(&key, b"v1").unwrap();
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
let held = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.unwrap();
assert_eq!(held.as_slice(), b"v1");
map.put(&key, b"v2").unwrap();
let after = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.unwrap();
assert_eq!(after.as_slice(), b"v2");
assert_eq!(held.as_slice(), b"v1");
assert_eq!(decodes.load(AtomicOrdering::Relaxed), 2);
}
#[test]
fn map_try_get_decoded_correct_after_compaction() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 9u64.to_be_bytes();
for _ in 0..64 {
map.put(&key, b"churn-churn-churn-churn").unwrap();
}
map.put(&key, b"final-value-payload-XX").unwrap();
map.flush_buffers().unwrap();
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
map.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.unwrap();
let shard = &map.engine.shards()[0];
let _ = compact_shard(shard, &map, 0.0).expect("compaction");
let v = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.expect("post-compaction read")
.expect("present");
assert_eq!(v.as_slice(), b"final-value-payload-XX");
}
#[test]
fn map_try_get_decoded_absent_after_delete_redecodes_on_reinsert() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 41u64.to_be_bytes();
map.put(&key, b"v1").unwrap();
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
let v1 = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.unwrap();
assert_eq!(v1.as_slice(), b"v1");
assert!(map.delete(&key).unwrap());
assert!(
map.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.is_none()
);
map.put(&key, b"v2").unwrap();
let before = decodes.load(AtomicOrdering::Relaxed);
let v2 = map
.try_get_decoded(&key, &cache, counting_decode(&decodes))
.unwrap()
.unwrap();
assert_eq!(v2.as_slice(), b"v2");
assert_eq!(
decodes.load(AtomicOrdering::Relaxed) - before,
1,
"re-insert must decode again"
);
}
#[test]
fn map_try_get_decoded_returns_stale_after_retry_limit() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 51u64.to_be_bytes();
map.put(&key, b"payload").expect("put");
{
let shard_id = map.shard_for(&key);
let shard = &map.engine.shards()[shard_id];
let mut inner = shard.lock();
inner.immutable = Vec::new();
inner.active.file_id = u32::MAX;
}
let cache = decoded_cache_on();
let decodes = AtomicUsize::new(0);
match map.try_get_decoded(&key, &cache, counting_decode(&decodes)) {
Err(DbError::StaleDiskLoc) => {}
Ok(Some(v)) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_slice()),
Ok(None) => panic!("expected StaleDiskLoc, got Ok(None)"),
Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
}
}
#[test]
fn map_try_get_decoded_does_not_cache_failed_decodes() {
let dir = tempdir().unwrap();
let map = open_test_map(dir.path());
let key = 31u64.to_be_bytes();
map.put(&key, b"whatever").unwrap();
let cache = decoded_cache_on();
let attempts = AtomicUsize::new(0);
let failing = |_b: &ByteView| -> DbResult<Vec<u8>> {
attempts.fetch_add(1, AtomicOrdering::Relaxed);
Err(DbError::CorruptedEntry { offset: 0 })
};
for _ in 0..2 {
assert!(matches!(
map.try_get_decoded(&key, &cache, failing),
Err(DbError::CorruptedEntry { .. })
));
}
assert_eq!(attempts.load(AtomicOrdering::Relaxed), 2);
}
}