armdb 0.7.2

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::marker::PhantomData;
use std::mem::size_of;

use crate::Key;
use crate::codec::Codec;
use crate::error::{DbError, DbResult};
use crate::lsm::map_fjall;

/// Decoding iterator over an `LsmTree`. Honors the tree's `reversed` flag by
/// pulling from the back of fjall's `DoubleEndedIterator`. Rows whose value
/// fails to decode are skipped (logged at debug), matching `VarTypedIter`.
pub struct LsmIter<'a, K, T, C> {
    inner: fjall::Iter,
    reversed: bool,
    codec: &'a C,
    _marker: PhantomData<fn() -> (K, T)>,
}

impl<'a, K, T, C> LsmIter<'a, K, T, C> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool, codec: &'a C) -> Self {
        Self {
            inner,
            reversed,
            codec,
            _marker: PhantomData,
        }
    }
}

impl<K: Key, T, C: Codec<T>> LsmIter<'_, K, T, C> {
    /// Decode a raw fjall guard into a `(K, T)` pair. Returns `None` on a guard
    /// error or value decode error (both are logged).
    fn decode_guard(&self, g: fjall::Guard) -> Option<(K, T)> {
        match g.into_inner() {
            Ok((uk, uv)) => match self.codec.decode_from(uv.as_ref()) {
                Ok(v) => Some((K::from_bytes(uk.as_ref()), v)),
                Err(_) => {
                    tracing::debug!(value_len = uv.len(), "lsm_iter: decode error, skipping");
                    None
                }
            },
            Err(e) => {
                tracing::warn!(%e, "lsm_iter: guard error, skipping");
                None
            }
        }
    }
}

impl<K: Key, T, C: Codec<T>> Iterator for LsmIter<'_, K, T, C> {
    type Item = (K, T);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let g = if self.reversed {
                self.inner.next_back()
            } else {
                self.inner.next()
            }?;
            if let Some(kv) = self.decode_guard(g) {
                return Some(kv);
            }
        }
    }
}

impl<K: Key, T, C: Codec<T>> DoubleEndedIterator for LsmIter<'_, K, T, C> {
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            let g = if self.reversed {
                self.inner.next()
            } else {
                self.inner.next_back()
            }?;
            if let Some(kv) = self.decode_guard(g) {
                return Some(kv);
            }
        }
    }
}

/// Strict decoding iterator over an `LsmTree`. Honors the tree's `reversed` flag
/// and surfaces storage/decode errors instead of skipping corrupt rows.
pub(crate) struct TryLsmIter<'a, K, T, C> {
    inner: fjall::Iter,
    codec: &'a C,
    reversed: bool,
    failed: bool,
    _marker: PhantomData<fn() -> (K, T)>,
}

impl<'a, K, T, C> TryLsmIter<'a, K, T, C> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool, codec: &'a C) -> Self {
        Self {
            inner,
            codec,
            reversed,
            failed: false,
            _marker: PhantomData,
        }
    }
}

impl<K: Key, T, C: Codec<T>> TryLsmIter<'_, K, T, C> {
    fn decode_guard(&self, g: fjall::Guard) -> DbResult<(K, T)> {
        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 value = self.codec.decode_from(uv.as_ref())?;
        Ok((key, value))
    }
}

impl<K: Key, T, C: Codec<T>> Iterator for TryLsmIter<'_, K, T, C> {
    type Item = DbResult<(K, T)>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        let raw = if self.reversed {
            self.inner.next_back()?
        } else {
            self.inner.next()?
        };
        let result = self.decode_guard(raw);
        if result.is_err() {
            self.failed = true;
        }
        Some(result)
    }
}