use std::marker::PhantomData;
use std::ops::Bound;
#[cfg(any(test, feature = "test-faults"))]
use std::sync::atomic::AtomicBool;
#[cfg(any(feature = "armour", test))]
use std::sync::atomic::AtomicU8;
use std::sync::atomic::{AtomicU64, Ordering};
#[cfg(feature = "armour")]
use std::sync::{Arc, Weak};
#[cfg(feature = "armour")]
use crate::armour::db::ManagedLsmLifetime;
use crate::codec::Codec;
use crate::error::{DbError, DbResult};
use crate::hook::TypedWriteHook;
use crate::lsm::config::{LsmConfig, PersistPolicy};
use crate::lsm::iter::{LsmIter, LsmKeyIter, TryLsmIter, TryLsmKeyIter};
use crate::lsm::map_fjall;
use crate::map_index::dir_from;
use crate::sync::{self, Mutex, RwLock};
#[cfg(feature = "armour")]
use crate::sync::{MutexGuard, RwLockReadGuard, RwLockWriteGuard};
use crate::{CollectionMeta, Dir, Key, NoHook};
#[cfg(test)]
pub(crate) static LSM_LAST_PERIODIC: AtomicU8 = AtomicU8::new(255);
#[cfg(test)]
pub(crate) static LSM_LAST_DURABLE: AtomicU8 = AtomicU8::new(255);
#[cfg(feature = "armour")]
pub(crate) struct ManagedLsmCleanMarker {
lifetime: Weak<ManagedLsmLifetime>,
state: AtomicU8,
activity: RwLock<()>,
transition: Mutex<()>,
#[cfg(any(test, feature = "test-faults"))]
fail_next_demotion: AtomicBool,
#[cfg(test)]
publication_gate: std::sync::Mutex<Option<Arc<CleanPublicationGate>>>,
#[cfg(test)]
mutation_entry_gate: std::sync::Mutex<Option<Arc<MutationEntryGate>>>,
#[cfg(test)]
finalization_entry_gate: std::sync::Mutex<Option<Arc<FinalizationEntryGate>>>,
#[cfg(test)]
mutation_operation_gate: std::sync::Mutex<Option<Arc<MutationOperationGate>>>,
#[cfg(test)]
panic_after_clean_persistence: AtomicBool,
}
#[cfg(feature = "armour")]
const LSM_DIRTY: u8 = 0;
#[cfg(feature = "armour")]
const LSM_FINALIZING: u8 = 1;
#[cfg(feature = "armour")]
const LSM_CLEAN: u8 = 2;
#[cfg(feature = "armour")]
pub(crate) struct LsmFinalizationGuard<'a> {
marker: &'a ManagedLsmCleanMarker,
_transition: MutexGuard<'a, ()>,
_activity: RwLockWriteGuard<'a, ()>,
was_clean: bool,
clean_may_have_persisted: bool,
finished: bool,
}
#[cfg(feature = "armour")]
struct ManagedLsmMutationGuard<'a> {
_activity: RwLockReadGuard<'a, ()>,
_lifetime: Arc<ManagedLsmLifetime>,
}
struct LsmMutationGuard<'a> {
#[cfg(feature = "armour")]
_managed: Option<ManagedLsmMutationGuard<'a>>,
}
#[cfg(all(feature = "armour", test))]
struct CleanPublicationGate {
state: std::sync::Mutex<(bool, bool)>,
cv: std::sync::Condvar,
}
#[cfg(all(feature = "armour", test))]
impl CleanPublicationGate {
fn wait_until_blocked(&self) {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
let mut state = self.state.lock().expect("clean publication gate poisoned");
while !state.0 {
let remaining = deadline
.checked_duration_since(std::time::Instant::now())
.expect("finalizer never reached clean publication gate");
let (next, timeout) = self
.cv
.wait_timeout(state, remaining)
.expect("clean publication gate poisoned");
assert!(
!timeout.timed_out() || next.0,
"finalizer never reached clean publication gate"
);
state = next;
}
}
fn block(&self) {
let mut state = self.state.lock().expect("clean publication gate poisoned");
state.0 = true;
self.cv.notify_all();
while !state.1 {
state = self
.cv
.wait(state)
.expect("clean publication gate poisoned");
}
}
fn release(&self) {
let mut state = self.state.lock().expect("clean publication gate poisoned");
state.1 = true;
self.cv.notify_all();
}
}
#[cfg(all(feature = "armour", test))]
pub(crate) struct CleanPublicationGateGuard<'a> {
marker: &'a ManagedLsmCleanMarker,
gate: Arc<CleanPublicationGate>,
}
#[cfg(all(feature = "armour", test))]
impl CleanPublicationGateGuard<'_> {
pub(crate) fn wait_until_blocked(&self) {
self.gate.wait_until_blocked();
}
pub(crate) fn release(&self) {
self.gate.release();
}
}
#[cfg(all(feature = "armour", test))]
impl Drop for CleanPublicationGateGuard<'_> {
fn drop(&mut self) {
self.gate.release();
let mut active = self
.marker
.publication_gate
.lock()
.expect("clean publication gate poisoned");
if active
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &self.gate))
{
*active = None;
}
}
}
#[cfg(all(feature = "armour", test))]
struct MutationEntryGate {
entered: std::sync::Mutex<bool>,
cv: std::sync::Condvar,
}
#[cfg(all(feature = "armour", test))]
struct FinalizationEntryGate {
entered: std::sync::Mutex<bool>,
cv: std::sync::Condvar,
}
#[cfg(all(feature = "armour", test))]
struct MutationOperationGate {
state: std::sync::Mutex<(bool, bool)>,
cv: std::sync::Condvar,
}
#[cfg(all(feature = "armour", test))]
impl MutationOperationGate {
fn block(&self) {
let mut state = self.state.lock().expect("mutation operation gate poisoned");
state.0 = true;
self.cv.notify_all();
while !state.1 {
state = self
.cv
.wait(state)
.expect("mutation operation gate poisoned");
}
}
fn wait_until_blocked(&self) {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
let mut state = self.state.lock().expect("mutation operation gate poisoned");
while !state.0 {
let remaining = deadline
.checked_duration_since(std::time::Instant::now())
.expect("writer never reached managed LSM operation gate");
let (next, timeout) = self
.cv
.wait_timeout(state, remaining)
.expect("mutation operation gate poisoned");
assert!(
!timeout.timed_out() || next.0,
"writer never reached managed LSM operation gate"
);
state = next;
}
}
fn release(&self) {
let mut state = self.state.lock().expect("mutation operation gate poisoned");
state.1 = true;
self.cv.notify_all();
}
}
#[cfg(all(feature = "armour", test))]
pub(crate) struct MutationOperationGateGuard<'a> {
marker: &'a ManagedLsmCleanMarker,
gate: Arc<MutationOperationGate>,
}
#[cfg(all(feature = "armour", test))]
impl MutationOperationGateGuard<'_> {
pub(crate) fn wait_until_blocked(&self) {
self.gate.wait_until_blocked();
}
pub(crate) fn release(&self) {
self.gate.release();
}
}
#[cfg(all(feature = "armour", test))]
impl Drop for MutationOperationGateGuard<'_> {
fn drop(&mut self) {
self.gate.release();
let mut active = self
.marker
.mutation_operation_gate
.lock()
.expect("mutation operation gate poisoned");
if active
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &self.gate))
{
*active = None;
}
}
}
#[cfg(all(feature = "armour", test))]
impl MutationEntryGate {
fn enter(&self) {
*self.entered.lock().expect("mutation entry gate poisoned") = true;
self.cv.notify_all();
}
fn wait_until_entered(&self) {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
let mut entered = self.entered.lock().expect("mutation entry gate poisoned");
while !*entered {
let remaining = deadline
.checked_duration_since(std::time::Instant::now())
.expect("writer never entered managed LSM mutation");
let (next, timeout) = self
.cv
.wait_timeout(entered, remaining)
.expect("mutation entry gate poisoned");
assert!(
!timeout.timed_out() || *next,
"writer never entered managed LSM mutation"
);
entered = next;
}
}
}
#[cfg(all(feature = "armour", test))]
impl FinalizationEntryGate {
fn enter(&self) {
*self
.entered
.lock()
.expect("finalization entry gate poisoned") = true;
self.cv.notify_all();
}
fn wait_until_entered(&self) {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
let mut entered = self
.entered
.lock()
.expect("finalization entry gate poisoned");
while !*entered {
let remaining = deadline
.checked_duration_since(std::time::Instant::now())
.expect("finalizer never entered managed LSM finalization");
let (next, timeout) = self
.cv
.wait_timeout(entered, remaining)
.expect("finalization entry gate poisoned");
assert!(
!timeout.timed_out() || *next,
"finalizer never entered managed LSM finalization"
);
entered = next;
}
}
}
#[cfg(all(feature = "armour", test))]
pub(crate) struct MutationEntryGateGuard<'a> {
marker: &'a ManagedLsmCleanMarker,
gate: Arc<MutationEntryGate>,
}
#[cfg(all(feature = "armour", test))]
impl MutationEntryGateGuard<'_> {
pub(crate) fn wait_until_entered(&self) {
self.gate.wait_until_entered();
}
}
#[cfg(all(feature = "armour", test))]
impl Drop for MutationEntryGateGuard<'_> {
fn drop(&mut self) {
let mut active = self
.marker
.mutation_entry_gate
.lock()
.expect("mutation entry gate poisoned");
if active
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &self.gate))
{
*active = None;
}
}
}
#[cfg(all(feature = "armour", test))]
pub(crate) struct FinalizationEntryGateGuard<'a> {
marker: &'a ManagedLsmCleanMarker,
gate: Arc<FinalizationEntryGate>,
}
#[cfg(all(feature = "armour", test))]
impl FinalizationEntryGateGuard<'_> {
pub(crate) fn wait_until_entered(&self) {
self.gate.wait_until_entered();
}
}
#[cfg(all(feature = "armour", test))]
impl Drop for FinalizationEntryGateGuard<'_> {
fn drop(&mut self) {
let mut active = self
.marker
.finalization_entry_gate
.lock()
.expect("finalization entry gate poisoned");
if active
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &self.gate))
{
*active = None;
}
}
}
#[cfg(feature = "armour")]
impl ManagedLsmCleanMarker {
pub(crate) fn new(lifetime: Weak<ManagedLsmLifetime>) -> Self {
Self {
lifetime,
state: AtomicU8::new(LSM_DIRTY),
activity: RwLock::new(()),
transition: Mutex::new(()),
#[cfg(any(test, feature = "test-faults"))]
fail_next_demotion: AtomicBool::new(false),
#[cfg(test)]
publication_gate: std::sync::Mutex::new(None),
#[cfg(test)]
mutation_entry_gate: std::sync::Mutex::new(None),
#[cfg(test)]
finalization_entry_gate: std::sync::Mutex::new(None),
#[cfg(test)]
mutation_operation_gate: std::sync::Mutex::new(None),
#[cfg(test)]
panic_after_clean_persistence: AtomicBool::new(false),
}
}
pub(crate) fn begin_finalization(&self) -> LsmFinalizationGuard<'_> {
#[cfg(test)]
self.note_finalization_entry_for_test();
let activity = sync::write(&self.activity);
let transition = sync::lock(&self.transition);
let was_clean = self.state.load(Ordering::Acquire) == LSM_CLEAN;
self.state.store(LSM_FINALIZING, Ordering::Release);
LsmFinalizationGuard {
marker: self,
_transition: transition,
_activity: activity,
was_clean,
clean_may_have_persisted: false,
finished: false,
}
}
fn persist_dirty_while_locked(&self) -> DbResult<()> {
#[cfg(any(test, feature = "test-faults"))]
if self.fail_next_demotion.swap(false, Ordering::SeqCst) {
return Err(DbError::Io(std::io::Error::other(
"injected managed LSM clean-marker demotion failure",
)));
}
let lifetime = self.lifetime.upgrade().ok_or(DbError::Internal(
"managed LSM tree outlived its Db clean marker",
))?;
lifetime.db_info.update(|info| info.clean = false)?;
self.state.store(LSM_DIRTY, Ordering::Release);
Ok(())
}
fn demote_before_mutation(&self) -> DbResult<ManagedLsmMutationGuard<'_>> {
#[cfg(test)]
self.note_mutation_entry_for_test();
let activity = sync::read(&self.activity);
let lifetime = self.lifetime.upgrade().ok_or(DbError::Internal(
"managed LSM tree outlived its Db clean marker",
))?;
if self.state.load(Ordering::Acquire) == LSM_DIRTY {
return Ok(ManagedLsmMutationGuard {
_activity: activity,
_lifetime: lifetime,
});
}
let _guard = sync::lock(&self.transition);
match self.state.load(Ordering::Acquire) {
LSM_DIRTY => Ok(()),
LSM_CLEAN => self.persist_dirty_while_locked(),
LSM_FINALIZING => Err(DbError::Internal(
"managed LSM clean marker remained finalizing after transition lock acquisition",
)),
_ => Err(DbError::Internal("invalid managed LSM clean marker state")),
}?;
Ok(ManagedLsmMutationGuard {
_activity: activity,
_lifetime: lifetime,
})
}
#[cfg(any(test, feature = "test-faults"))]
pub(crate) fn fail_next_demotion_for_test(&self) {
self.fail_next_demotion.store(true, Ordering::SeqCst);
}
#[cfg(test)]
fn wait_clean_publication_gate(&self) {
let gate = self
.publication_gate
.lock()
.expect("clean publication gate poisoned")
.clone();
if let Some(gate) = gate {
gate.block();
}
}
#[cfg(test)]
pub(crate) fn hold_clean_publication_for_test(&self) -> CleanPublicationGateGuard<'_> {
let gate = Arc::new(CleanPublicationGate {
state: std::sync::Mutex::new((false, false)),
cv: std::sync::Condvar::new(),
});
let mut active = self
.publication_gate
.lock()
.expect("clean publication gate poisoned");
assert!(
active.is_none(),
"clean publication gate already registered"
);
*active = Some(Arc::clone(&gate));
CleanPublicationGateGuard { marker: self, gate }
}
#[cfg(test)]
fn note_mutation_entry_for_test(&self) {
let gate = self
.mutation_entry_gate
.lock()
.expect("mutation entry gate poisoned")
.clone();
if let Some(gate) = gate {
gate.enter();
}
}
#[cfg(test)]
fn note_finalization_entry_for_test(&self) {
let gate = self
.finalization_entry_gate
.lock()
.expect("finalization entry gate poisoned")
.clone();
if let Some(gate) = gate {
gate.enter();
}
}
#[cfg(test)]
fn wait_before_keyspace_mutation_for_test(&self) {
let gate = self
.mutation_operation_gate
.lock()
.expect("mutation operation gate poisoned")
.clone();
if let Some(gate) = gate {
gate.block();
}
}
#[cfg(test)]
pub(crate) fn hold_mutation_entry_for_test(&self) -> MutationEntryGateGuard<'_> {
let gate = Arc::new(MutationEntryGate {
entered: std::sync::Mutex::new(false),
cv: std::sync::Condvar::new(),
});
let mut active = self
.mutation_entry_gate
.lock()
.expect("mutation entry gate poisoned");
assert!(active.is_none(), "mutation entry gate already registered");
*active = Some(Arc::clone(&gate));
MutationEntryGateGuard { marker: self, gate }
}
#[cfg(test)]
pub(crate) fn hold_finalization_entry_for_test(&self) -> FinalizationEntryGateGuard<'_> {
let gate = Arc::new(FinalizationEntryGate {
entered: std::sync::Mutex::new(false),
cv: std::sync::Condvar::new(),
});
let mut active = self
.finalization_entry_gate
.lock()
.expect("finalization entry gate poisoned");
assert!(
active.is_none(),
"finalization entry gate already registered"
);
*active = Some(Arc::clone(&gate));
FinalizationEntryGateGuard { marker: self, gate }
}
#[cfg(test)]
pub(crate) fn hold_mutation_operation_for_test(&self) -> MutationOperationGateGuard<'_> {
let gate = Arc::new(MutationOperationGate {
state: std::sync::Mutex::new((false, false)),
cv: std::sync::Condvar::new(),
});
let mut active = self
.mutation_operation_gate
.lock()
.expect("mutation operation gate poisoned");
assert!(
active.is_none(),
"mutation operation gate already registered"
);
*active = Some(Arc::clone(&gate));
MutationOperationGateGuard { marker: self, gate }
}
#[cfg(test)]
pub(crate) fn panic_after_clean_persistence_for_test(&self) {
self.panic_after_clean_persistence
.store(true, Ordering::SeqCst);
}
}
#[cfg(feature = "armour")]
impl LsmFinalizationGuard<'_> {
pub(crate) fn note_clean_persisted(&mut self) {
self.clean_may_have_persisted = true;
#[cfg(test)]
assert!(
!self
.marker
.panic_after_clean_persistence
.swap(false, Ordering::SeqCst),
"injected panic after managed LSM clean persistence"
);
}
pub(crate) fn publish_clean(&mut self) {
#[cfg(test)]
self.marker.wait_clean_publication_gate();
self.marker.state.store(LSM_CLEAN, Ordering::Release);
self.finished = true;
}
pub(crate) fn finish_dirty(&mut self) -> DbResult<()> {
let result = self.restore_dirty();
self.finished = true;
result
}
fn restore_dirty(&mut self) -> DbResult<()> {
if self.was_clean || self.clean_may_have_persisted {
if let Err(error) = self.marker.persist_dirty_while_locked() {
self.marker.state.store(LSM_CLEAN, Ordering::Release);
return Err(error);
}
} else {
self.marker.state.store(LSM_DIRTY, Ordering::Release);
}
Ok(())
}
}
#[cfg(feature = "armour")]
impl Drop for LsmFinalizationGuard<'_> {
fn drop(&mut self) {
if !self.finished {
if let Err(error) = self.restore_dirty() {
tracing::error!(error = %error, "failed to demote managed LSM clean marker during finalization unwind");
}
}
}
}
#[allow(dead_code)]
pub struct LsmTree<K, T, C, H = NoHook>
where
K: Key + Ord,
T: CollectionMeta<SelfId = K> + Send + Sync,
C: Codec<T> + Clone,
H: TypedWriteHook<K, T>,
{
db: fjall::Database,
ks: fjall::Keyspace,
codec: C,
hook: H,
config: LsmConfig,
reversed: bool,
persist: PersistPolicy,
stripes: Vec<Mutex<()>>,
count: AtomicU64,
#[cfg(feature = "armour")]
clean_marker: Option<Arc<ManagedLsmCleanMarker>>,
#[cfg(test)]
force_next_key_guard_error: AtomicBool,
_marker: PhantomData<fn() -> (K, T)>,
}
impl<K, T, C, H> LsmTree<K, T, C, H>
where
K: Key + Ord,
T: CollectionMeta<SelfId = K> + Send + Sync,
C: Codec<T> + Clone,
H: TypedWriteHook<K, T>,
{
pub fn open(db: fjall::Database, config: LsmConfig, codec: C, hook: H) -> DbResult<Self> {
Self::open_with_initial_count(db, config, codec, hook, None)
}
pub(crate) fn open_with_initial_count(
db: fjall::Database,
config: LsmConfig,
codec: C,
hook: H,
initial_count: Option<u64>,
) -> DbResult<Self> {
if H::WATCHES {
return Err(DbError::WatchUnsupported);
}
let memtable = config.memtable_size;
let ks = map_fjall(db.keyspace(T::NAME, move || {
fjall::KeyspaceCreateOptions::default().max_memtable_size(memtable)
}))?;
let count = match initial_count {
Some(c) => {
let empty = map_fjall(ks.is_empty())?;
if empty && c != 0 {
tracing::warn!(
collection = T::NAME,
seeded = c,
"persisted LSM count is non-zero over an empty keyspace — re-seeding to 0"
);
0
} else if !empty && c == 0 {
let actual = map_fjall(ks.len())? as u64;
tracing::warn!(
collection = T::NAME,
recovered = actual,
"persisted LSM count is 0 over a non-empty keyspace — rescanned"
);
actual
} else {
c
}
}
None => map_fjall(ks.len())? as u64,
};
let stripes = (0..config.lock_stripes.max(1))
.map(|_| Mutex::new(()))
.collect();
Ok(Self {
db,
ks,
codec,
hook,
reversed: config.reversed,
persist: config.persist_or_default(),
config,
stripes,
count: AtomicU64::new(count),
#[cfg(feature = "armour")]
clean_marker: None,
#[cfg(test)]
force_next_key_guard_error: AtomicBool::new(false),
_marker: PhantomData,
})
}
#[cfg(feature = "armour")]
pub(crate) fn open_with_persist(
db: fjall::Database,
config: LsmConfig,
persist: PersistPolicy,
codec: C,
hook: H,
initial_count: Option<u64>,
clean_marker: Arc<ManagedLsmCleanMarker>,
) -> DbResult<Self> {
let mut t = Self::open_with_initial_count(db, config, codec, hook, initial_count)?;
t.persist = persist;
t.clean_marker = Some(clean_marker);
Ok(t)
}
#[allow(dead_code)]
#[inline]
pub(crate) fn keyspace(&self) -> &fjall::Keyspace {
&self.ks
}
#[allow(dead_code)]
#[inline]
pub(crate) fn codec(&self) -> &C {
&self.codec
}
#[allow(dead_code)]
#[inline]
pub(crate) fn persist_handle(&self) -> (&fjall::Database, PersistPolicy) {
(&self.db, self.persist)
}
#[inline]
pub fn config(&self) -> &LsmConfig {
&self.config
}
#[inline]
pub fn shard_for(&self, key: &K) -> usize {
(xxhash_rust::xxh3::xxh3_64(key.as_bytes()) as usize) % self.stripes.len()
}
#[inline]
pub(crate) fn stripe_for(&self, key: &K) -> &Mutex<()> {
&self.stripes[self.shard_for(key)]
}
#[inline]
fn prepare_mutation(&self) -> DbResult<LsmMutationGuard<'_>> {
#[cfg(feature = "armour")]
if let Some(marker) = &self.clean_marker {
return Ok(LsmMutationGuard {
_managed: Some(marker.demote_before_mutation()?),
});
}
Ok(LsmMutationGuard {
#[cfg(feature = "armour")]
_managed: None,
})
}
#[inline]
fn before_keyspace_mutation(&self) {
#[cfg(all(feature = "armour", test))]
if let Some(marker) = &self.clean_marker {
marker.wait_before_keyspace_mutation_for_test();
}
}
#[cfg(any(test, feature = "test-faults"))]
pub(crate) fn fail_next_clean_marker_demotion_for_test(&self) {
#[cfg(feature = "armour")]
if let Some(marker) = &self.clean_marker {
marker.fail_next_demotion_for_test();
}
}
#[cfg(all(test, feature = "armour"))]
pub(crate) fn hold_clean_marker_publication_for_test(&self) -> CleanPublicationGateGuard<'_> {
self.clean_marker
.as_ref()
.expect("direct LSM trees do not have a managed clean marker")
.hold_clean_publication_for_test()
}
#[cfg(all(test, feature = "armour"))]
pub(crate) fn hold_clean_marker_mutation_entry_for_test(&self) -> MutationEntryGateGuard<'_> {
self.clean_marker
.as_ref()
.expect("direct LSM trees do not have a managed clean marker")
.hold_mutation_entry_for_test()
}
#[cfg(all(test, feature = "armour"))]
pub(crate) fn hold_clean_marker_finalization_entry_for_test(
&self,
) -> FinalizationEntryGateGuard<'_> {
self.clean_marker
.as_ref()
.expect("direct LSM trees do not have a managed clean marker")
.hold_finalization_entry_for_test()
}
#[cfg(all(test, feature = "armour"))]
pub(crate) fn hold_clean_marker_mutation_operation_for_test(
&self,
) -> MutationOperationGateGuard<'_> {
self.clean_marker
.as_ref()
.expect("direct LSM trees do not have a managed clean marker")
.hold_mutation_operation_for_test()
}
#[cfg(all(test, feature = "armour"))]
pub(crate) fn panic_after_clean_marker_persistence_for_test(&self) {
self.clean_marker
.as_ref()
.expect("direct LSM trees do not have a managed clean marker")
.panic_after_clean_persistence_for_test();
}
fn count_dec(&self) {
loop {
let cur = self.count.load(Ordering::Relaxed);
if cur == 0 {
tracing::warn!(
collection = T::NAME,
"lsm count underflow suppressed — the seeded count under-counts live rows"
);
return;
}
if self
.count
.compare_exchange_weak(cur, cur - 1, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
return;
}
}
}
#[inline]
pub(crate) fn get_raw(&self, key: &K) -> DbResult<Option<fjall::Slice>> {
map_fjall(self.ks.get(key.as_bytes()))
}
pub fn get(&self, key: &K) -> Option<T> {
let raw = self.ks.get(key.as_bytes()).ok()??;
self.codec.decode_from(raw.as_ref()).ok()
}
pub fn contains(&self, key: &K) -> bool {
self.ks.contains_key(key.as_bytes()).unwrap_or(false)
}
pub fn try_contains(&self, key: &K) -> DbResult<bool> {
map_fjall(self.ks.contains_key(key.as_bytes()))
}
pub(crate) fn put_locked(&self, key: &K, value: &T) -> DbResult<()> {
self.put_locked_inner::<true>(key, value)
}
pub(crate) fn put_locked_inner<const HOOKS: bool>(&self, key: &K, value: &T) -> DbResult<()> {
let fire = HOOKS && H::NEEDS_WRITE;
let mut buf = Vec::new();
self.codec.encode_to(value, &mut buf)?;
let prev = self.get_raw(key)?;
let old_t = if fire {
prev.as_ref()
.map(|b| self.codec.decode_from(b.as_ref()))
.transpose()?
} else {
None
};
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.insert(key.as_bytes(), buf.as_slice()))?;
if prev.is_none() {
self.count.fetch_add(1, Ordering::Relaxed);
}
if fire {
self.hook.on_write(key, old_t.as_ref(), Some(value));
}
Ok(())
}
pub(crate) fn insert_locked(&self, key: &K, value: &T) -> DbResult<()> {
if map_fjall(self.ks.contains_key(key.as_bytes()))? {
return Err(DbError::KeyExists);
}
let mut buf = Vec::new();
self.codec.encode_to(value, &mut buf)?;
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.insert(key.as_bytes(), buf.as_slice()))?;
self.count.fetch_add(1, Ordering::Relaxed);
if H::NEEDS_WRITE {
self.hook.on_write(key, None, Some(value));
}
Ok(())
}
pub(crate) fn delete_locked(&self, key: &K) -> DbResult<bool> {
self.delete_locked_inner::<true>(key)
}
pub(crate) fn delete_locked_inner<const HOOKS: bool>(&self, key: &K) -> DbResult<bool> {
let fire = HOOKS && H::NEEDS_WRITE;
let Some(prev) = self.get_raw(key)? else {
return Ok(false);
};
let old_t = if fire {
Some(self.codec.decode_from(prev.as_ref())?)
} else {
None
};
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.remove(key.as_bytes()))?;
self.count_dec();
if fire {
self.hook.on_write(key, old_t.as_ref(), None);
}
Ok(true)
}
pub fn put(&self, key: &K, value: &T) -> DbResult<()> {
let _g = sync::lock(self.stripe_for(key));
self.put_locked(key, value)
}
pub fn insert(&self, key: &K, value: &T) -> DbResult<()> {
let _g = sync::lock(self.stripe_for(key));
self.insert_locked(key, value)
}
pub fn delete(&self, key: &K) -> DbResult<bool> {
let _g = sync::lock(self.stripe_for(key));
self.delete_locked(key)
}
pub fn atomic<R>(
&self,
shard_key: &K,
f: impl FnOnce(&mut crate::lsm::shard::LsmShard<'_, K, T, C, H>) -> DbResult<R>,
) -> DbResult<R> {
let _atomic_depth = crate::AtomicDepthGuard::try_enter()?;
let stripe = self.shard_for(shard_key);
let _g = sync::lock(&self.stripes[stripe]);
let mut shard = crate::lsm::shard::LsmShard { tree: self, stripe };
f(&mut shard)
}
pub fn get_or_err(&self, key: &K) -> DbResult<T> {
let raw = self.get_raw(key)?.ok_or(DbError::KeyNotFound)?;
self.codec.decode_from(raw.as_ref())
}
pub fn try_get(&self, key: &K) -> DbResult<Option<T>> {
match self.get_raw(key)? {
Some(raw) => Ok(Some(self.codec.decode_from(raw.as_ref())?)),
None => Ok(None),
}
}
pub fn entry_len(&self, key: &K) -> Option<u32> {
let raw = self.ks.get(key.as_bytes()).ok()??;
Some(raw.as_ref().len() as u32)
}
pub fn update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
let _g = sync::lock(self.stripe_for(key));
let Some(raw) = self.get_raw(key)? else {
return Ok(None);
};
let old = self.codec.decode_from(raw.as_ref())?;
let new = f(&old);
let mut buf = Vec::new();
self.codec.encode_to(&new, &mut buf)?;
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.insert(key.as_bytes(), buf.as_slice()))?;
if H::NEEDS_WRITE {
self.hook.on_write(key, Some(&old), Some(&new));
}
Ok(Some(new))
}
pub fn fetch_update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
let _g = sync::lock(self.stripe_for(key));
let Some(raw) = self.get_raw(key)? else {
return Ok(None);
};
let old = self.codec.decode_from(raw.as_ref())?;
let new = f(&old);
let mut buf = Vec::new();
self.codec.encode_to(&new, &mut buf)?;
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.insert(key.as_bytes(), buf.as_slice()))?;
if H::NEEDS_WRITE {
self.hook.on_write(key, Some(&old), Some(&new));
}
Ok(Some(old))
}
pub fn upsert<U: crate::IntoUpsert<T>>(
&self,
key: &K,
f: impl FnOnce(Option<&T>) -> U,
) -> crate::DbResult<U::Out> {
self.atomic(key, |s| crate::upsert::upsert_via_held(s, key, f))
}
pub fn cas(&self, key: &K, expected: &T, new_value: &T) -> DbResult<()>
where
T: PartialEq,
{
let _g = sync::lock(self.stripe_for(key));
let raw = self.get_raw(key)?.ok_or(DbError::CasMismatch)?;
let cur = self.codec.decode_from(raw.as_ref())?;
if &cur != expected {
return Err(DbError::CasMismatch);
}
let mut buf = Vec::new();
self.codec.encode_to(new_value, &mut buf)?;
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.insert(key.as_bytes(), buf.as_slice()))?;
if H::NEEDS_WRITE {
self.hook.on_write(key, Some(&cur), Some(new_value));
}
Ok(())
}
pub fn compare_delete(&self, key: &K, expected: &T) -> DbResult<()>
where
T: PartialEq,
{
let _g = sync::lock(self.stripe_for(key));
let raw = self.get_raw(key)?.ok_or(DbError::CasMismatch)?;
let cur = self.codec.decode_from(raw.as_ref())?;
if &cur != expected {
return Err(DbError::CasMismatch);
}
let _mutation = self.prepare_mutation()?;
self.before_keyspace_mutation();
map_fjall(self.ks.remove(key.as_bytes()))?;
self.count_dec();
if H::NEEDS_WRITE {
self.hook.on_write(key, Some(&cur), None);
}
Ok(())
}
pub fn iter(&self) -> LsmIter<'_, K, T, C> {
LsmIter::new(self.ks.iter(), self.reversed, &self.codec)
}
pub fn range(&self, start: &K, end: &K) -> LsmIter<'_, K, T, C> {
let lo = start.as_bytes().to_vec();
let hi = end.as_bytes().to_vec();
LsmIter::new(self.ks.range(lo..hi), self.reversed, &self.codec)
}
pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> LsmIter<'_, K, T, C> {
let map = |b: Bound<&K>| match b {
Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
Bound::Unbounded => Bound::Unbounded,
};
LsmIter::new(
self.ks.range((map(start), map(end))),
self.reversed,
&self.codec,
)
}
pub(crate) fn range_bounds_dir(
&self,
start: Bound<&K>,
end: Bound<&K>,
dir: Dir,
) -> LsmIter<'_, K, T, C> {
let reversed = matches!(dir, Dir::Desc);
let map = |b: Bound<&K>| match b {
Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
Bound::Unbounded => Bound::Unbounded,
};
LsmIter::new(self.ks.range((map(start), map(end))), reversed, &self.codec)
}
pub(crate) fn keys_range_bounds_dir(
&self,
start: Bound<&K>,
end: Bound<&K>,
dir: Dir,
) -> LsmKeyIter<K> {
let reversed = matches!(dir, Dir::Desc);
let map = |b: Bound<&K>| match b {
Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
Bound::Unbounded => Bound::Unbounded,
};
LsmKeyIter::new(self.ks.range((map(start), map(end))), reversed)
}
pub(crate) fn try_keys_range_bounds_dir(
&self,
start: Bound<&K>,
end: Bound<&K>,
dir: Dir,
) -> TryLsmKeyIter<K> {
let reversed = matches!(dir, Dir::Desc);
let map = |b: Bound<&K>| match b {
Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
Bound::Unbounded => Bound::Unbounded,
};
let inner = self.ks.range((map(start), map(end)));
#[cfg(test)]
if self
.force_next_key_guard_error
.swap(false, Ordering::Relaxed)
{
return TryLsmKeyIter::new_with_forced_guard_error(inner, reversed);
}
TryLsmKeyIter::new(inner, reversed)
}
pub fn scope<const N: usize>(
&self,
key: &K,
) -> crate::scope::Scope<'_, crate::scope::TreeScan<&Self>, N> {
crate::scope::tree_scope::<_, _, N>(self, key, dir_from(self.config().reversed))
}
pub(crate) fn try_range_bounds(
&self,
start: Bound<&K>,
end: Bound<&K>,
) -> TryLsmIter<'_, K, T, C> {
self.try_range_bounds_dir(start, end, dir_from(self.reversed))
}
pub(crate) fn try_range_bounds_dir(
&self,
start: Bound<&K>,
end: Bound<&K>,
dir: Dir,
) -> TryLsmIter<'_, K, T, C> {
let map = |b: Bound<&K>| match b {
Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
Bound::Unbounded => Bound::Unbounded,
};
TryLsmIter::new(
self.ks.range((map(start), map(end))),
matches!(dir, Dir::Desc),
&self.codec,
)
}
#[cfg(test)]
pub(crate) fn force_next_key_guard_error_for_test(&self) {
self.force_next_key_guard_error
.store(true, Ordering::Relaxed);
}
pub fn prefix_iter(&self, prefix: &[u8]) -> LsmIter<'_, K, T, C> {
LsmIter::new(self.ks.prefix(prefix), self.reversed, &self.codec)
}
pub fn first(&self) -> Option<(K, T)> {
self.iter().next()
}
pub fn last(&self) -> Option<(K, T)> {
self.iter().next_back()
}
pub fn try_first(&self) -> DbResult<Option<(K, T)>> {
let g = if self.reversed {
self.ks.last_key_value()
} else {
self.ks.first_key_value()
};
match g {
Some(g) => {
let (uk, uv) = crate::lsm::map_fjall(g.into_inner())?;
Ok(Some((
K::from_bytes(uk.as_ref()),
self.codec.decode_from(uv.as_ref())?,
)))
}
None => Ok(None),
}
}
pub fn try_last(&self) -> DbResult<Option<(K, T)>> {
let g = if self.reversed {
self.ks.first_key_value()
} else {
self.ks.last_key_value()
};
match g {
Some(g) => {
let (uk, uv) = crate::lsm::map_fjall(g.into_inner())?;
Ok(Some((
K::from_bytes(uk.as_ref()),
self.codec.decode_from(uv.as_ref())?,
)))
}
None => Ok(None),
}
}
pub(crate) fn replay_init(&self) {
if H::NEEDS_INIT {
for (k, v) in self.iter() {
self.hook.on_init(&k, &v);
}
}
}
pub fn flush_buffers(&self) -> DbResult<()> {
self.flush()
}
pub fn sync_hints(&self) -> DbResult<()> {
self.periodic_flush()
}
pub fn len(&self) -> usize {
self.count.load(Ordering::Relaxed) as usize
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn recount(&self) -> DbResult<u64> {
let actual = map_fjall(self.ks.len())? as u64;
let prev = self.count.swap(actual, Ordering::Relaxed);
if prev != actual {
tracing::info!(
collection = T::NAME,
previous = prev,
recounted = actual,
"LSM live-entry counter re-derived from the keyspace"
);
}
Ok(actual)
}
pub fn flush(&self) -> DbResult<()> {
#[cfg(test)]
LSM_LAST_DURABLE.store(2, Ordering::SeqCst);
map_fjall(self.db.persist(fjall::PersistMode::SyncAll))
}
pub fn clean_shutdown(&self) -> DbResult<()> {
self.flush()
}
pub fn periodic_flush(&self) -> DbResult<()> {
#[cfg(test)]
{
let tag = match self.persist {
PersistPolicy::Buffer => 0,
PersistPolicy::SyncData => 1,
PersistPolicy::SyncAll => 2,
};
LSM_LAST_PERIODIC.store(tag, Ordering::SeqCst);
}
map_fjall(self.db.persist(self.persist.to_fjall()))
}
pub fn compact(&self) -> DbResult<usize> {
map_fjall(self.ks.major_compact())?;
Ok(0)
}
pub fn warmup(&self) -> DbResult<()> {
for _ in self.iter() {}
Ok(())
}
pub fn migrate(&self, f: impl Fn(&K, &T) -> crate::MigrateAction<T>) -> DbResult<usize> {
self.migrate_inner(f, true)
}
pub(crate) fn migrate_inner(
&self,
f: impl Fn(&K, &T) -> crate::MigrateAction<T>,
fire_init: bool,
) -> DbResult<usize> {
let init = fire_init && H::NEEDS_INIT;
let mut changed = 0usize;
for g in self.ks.iter() {
let (uk, uv) = map_fjall(g.into_inner())?;
if uk.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let key = K::from_bytes(uk.as_ref());
let val = self.codec.decode_from(uv.as_ref())?;
match f(&key, &val) {
crate::MigrateAction::Keep => {
if init {
self.hook.on_init(&key, &val);
}
}
crate::MigrateAction::Update(new) => {
{
let _g = sync::lock(self.stripe_for(&key));
self.put_locked_inner::<false>(&key, &new)?;
}
if init {
self.hook.on_init(&key, &new);
}
changed += 1;
}
crate::MigrateAction::Delete => {
let _g = sync::lock(self.stripe_for(&key));
if self.delete_locked_inner::<false>(&key)? {
changed += 1;
}
}
}
}
Ok(changed)
}
}
impl<K, T, C, H> crate::armour::Collection for LsmTree<K, T, C, H>
where
K: Key + Ord + 'static,
T: CollectionMeta<SelfId = K> + Send + Sync + 'static,
C: Codec<T> + Clone + 'static,
H: TypedWriteHook<K, T> + 'static,
{
fn name(&self) -> &str {
T::NAME
}
fn len(&self) -> usize {
LsmTree::len(self)
}
fn compact(&self) -> DbResult<usize> {
LsmTree::compact(self)
}
fn flush(&self) -> DbResult<()> {
LsmTree::flush(self)
}
fn periodic_flush(&self) -> DbResult<()> {
LsmTree::periodic_flush(self)
}
fn durable_flush(&self) -> DbResult<()> {
LsmTree::flush(self)
}
fn clean_shutdown(&self) -> DbResult<()> {
LsmTree::clean_shutdown(self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{NoHook, RapiraCodec, WatchHook};
#[cfg(feature = "armour")]
use armour_core::persist::Persist;
use armour_core::{Fuid, GetRefs, GetType};
use rapira::Rapira;
armour_core::const_hasher_or!(doc_id, "DOC", "dev-key-doc", brand = "DOC");
type DocId = Fuid<doc_id::Hasher>;
#[derive(Clone, Debug, PartialEq, GetType, GetRefs, Rapira)]
struct Doc {
name: String,
}
impl crate::CollectionMeta for Doc {
type SelfId = DocId;
const NAME: &'static str = "docs";
}
fn k(n: u64) -> DocId {
DocId::from_be_bytes(n.to_be_bytes())
}
#[cfg(feature = "armour")]
fn managed_tree() -> (
tempfile::TempDir,
Arc<crate::armour::db::ManagedLsmLifetime>,
Arc<ManagedLsmCleanMarker>,
LsmTree<DocId, Doc, RapiraCodec, NoHook>,
) {
let dir = tempfile::tempdir().unwrap();
let lifetime = Arc::new(crate::armour::db::ManagedLsmLifetime {
db_info: Persist::open(dir.path().join("db.info")).unwrap(),
_root_lease: Arc::new(crate::lock::acquire(dir.path()).unwrap()),
});
let marker = Arc::new(ManagedLsmCleanMarker::new(Arc::downgrade(&lifetime)));
let db = fjall::Database::builder(dir.path().join("lsm"))
.open()
.unwrap();
let tree = LsmTree::open_with_persist(
db,
LsmConfig::default(),
PersistPolicy::SyncAll,
RapiraCodec,
NoHook,
None,
Arc::clone(&marker),
)
.unwrap();
(dir, lifetime, marker, tree)
}
#[cfg(feature = "armour")]
fn arm_managed_clean(info: &Persist<crate::armour::DbInfo>, marker: &ManagedLsmCleanMarker) {
info.update(|db_info| db_info.clean = true).unwrap();
let mut finalization = marker.begin_finalization();
finalization.publish_clean();
}
#[cfg(feature = "armour")]
#[test]
fn every_lsm_write_api_demotes_a_managed_clean_marker() {
let (_dir, lifetime, marker, tree) = managed_tree();
let assert_dirty = || {
assert!(
!lifetime.db_info.cloned().clean,
"write bypassed clean-marker demotion"
)
};
arm_managed_clean(&lifetime.db_info, &marker);
tree.put(&k(1), &Doc { name: "put".into() }).unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
tree.insert(
&k(2),
&Doc {
name: "insert".into(),
},
)
.unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
tree.update(&k(1), |_| Doc {
name: "update".into(),
})
.unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
tree.fetch_update(&k(1), |_| Doc {
name: "fetch-update".into(),
})
.unwrap();
assert_dirty();
let expected = tree.get(&k(1)).unwrap();
arm_managed_clean(&lifetime.db_info, &marker);
tree.cas(&k(1), &expected, &Doc { name: "cas".into() })
.unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
tree.atomic(&k(3), |shard| {
shard.put(
&k(3),
&Doc {
name: "atomic".into(),
},
)
})
.unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
tree.upsert(&k(4), |_| {
crate::Upsert::Put(Doc {
name: "upsert".into(),
})
})
.unwrap();
assert_dirty();
let expected = tree.get(&k(2)).unwrap();
arm_managed_clean(&lifetime.db_info, &marker);
tree.compare_delete(&k(2), &expected).unwrap();
assert_dirty();
arm_managed_clean(&lifetime.db_info, &marker);
assert!(tree.delete(&k(3)).unwrap());
assert_dirty();
tree.put(
&k(5),
&Doc {
name: "bulk-delete".into(),
},
)
.unwrap();
arm_managed_clean(&lifetime.db_info, &marker);
assert!(tree.scope::<0>(&k(0)).delete_up_to(1).unwrap().deleted > 0);
assert_dirty();
tree.put(
&k(6),
&Doc {
name: "bulk-retain".into(),
},
)
.unwrap();
arm_managed_clean(&lifetime.db_info, &marker);
assert!(tree.scope::<0>(&k(0)).retain(|_, _| false).unwrap() > 0);
assert_dirty();
}
#[cfg(feature = "armour")]
#[test]
fn retained_managed_mutation_guard_keeps_root_lease_until_keyspace_operation_ends() {
use std::sync::mpsc;
let dir = tempfile::tempdir().unwrap();
let lifetime = Arc::new(crate::armour::db::ManagedLsmLifetime {
db_info: Persist::open(dir.path().join("db.info")).unwrap(),
_root_lease: Arc::new(crate::lock::acquire(dir.path()).unwrap()),
});
let marker = Arc::new(ManagedLsmCleanMarker::new(Arc::downgrade(&lifetime)));
let fjall = fjall::Database::builder(dir.path().join("lsm"))
.open()
.unwrap();
let tree = Arc::new(
LsmTree::open_with_persist(
fjall,
LsmConfig::default(),
PersistPolicy::SyncAll,
RapiraCodec,
NoHook,
None,
marker,
)
.unwrap(),
);
let operation = tree.hold_clean_marker_mutation_operation_for_test();
let (done_tx, done_rx) = mpsc::channel();
std::thread::scope(|scope| {
let writer = Arc::clone(&tree);
scope.spawn(move || {
done_tx
.send(writer.put(
&k(1),
&Doc {
name: "held".into(),
},
))
.unwrap();
});
operation.wait_until_blocked();
drop(lifetime);
assert!(
matches!(
crate::lock::acquire(dir.path()),
Err(DbError::Locked { .. })
),
"the admitted mutation must retain the root lease after Db lifetime release"
);
operation.release();
done_rx
.recv_timeout(std::time::Duration::from_secs(1))
.unwrap()
.unwrap();
});
crate::lock::acquire(dir.path())
.expect("the lease must release after the keyspace operation ends");
}
fn temp_tree() -> (tempfile::TempDir, LsmTree<DocId, Doc, RapiraCodec, NoHook>) {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let tree = LsmTree::open(db, LsmConfig::default(), RapiraCodec, NoHook).unwrap();
(dir, tree)
}
fn reopen_seeded(
dir: &tempfile::TempDir,
seed: Option<u64>,
) -> LsmTree<DocId, Doc, RapiraCodec, NoHook> {
let db = fjall::Database::builder(dir.path()).open().unwrap();
LsmTree::open_with_initial_count(db, LsmConfig::default(), RapiraCodec, NoHook, seed)
.unwrap()
}
#[test]
fn lsm_watch_hook_is_refused_from_direct_and_managed_open() {
let direct_dir = tempfile::tempdir().unwrap();
let direct_db = fjall::Database::builder(direct_dir.path()).open().unwrap();
let direct = LsmTree::<DocId, Doc, RapiraCodec, WatchHook<Doc, NoHook, true>>::open(
direct_db,
LsmConfig::default(),
RapiraCodec,
WatchHook::with_old(NoHook),
);
assert!(matches!(direct, Err(DbError::WatchUnsupported)));
let managed_dir = tempfile::tempdir().unwrap();
let managed_db = crate::armour::Db::open_test(managed_dir.path()).unwrap();
let managed = managed_db.open_lsm_tree::<Doc, RapiraCodec, _>(
None,
LsmConfig::default(),
WatchHook::with_old(NoHook),
&[],
);
assert!(matches!(managed, Err(DbError::WatchUnsupported)));
assert!(
!managed_db.has_lsm_database_for_test(),
"rejected watch must not open or record the managed fjall database"
);
}
#[test]
fn lsm_seeded_count_self_heals_the_two_provably_wrong_states() {
let dir = tempfile::tempdir().unwrap();
{
let t = reopen_seeded(&dir, None);
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
t.put(&k(2), &Doc { name: "b".into() }).unwrap();
t.flush().unwrap();
}
let t = reopen_seeded(&dir, Some(0));
assert_eq!(t.len(), 2, "count 0 over a non-empty keyspace must rescan");
drop(t);
let dir2 = tempfile::tempdir().unwrap();
{
let t = reopen_seeded(&dir2, None);
t.flush().unwrap();
}
let t = reopen_seeded(&dir2, Some(7));
assert_eq!(
t.len(),
0,
"non-zero count over an empty keyspace must reset"
);
}
#[test]
fn lsm_recount_rederives_a_drifted_counter() {
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
t.put(&k(2), &Doc { name: "b".into() }).unwrap();
t.count.store(99, Ordering::Relaxed);
assert_eq!(t.len(), 99);
assert_eq!(t.recount().unwrap(), 2);
assert_eq!(t.len(), 2, "recount must publish the re-derived value");
}
fn names(it: impl Iterator<Item = (DocId, Doc)>) -> Vec<String> {
it.map(|(_, d)| d.name).collect()
}
#[test]
fn lsm_iter_reversed_default_is_desc() {
let (_d, t) = temp_tree();
for i in 1..=3u64 {
t.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
assert_eq!(names(t.iter()), vec!["3", "2", "1"]);
assert_eq!(t.first().map(|(_, d)| d.name), Some("3".into()));
assert_eq!(t.last().map(|(_, d)| d.name), Some("1".into()));
}
#[test]
fn lsm_keys_use_guard_key() {
let (_d, t_rev) = temp_tree();
for i in 1..=3u64 {
t_rev
.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
let keys: Vec<_> = t_rev
.keys_range_bounds_dir(Bound::Unbounded, Bound::Unbounded, Dir::Desc)
.collect();
assert_eq!(keys, vec![k(3), k(2), k(1)]);
let keys_asc: Vec<_> = t_rev
.keys_range_bounds_dir(Bound::Unbounded, Bound::Unbounded, Dir::Asc)
.collect();
assert_eq!(keys_asc, vec![k(1), k(2), k(3)]);
let names_asc: Vec<_> = t_rev
.range_bounds_dir(Bound::Unbounded, Bound::Unbounded, Dir::Asc)
.map(|(_, d)| d.name)
.collect();
assert_eq!(names_asc, vec!["1", "2", "3"]);
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let cfg = LsmConfig {
reversed: false,
..LsmConfig::default()
};
let t_fwd: LsmTree<DocId, Doc, RapiraCodec, NoHook> =
LsmTree::open(db, cfg, RapiraCodec, NoHook).unwrap();
for i in 1..=3u64 {
t_fwd
.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
let keys_desc: Vec<_> = t_fwd
.keys_range_bounds_dir(Bound::Unbounded, Bound::Unbounded, Dir::Desc)
.collect();
assert_eq!(keys_desc, vec![k(3), k(2), k(1)]);
}
#[test]
fn lsm_iter_ascending_when_not_reversed() {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let cfg = LsmConfig {
reversed: false,
..LsmConfig::default()
};
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> =
LsmTree::open(db, cfg, RapiraCodec, NoHook).unwrap();
for i in 1..=3u64 {
t.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
assert_eq!(names(t.iter()), vec!["1", "2", "3"]);
assert_eq!(t.first().map(|(_, d)| d.name), Some("1".into()));
}
#[test]
fn lsm_reversed_tunable_across_reopen() {
let dir = tempfile::tempdir().unwrap();
{
let db = fjall::Database::builder(dir.path()).open().unwrap();
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> =
LsmTree::open(db, LsmConfig::default(), RapiraCodec, NoHook).unwrap();
for i in 1..=3u64 {
t.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
}
let db = fjall::Database::builder(dir.path()).open().unwrap();
let cfg = LsmConfig {
reversed: false,
..LsmConfig::default()
};
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> =
LsmTree::open(db, cfg, RapiraCodec, NoHook).unwrap();
assert_eq!(names(t.iter()), vec!["1", "2", "3"]); assert_eq!(t.len(), 3); }
#[test]
fn lsm_put_get_roundtrip() {
let (_d, t) = temp_tree();
let doc = Doc { name: "a".into() };
t.put(&k(1), &doc).unwrap();
assert_eq!(t.get(&k(1)), Some(doc));
assert!(t.get(&k(2)).is_none());
assert!(t.contains(&k(1)));
}
#[test]
fn lsm_insert_rejects_existing() {
let (_d, t) = temp_tree();
t.insert(&k(1), &Doc { name: "a".into() }).unwrap();
assert!(matches!(
t.insert(&k(1), &Doc { name: "b".into() }),
Err(DbError::KeyExists)
));
}
#[test]
fn lsm_delete_reports_existence_and_count() {
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
t.put(&k(2), &Doc { name: "b".into() }).unwrap();
assert_eq!(t.len(), 2);
assert!(t.delete(&k(1)).unwrap());
assert!(!t.delete(&k(1)).unwrap()); assert_eq!(t.len(), 1);
t.put(&k(2), &Doc { name: "b2".into() }).unwrap();
assert_eq!(t.len(), 1);
}
#[test]
fn lsm_update_returns_new_value_and_fires_hook() {
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
let new = t
.update(&k(1), |d| Doc {
name: format!("{}!", d.name),
})
.unwrap();
assert_eq!(new, Some(Doc { name: "a!".into() }));
assert_eq!(t.get(&k(1)), Some(Doc { name: "a!".into() }));
assert_eq!(t.update(&k(99), |d| d.clone()).unwrap(), None);
}
#[test]
fn lsm_cas_and_compare_delete() {
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
assert!(matches!(
t.cas(
&k(1),
&Doc {
name: "WRONG".into()
},
&Doc { name: "b".into() }
),
Err(DbError::CasMismatch)
));
t.cas(&k(1), &Doc { name: "a".into() }, &Doc { name: "b".into() })
.unwrap();
assert_eq!(t.get(&k(1)), Some(Doc { name: "b".into() }));
assert!(matches!(
t.compare_delete(&k(1), &Doc { name: "a".into() }),
Err(DbError::CasMismatch)
));
t.compare_delete(&k(1), &Doc { name: "b".into() }).unwrap();
assert!(!t.contains(&k(1)));
assert_eq!(t.len(), 0);
}
#[test]
fn lsm_entry_len_and_typed_reads() {
let (_d, t) = temp_tree();
assert!(matches!(t.get_or_err(&k(1)), Err(DbError::KeyNotFound)));
assert_eq!(t.try_get(&k(1)).unwrap(), None);
t.put(&k(1), &Doc { name: "abc".into() }).unwrap();
assert!(t.entry_len(&k(1)).unwrap() > 0);
assert_eq!(t.get_or_err(&k(1)).unwrap(), Doc { name: "abc".into() });
}
#[test]
fn lsm_hook_fires_with_old_and_new() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Clone)]
struct CountHook(Arc<AtomicUsize>);
impl crate::hook::TypedWriteHook<DocId, Doc> for CountHook {
fn on_write(&self, _k: &DocId, _old: Option<&Doc>, _new: Option<&Doc>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let n = Arc::new(AtomicUsize::new(0));
let t: LsmTree<DocId, Doc, RapiraCodec, CountHook> =
LsmTree::open(db, LsmConfig::default(), RapiraCodec, CountHook(n.clone())).unwrap();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
t.delete(&k(1)).unwrap();
assert_eq!(n.load(Ordering::Relaxed), 2);
}
#[test]
fn lsm_migrate_rewrites_and_deletes() {
use crate::MigrateAction;
let (_d, t) = temp_tree();
for i in 1..=4u64 {
t.put(
&k(i),
&Doc {
name: i.to_string(),
},
)
.unwrap();
}
let changed = t
.migrate(|_id, d| {
if d.name == "2" {
MigrateAction::Delete
} else {
MigrateAction::Update(Doc {
name: format!("v_{}", d.name),
})
}
})
.unwrap();
assert_eq!(changed, 4); assert_eq!(t.get(&k(1)), Some(Doc { name: "v_1".into() }));
assert!(t.get(&k(2)).is_none());
assert_eq!(t.len(), 3);
}
#[test]
fn lsm_collection_trait_object() {
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
let c: &dyn crate::armour::Collection = &t;
assert_eq!(c.name(), "docs");
assert_eq!(c.len(), 1);
c.flush().unwrap();
c.periodic_flush().unwrap();
assert_eq!(c.compact().unwrap(), 0); }
#[test]
fn lsm_atomic_runs_multiple_writes_under_one_lock() {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> = LsmTree::open(
db,
LsmConfig {
lock_stripes: 1,
..LsmConfig::default()
},
RapiraCodec,
NoHook,
)
.unwrap();
t.put(&k(1), &Doc { name: "a".into() }).unwrap();
let out = t
.atomic(&k(1), |sh| {
let cur = sh.get(&k(1)).unwrap();
sh.put(
&k(1),
&Doc {
name: format!("{}-x", cur.name),
},
)?;
sh.insert(&k(2), &Doc { name: "two".into() })?;
Ok(cur.name)
})
.unwrap();
assert_eq!(out, "a");
assert_eq!(t.get(&k(1)), Some(Doc { name: "a-x".into() }));
assert_eq!(t.get(&k(2)), Some(Doc { name: "two".into() }));
assert_eq!(t.len(), 2);
}
#[test]
fn lsm_atomic_rejects_foreign_stripe() {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> = LsmTree::open(
db,
LsmConfig {
lock_stripes: 4,
..LsmConfig::default()
},
RapiraCodec,
NoHook,
)
.unwrap();
let a = k(0);
let a_stripe = t.shard_for(&a);
let b = (1u64..=1000)
.map(k)
.find(|key| t.shard_for(key) != a_stripe)
.expect("should find a key on a different stripe within first 1000");
assert!(matches!(
t.atomic(&a, |sh| sh.put(
&b,
&Doc {
name: "foreign".into()
}
)),
Err(DbError::ShardMismatch)
));
}
#[test]
fn lsm_count_saturating_no_underflow() {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let t: LsmTree<DocId, Doc, RapiraCodec, NoHook> = LsmTree::open_with_initial_count(
db,
LsmConfig::default(),
RapiraCodec,
NoHook,
Some(0),
)
.unwrap();
t.keyspace().insert(k(1).as_bytes(), b"raw").unwrap();
t.delete(&k(1)).unwrap();
assert_eq!(t.len(), 0, "count must saturate at 0, not underflow");
assert!(t.is_empty());
}
#[test]
fn lsm_migrate_errors_on_undecodable_row() {
use crate::MigrateAction;
let (_d, t) = temp_tree();
t.put(&k(1), &Doc { name: "one".into() }).unwrap();
t.keyspace().insert(k(2).as_bytes(), b"not-a-doc").unwrap();
let err = t
.migrate(|_id, _d| MigrateAction::Update(Doc { name: "v".into() }))
.unwrap_err();
assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
}
#[test]
fn lsm_put_delete_propagate_decode_error_for_corrupt_prev() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Clone)]
struct CountHook(Arc<AtomicUsize>);
impl crate::hook::TypedWriteHook<DocId, Doc> for CountHook {
fn on_write(&self, _k: &DocId, _old: Option<&Doc>, _new: Option<&Doc>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
fn corrupt_tree() -> (
tempfile::TempDir,
LsmTree<DocId, Doc, RapiraCodec, CountHook>,
Arc<AtomicUsize>,
) {
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let hits = Arc::new(AtomicUsize::new(0));
let t: LsmTree<DocId, Doc, RapiraCodec, CountHook> = LsmTree::open(
db,
LsmConfig::default(),
RapiraCodec,
CountHook(Arc::clone(&hits)),
)
.unwrap();
t.keyspace().insert(k(1).as_bytes(), b"not-a-doc").unwrap();
(dir, t, hits)
}
let (_d, t, hits) = corrupt_tree();
let err = t.put(&k(1), &Doc { name: "x".into() }).unwrap_err();
assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
assert_eq!(
t.keyspace().get(k(1).as_bytes()).unwrap().as_deref(),
Some(&b"not-a-doc"[..]),
"failed put must not have overwritten the stored value"
);
assert_eq!(
hits.load(Ordering::Relaxed),
0,
"hook must not fire for a failed put"
);
let (_d, t, hits) = corrupt_tree();
let err = t.delete(&k(1)).unwrap_err();
assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
assert!(
t.contains(&k(1)),
"failed delete must not have removed the row"
);
assert_eq!(
hits.load(Ordering::Relaxed),
0,
"hook must not fire for a failed delete"
);
}
#[test]
fn strict_lsm_range_bounds_surfaces_decode_error_without_skipping() {
use crate::error::{DbError, DbResult};
use std::ops::Bound;
let (_d, tree) = temp_tree();
tree.put(&k(1), &Doc { name: "one".into() }).unwrap();
tree.put(
&k(3),
&Doc {
name: "three".into(),
},
)
.unwrap();
tree.keyspace()
.insert(k(2).as_bytes(), b"not-a-doc")
.unwrap();
let result = tree
.try_range_bounds(Bound::Unbounded, Bound::Unbounded)
.collect::<DbResult<Vec<_>>>();
assert!(matches!(result, Err(DbError::CorruptedEntry { .. })));
assert_eq!(
tree.iter().count(),
2,
"lenient LSM iteration remains unchanged"
);
}
#[test]
fn strict_lsm_range_rejects_short_key_before_from_bytes() {
use crate::error::DbError;
use std::ops::Bound;
let dir = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let tree: LsmTree<DocId, Doc, RapiraCodec, NoHook> = LsmTree::open(
db,
LsmConfig {
reversed: false,
..LsmConfig::default()
},
RapiraCodec,
NoHook,
)
.unwrap();
tree.put(&k(1), &Doc { name: "one".into() }).unwrap();
tree.keyspace().insert(b"short", b"not-a-doc").unwrap();
let mut iter = tree.try_range_bounds(Bound::Unbounded, Bound::Unbounded);
match iter.next() {
Some(Ok((key, d))) => {
assert_eq!(key, k(1));
assert_eq!(d, Doc { name: "one".into() });
}
other => panic!("expected Ok((k(1), doc)), got {other:?}"),
}
assert!(matches!(
iter.next(),
Some(Err(DbError::CorruptedEntry { .. }))
));
assert!(iter.next().is_none());
}
#[test]
fn lsm_try_first_last_and_aliases() {
let (_d, t) = temp_tree();
assert_eq!(t.try_first().unwrap(), None);
assert_eq!(t.try_last().unwrap(), None);
t.put(&k(1), &Doc { name: "one".into() }).unwrap();
t.put(&k(2), &Doc { name: "two".into() }).unwrap();
assert_eq!(t.try_first().unwrap(), t.first());
assert_eq!(t.try_last().unwrap(), t.last());
assert_eq!(
t.try_first().unwrap(),
Some((k(2), Doc { name: "two".into() }))
);
assert_eq!(
t.try_last().unwrap(),
Some((k(1), Doc { name: "one".into() }))
);
t.flush_buffers().unwrap();
t.sync_hints().unwrap();
}
mod durable_tick_lsm {
use super::*;
use crate::lsm::config::{LsmConfig, LsmDbConfig, PersistPolicy};
use std::time::Duration;
fn wait_periodic(expected: u8) {
let start = std::time::Instant::now();
loop {
let got = LSM_LAST_PERIODIC.load(Ordering::SeqCst);
if got == expected {
return;
}
if start.elapsed() >= Duration::from_secs(2) {
panic!(
"tick must persist with the configured policy: got {got}, expected {expected}"
);
}
std::thread::sleep(Duration::from_millis(20));
}
}
fn open_db(
policy: PersistPolicy,
) -> (
tempfile::TempDir,
crate::armour::Db,
std::sync::Arc<LsmTree<DocId, Doc, RapiraCodec, NoHook>>,
) {
let dir = tempfile::tempdir().unwrap();
let db = crate::armour::Db::builder(dir.path())
.for_tests()
.workers(crate::armour::WorkerIntervals {
compaction: None,
flush: Some(Duration::from_millis(50)),
ttl_sweep: None,
ttl_budget: 4096,
})
.open()
.unwrap();
let tree = db
.open_lsm_tree::<Doc, RapiraCodec, _>(
Some(LsmDbConfig {
persist: policy,
..LsmDbConfig::default()
}),
LsmConfig::default(),
NoHook,
&[],
)
.unwrap();
(dir, db, tree)
}
#[test]
fn durable_tick_lsm_buffer() {
LSM_LAST_PERIODIC.store(255, Ordering::SeqCst);
LSM_LAST_DURABLE.store(255, Ordering::SeqCst);
let (_dir, _db, tree) = open_db(PersistPolicy::Buffer);
wait_periodic(0);
assert_eq!(LSM_LAST_DURABLE.load(Ordering::SeqCst), 255);
tree.flush().unwrap();
assert_eq!(LSM_LAST_DURABLE.load(Ordering::SeqCst), 2);
}
#[test]
fn durable_tick_lsm_sync_data() {
LSM_LAST_PERIODIC.store(255, Ordering::SeqCst);
let (_dir, _db, _tree) = open_db(PersistPolicy::SyncData);
wait_periodic(1);
}
#[test]
fn durable_tick_lsm_sync_all() {
LSM_LAST_PERIODIC.store(255, Ordering::SeqCst);
LSM_LAST_DURABLE.store(255, Ordering::SeqCst);
let (_dir, _db, tree) = open_db(PersistPolicy::SyncAll);
wait_periodic(2);
tree.flush().unwrap();
assert_eq!(LSM_LAST_DURABLE.load(Ordering::SeqCst), 2);
}
}
}