lsm_tree/memtable/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2024-present, fjall-rs
// Copyright (c) 2026-present, Dmitry Prudnikov

pub(crate) mod arena;
pub(crate) mod interval_tree;
pub(crate) mod own_line;
use own_line::OwnLine;
pub(crate) mod skiplist;
pub(crate) mod value_store;

use crate::comparator::SharedComparator;
use crate::key::InternalKey;
use crate::range_tombstone::RangeTombstone;
use crate::{
    UserKey,
    value::{InternalValue, SeqNo},
};
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
use core::ops::RangeBounds;
use core::sync::atomic::AtomicBool;
use portable_atomic::AtomicU64;
// `parking_lot::RwLock` (std: small, userspace fast-path, no poisoning) /
// `spin::RwLock` (no_std). Neither poisons on a panicked holder, so the read/
// write guards are taken without a `LockResult` unwrap.
#[cfg(feature = "std")]
use parking_lot::RwLock;
#[cfg(not(feature = "std"))]
use spin::RwLock;

pub use crate::tree::inner::MemtableId;

/// The distinct keys among `keys`, in the table's key order, and the distinct
/// hashes a table filter built over them holds: one per key and, with
/// `prefixes`, one per prefix, as a full filter deduplicates them before it
/// builds. Versions of a key add none.
///
/// Counted over all of `keys` at once: a flush rotating into several tables
/// hashes a prefix spanning two of them in each, so the count runs short by
/// at most one prefix per rotation.
pub(crate) fn filter_count(
    keys: impl Iterator<Item = crate::UserKey>,
    prefixes: Option<&dyn crate::PrefixExtractor>,
) -> crate::filter_budget::FilterCount {
    let mut distinct = 0u64;
    let mut previous: Option<crate::UserKey> = None;
    // Without an extractor a filter holds one hash a key, and needs no list.
    let mut hashes: Vec<u64> = Vec::new();
    let mut previous_prefixes: Vec<u64> = Vec::new();
    for key in keys {
        if previous
            .as_ref()
            .is_some_and(|previous| crate::comparator::same_user_key(previous, &key))
        {
            continue;
        }
        distinct += 1;
        if let Some(extractor) = prefixes {
            hashes.push(crate::hash::hash64(&key));
            // A prefix the key before had at the same position is a repeat,
            // dropped here to keep the list short, as the writer does.
            for (position, prefix) in extractor.prefixes(&key).enumerate() {
                let hash = crate::hash::hash64(prefix);
                match previous_prefixes.get_mut(position) {
                    Some(previous) if *previous == hash => {}
                    Some(previous) => {
                        *previous = hash;
                        hashes.push(hash);
                    }
                    None => {
                        previous_prefixes.push(hash);
                        hashes.push(hash);
                    }
                }
            }
        }
        previous = Some(key);
    }
    let hashes = if prefixes.is_some() {
        // Tokens equal to a key, or repeated out of position, the writer
        // drops when it sorts its hashes.
        hashes.sort_unstable();
        hashes.dedup();
        u64::try_from(hashes.len()).unwrap_or(u64::MAX)
    } else {
        distinct
    };
    crate::filter_budget::FilterCount {
        keys: distinct,
        hashes,
    }
}

/// The memtable serves as an intermediary, ephemeral, sorted storage for new items
///
/// When the Memtable exceeds some size, it should be flushed to a table.
pub struct Memtable {
    #[doc(hidden)]
    pub id: MemtableId,

    /// The user key comparator used for ordering entries.
    pub(crate) comparator: SharedComparator,

    /// The actual content, stored in an arena-based skiplist with lock-free traversal.
    ///
    /// Nodes are allocated from a contiguous byte arena for cache locality
    /// and O(1) bulk deallocation when the memtable is dropped.  Traversal of
    /// the skiplist index uses atomic loads and CAS for inserts.
    pub(crate) items: skiplist::SkipMap,

    /// Range tombstones stored in an interval tree.
    ///
    /// Protected by `RwLock` — read-heavy suppression queries (`query_suppression`,
    /// `range_tombstones_sorted`) take a shared read lock, while `insert_range_tombstone`
    /// takes an exclusive write lock. After a rotation has been requested via
    /// `requested_rotation`, the interval tree is treated as read-only by convention,
    /// and only readers are expected to access this field (the `RwLock` is still used
    /// for synchronization, but there should be no further writes).
    ///
    /// `std::sync::RwLock` may be reader-biased on some platforms, but writer
    /// starvation is not a concern here: range deletes are rare, the write-side
    /// critical section is O(log n) with n typically small, and the memtable
    /// rotates (becoming read-only) well before contention could accumulate.
    pub(crate) range_tombstones: RwLock<interval_tree::IntervalTree>,

    /// Approximate active memtable size.
    ///
    /// If this grows too large, a flush is triggered. Updated with `Relaxed`
    /// RMWs: it is a rotation heuristic, and nothing derives the visibility
    /// of other memory from it (entry visibility comes from the skiplist's
    /// release/acquire tower stores; `highest_seqno` keeps `AcqRel` so a
    /// reader that observes a seqno also observes its entry).
    ///
    /// Each counter every insert writes sits on a line of its own, apart from
    /// the fields reads and inserts only load.
    pub(crate) approximate_size: OwnLine<AtomicU64>,

    /// Highest encountered sequence number.
    ///
    /// This is used so that `get_highest_seqno` has O(1) complexity.
    pub(crate) highest_seqno: OwnLine<AtomicU64>,

    pub(crate) requested_rotation: AtomicBool,

    /// Whether any insert-time per-KV digest (`KvChecksumComputePoint::AtInsert`)
    /// has been stored in this memtable. Set on the first digest-bearing insert
    /// and read once at flush by [`Self::verify_kv_residence`] to skip walking
    /// the nodes entirely when there is nothing to verify (the default `Off` /
    /// `AtBlockCompile` path). The per-node digest carries its own algorithm, so
    /// no memtable-wide algorithm is tracked here.
    has_at_insert_digests: AtomicBool,

    /// Blob files the cell rows written here reference. Until this memtable
    /// is flushed, no table links them for these rows, so a blob file listed
    /// here is neither removed nor punched.
    blob_refs: RwLock<crate::HashSet<crate::vlog::BlobFileId>>,
}

impl Memtable {
    /// Returns the memtable ID.
    pub fn id(&self) -> MemtableId {
        self.id
    }

    /// Returns `true` if the memtable was already flagged for rotation.
    pub fn is_flagged_for_rotation(&self) -> bool {
        self.requested_rotation
            .load(core::sync::atomic::Ordering::Relaxed)
    }

    /// Flags the memtable as requested for rotation.
    pub fn flag_rotated(&self) {
        self.requested_rotation
            .store(true, core::sync::atomic::Ordering::Relaxed);
    }

    // `pub` + `#[doc(hidden)]`: used by the host crate (fjall) to construct
    // ephemeral memtables. Not part of the semver-stable API.
    // Keep the comparator by-value for hidden-public API compatibility while
    // still requiring callers to pass the tree comparator explicitly.
    #[doc(hidden)]
    #[expect(
        clippy::needless_pass_by_value,
        reason = "hidden-public constructor keeps the preexisting by-value signature for compatibility"
    )]
    #[must_use]
    pub fn new(id: MemtableId, comparator: SharedComparator) -> Self {
        Self {
            id,
            items: skiplist::SkipMap::new(comparator.clone()),
            comparator: comparator.clone(),
            range_tombstones: RwLock::new(interval_tree::IntervalTree::new_with_comparator(
                comparator.clone(),
            )),
            approximate_size: OwnLine(AtomicU64::default()),
            highest_seqno: OwnLine(AtomicU64::default()),
            requested_rotation: AtomicBool::default(),
            has_at_insert_digests: AtomicBool::default(),
            blob_refs: RwLock::new(crate::HashSet::default()),
        }
    }

    /// Records the blob files `item` references when it is a cell row.
    ///
    /// Every insert path calls this, so a row replayed from an external log
    /// registers its files like a row written through
    /// [`BlobTree::insert_cells`](crate::BlobTree::insert_cells). A row that
    /// does not decode registers nothing: it fails at flush or read, which
    /// decode it again.
    fn note_blob_refs(&self, item: &InternalValue) {
        if !item.key.value_type.is_cell_row() {
            return;
        }
        let Ok(refs) = crate::blob_tree::field_row::row_refs(&item.value) else {
            return;
        };
        if refs.is_empty() {
            return;
        }
        let mut set = self.blob_refs.write();
        for (indirection, _) in refs {
            set.insert(indirection.vhandle.blob_file_id);
        }
    }

    /// Whether a cell row in this memtable references `blob_file_id`.
    pub(crate) fn references_blob_file(&self, blob_file_id: crate::vlog::BlobFileId) -> bool {
        self.blob_refs.read().contains(&blob_file_id)
    }

    /// Adds the blob files cell rows in this memtable reference to `files`.
    pub(crate) fn collect_blob_files(&self, files: &mut crate::HashSet<crate::vlog::BlobFileId>) {
        files.extend(self.blob_refs.read().iter().copied());
    }

    /// Creates an iterator over all items.
    pub fn iter(&self) -> impl DoubleEndedIterator<Item = InternalValue> + '_ {
        self.items.iter().map(|entry| InternalValue {
            key: entry.key(),
            value: entry.value(),
        })
    }

    /// Creates an iterator over a range of items.
    ///
    /// Accepts `InternalKey`-based bounds.
    pub(crate) fn range_internal<'a, R: RangeBounds<InternalKey> + 'a>(
        &'a self,
        range: R,
    ) -> impl DoubleEndedIterator<Item = InternalValue> + 'a {
        self.items.range(range).map(|entry| InternalValue {
            key: entry.key(),
            value: entry.value(),
        })
    }

    /// Returns the item by key if it exists.
    ///
    /// Returns the version with the highest seqno that is strictly less than
    /// the given `seqno`.  Pass [`MAX_SEQNO`](crate::MAX_SEQNO) to retrieve the latest version.
    #[doc(hidden)]
    pub fn get(&self, key: &[u8], seqno: SeqNo) -> Option<InternalValue> {
        if seqno == 0 {
            return None;
        }

        // NOTE: This range start deserves some explanation...
        // InternalKeys are multi-sorted by 2 categories: user_key and Reverse(seqno). (tombstone doesn't really matter)
        // We search for the lowest entry that is greater or equal the user's prefix key
        // and has the seqno (or lower) we want (because the seqno is stored in reverse order)
        //
        // Example: We search for "abc"
        //
        // key -> seqno
        //
        // a   -> 7
        // abc -> 5 <<< This is the lowest key (highest seqno) that matches the key with seqno=MAX
        // abc -> 4
        // abc -> 3 <<< If searching for abc and seqno=4, we would get this
        // abcdef -> 6
        // abcdef -> 5
        //
        // Borrowed-parts point lookup: no `InternalKey` bound is built, so the
        // probe key is never copied (the old `range(lower_bound..)` path
        // heap-copied `key` into a fresh bound per probed memtable).
        self.items.point_get(key, seqno - 1)
    }

    /// Gets approximate size of memtable in bytes.
    pub fn size(&self) -> u64 {
        self.approximate_size
            .load(core::sync::atomic::Ordering::Acquire)
    }

    /// Counts the number of items in the memtable.
    pub fn len(&self) -> usize {
        self.items.len()
    }

    /// How often an insert had to retry linking an item because a concurrent
    /// insert changed its neighbour first: the contention writers meet in
    /// this memtable, counted once per retried level of the skiplist.
    #[cfg(feature = "metrics")]
    #[must_use]
    pub fn insert_cas_retries(&self) -> usize {
        self.items.cas_retries()
    }

    /// Returns `true` if the memtable has no KV items and no range tombstones.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.items.is_empty() && self.range_tombstone_count() == 0
    }

    /// Inserts multiple items into the memtable in bulk.
    ///
    /// More efficient than calling [`Memtable::insert`] in a loop because it
    /// performs a single `fetch_add` for the total size and a single
    /// `fetch_max` for the highest seqno.
    ///
    /// Returns `(total_bytes_added, new_memtable_size)`.
    #[doc(hidden)]
    pub fn insert_batch(&self, items: Vec<InternalValue>) -> (u64, u64) {
        self.insert_batch_with_kv_algo(items, None)
    }

    /// Bulk insert, optionally computing an insert-time per-KV digest per item
    /// under `kv_algo` (`KvChecksumComputePoint::AtInsert`).
    ///
    /// `kv_algo` is `Some(algo)` (a 4-byte algorithm) to fix each entry's
    /// digest at insert for the flush-time residence check, or `None` for the
    /// plain bulk path. Same single-`fetch_add` / single-`fetch_max` accounting
    /// as [`Self::insert_batch`].
    #[doc(hidden)]
    pub fn insert_batch_with_kv_algo(
        &self,
        items: Vec<InternalValue>,
        kv_algo: Option<crate::runtime_config::ChecksumAlgorithm>,
    ) -> (u64, u64) {
        if items.is_empty() {
            let size = self
                .approximate_size
                .load(core::sync::atomic::Ordering::Acquire);
            return (0, size);
        }

        let mut total_size: u64 = 0;
        let mut max_seqno: u64 = 0;

        let overhead =
            core::mem::size_of::<InternalValue>() + core::mem::size_of::<SharedComparator>();

        for item in &items {
            #[expect(
                clippy::expect_used,
                reason = "keys are limited to 16-bit length + values are limited to 32-bit length"
            )]
            let item_size: u64 = (item.key.user_key.len() + item.value.len() + overhead)
                .try_into()
                .expect("should fit into u64");

            // Running memtable byte total, bounded by the in-memory data size;
            // a plain add cannot overflow u64.
            total_size += item_size;

            if item.key.seqno > max_seqno {
                max_seqno = item.key.seqno;
            }
        }

        let size_before = self
            .approximate_size
            .fetch_add(total_size, core::sync::atomic::Ordering::Relaxed);

        if kv_algo.is_some() {
            // Flag that this memtable carries residence digests for the
            // flush-time verify. The algorithm lives per node.
            self.has_at_insert_digests
                .store(true, core::sync::atomic::Ordering::Relaxed);
        }

        for item in items {
            self.note_blob_refs(&item);
            let digest = kv_algo.and_then(|algo| {
                crate::table::block::kv_checksum::kv_digest(&item, algo).map(|d| {
                    #[expect(
                        clippy::cast_possible_truncation,
                        reason = "AtInsert is config-validated to a 4-byte algorithm; the digest fits u32"
                    )]
                    let lo = d as u32;
                    (lo, algo)
                })
            });
            self.items
                .insert_with_kv_digest(&item.key, item.value, digest);
        }

        self.highest_seqno
            .fetch_max(max_seqno, core::sync::atomic::Ordering::AcqRel);

        // fetch_add returns value BEFORE the add, so size_before + total_size
        // = value AFTER add = new memtable size. Same pattern as Memtable::insert().
        (total_size, size_before + total_size)
    }

    /// Inserts an item into the memtable
    #[doc(hidden)]
    pub fn insert(&self, item: InternalValue) -> (u64, u64) {
        #[expect(
            clippy::expect_used,
            reason = "keys are limited to 16-bit length + values are limited to 32-bit length"
        )]
        // Account for MemtableKey overhead (InternalKey + Arc<dyn UserComparator>)
        let item_size = (item.key.user_key.len()
            + item.value.len()
            + core::mem::size_of::<InternalValue>()
            + core::mem::size_of::<SharedComparator>())
        .try_into()
        .expect("should fit into u64");

        let size_before = self
            .approximate_size
            .fetch_add(item_size, core::sync::atomic::Ordering::Relaxed);

        self.note_blob_refs(&item);

        // `item.key` is borrowed and `item.value` moved: the key bytes are
        // copied into the arena, the value handle is stored as-is.
        self.items.insert(&item.key, item.value);

        self.highest_seqno
            .fetch_max(item.key.seqno, core::sync::atomic::Ordering::AcqRel);

        (item_size, size_before + item_size)
    }

    /// Inserts an item, optionally carrying a precomputed insert-time per-KV
    /// digest (`KvChecksumComputePoint::AtInsert`).
    ///
    /// `kv_digest` is `Some((digest, algo))` when the caller computed the
    /// entry's 4-byte logical-content digest at insert (under `AtInsert` with a
    /// 4-byte algorithm), or `None` for the plain path. When present, the digest
    /// and its algorithm are stored in the skiplist node (per node, so a later
    /// config change cannot misverify it) and the memtable flags that it carries
    /// at least one digest so [`Self::verify_kv_residence`] knows to walk at
    /// flush. Mixed inserts (some with, some without a digest) are supported for
    /// the `Off` -> `AtInsert` live toggle.
    #[doc(hidden)]
    pub fn insert_with_kv_digest(
        &self,
        item: InternalValue,
        kv_digest: Option<(u32, crate::runtime_config::ChecksumAlgorithm)>,
    ) -> (u64, u64) {
        #[expect(
            clippy::expect_used,
            reason = "keys are limited to 16-bit length + values are limited to 32-bit length"
        )]
        let item_size = (item.key.user_key.len()
            + item.value.len()
            + core::mem::size_of::<InternalValue>()
            + core::mem::size_of::<SharedComparator>())
        .try_into()
        .expect("should fit into u64");

        let size_before = self
            .approximate_size
            .fetch_add(item_size, core::sync::atomic::Ordering::Relaxed);

        if kv_digest.is_some() {
            // Flag that this memtable carries at least one residence digest so
            // the flush-time verify walks the nodes. The algorithm lives per
            // node, not here.
            self.has_at_insert_digests
                .store(true, core::sync::atomic::Ordering::Relaxed);
        }

        self.note_blob_refs(&item);

        self.items
            .insert_with_kv_digest(&item.key, item.value, kv_digest);

        self.highest_seqno
            .fetch_max(item.key.seqno, core::sync::atomic::Ordering::AcqRel);

        (item_size, size_before + item_size)
    }

    /// Verifies every insert-time per-KV digest in this memtable against a
    /// recompute over the entry's current bytes (the
    /// [`KvChecksumComputePoint::AtInsert`](crate::runtime_config::KvChecksumComputePoint::AtInsert)
    /// residence check), called once at flush.
    ///
    /// Returns `Ok` immediately when no `AtInsert` digest was ever inserted, so
    /// the default path pays nothing.
    ///
    /// # Errors
    ///
    /// - [`crate::Error::MemtableKvChecksumMismatch`] when an entry's stored
    ///   digest diverges from the recompute (a RAM bit-flip during residence).
    /// - [`crate::Error::FeatureUnsupported`] when a node's algorithm is not
    ///   compiled into this build.
    pub fn verify_kv_residence(&self) -> crate::Result<()> {
        if !self
            .has_at_insert_digests
            .load(core::sync::atomic::Ordering::Relaxed)
        {
            return Ok(());
        }
        self.items.verify_kv_digests()
    }

    /// Inserts a range tombstone covering `[start, end)` at the given seqno.
    ///
    /// Returns the approximate size added to the memtable.
    ///
    /// Returns 0 if `start >= end` or if either bound exceeds `u16::MAX` bytes.
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned.
    #[must_use]
    pub fn insert_range_tombstone(&self, start: UserKey, end: UserKey, seqno: SeqNo) -> u64 {
        // flag_rotated() (which sets requested_rotation) is called by the host
        // crate (fjall) before rotation; this crate never sets it directly.
        // The assert catches misuse by callers
        // in debug builds — intentionally debug-only because post-rotation writes
        // are structurally prevented by the host (sealed memtables are behind Arc
        // with no write path exposed), and an atomic load here would add overhead
        // on the hot insert path in release builds for no practical benefit.
        debug_assert!(
            !self.is_flagged_for_rotation(),
            "insert_range_tombstone called after memtable was flagged for rotation"
        );

        // Reject invalid intervals in release builds (debug_assert is not enough)
        if self.comparator.compare(&start, &end) != core::cmp::Ordering::Less {
            return 0;
        }

        // On-disk RT format writes key lengths as u16, enforce at insertion time.
        // Emit a warning when rejecting an oversized bound so this failure is diagnosable.
        if u16::try_from(start.len()).is_err() || u16::try_from(end.len()).is_err() {
            log::warn!(
                "insert_range_tombstone: rejecting oversized range tombstone \
                 bounds (start_len = {}, end_len = {}, max = {})",
                start.len(),
                end.len(),
                u16::MAX,
            );
            return 0;
        }

        let size = (start.len() + end.len() + core::mem::size_of::<RangeTombstone>()) as u64;

        self.range_tombstones
            .write()
            .insert(RangeTombstone::new(start, end, seqno));

        self.approximate_size
            .fetch_add(size, core::sync::atomic::Ordering::AcqRel);

        self.highest_seqno
            .fetch_max(seqno, core::sync::atomic::Ordering::AcqRel);

        size
    }

    /// Returns `true` if the key at `key_seqno` is suppressed by a range tombstone
    /// visible at `read_seqno`.
    pub(crate) fn is_key_suppressed_by_range_tombstone(
        &self,
        key: &[u8],
        key_seqno: SeqNo,
        read_seqno: SeqNo,
    ) -> bool {
        self.range_tombstones
            .read()
            .query_suppression(key, key_seqno, read_seqno)
    }

    /// Hands `visit` the range tombstones holding `key`.
    pub(crate) fn for_each_range_tombstone_containing(
        &self,
        key: &[u8],
        visit: impl FnMut(&RangeTombstone),
    ) {
        self.range_tombstones.read().for_each_containing(key, visit);
    }

    /// Returns all range tombstones in sorted order (for flush).
    pub(crate) fn range_tombstones_sorted(&self) -> Vec<RangeTombstone> {
        self.range_tombstones.read().iter_sorted()
    }

    /// Returns the number of range tombstones.
    #[must_use]
    pub fn range_tombstone_count(&self) -> usize {
        self.range_tombstones.read().len()
    }

    /// Returns the highest sequence number in the memtable.
    pub fn get_highest_seqno(&self) -> Option<SeqNo> {
        if self.is_empty() {
            None
        } else {
            Some(
                self.highest_seqno
                    .load(core::sync::atomic::Ordering::Acquire),
            )
        }
    }
}

#[cfg(test)]
mod tests;