coordinode-lsm-tree 7.4.2

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

//! Size-adaptive block-index writer (streaming spill).
//!
//! Picks the index layout by actual index size instead of an
//! unconditional policy:
//!
//! - While the accumulated index stays at or below `spill_threshold`
//!   bytes, entries are buffered and the index is written as a
//!   **single-level** ([`FullIndexWriter`]) index at `finish`. The
//!   whole index is one block; the reader pins it in `FullBlockIndex`
//!   and a point read costs one index-block lookup.
//! - The moment the buffer would exceed `spill_threshold`, the writer
//!   **spills**: it replays the buffered handles into a streaming
//!   [`PartitionedIndexWriter`] and forwards every subsequent handle to
//!   it, producing a **two-level** (partitioned) index. The reader then
//!   keeps only the top-level index resident and pages bottom partitions
//!   through the block cache.
//!
//! Rationale: a two-level index bounds resident index RAM per open SST
//! (only the TLI stays pinned), which matters for large datasets — but
//! it costs an extra index level on every point read. For small/medium
//! SSTs, pinning the whole index is cheap and the extra level is pure
//! overhead, so single-level is both faster and memory-fine. This writer
//! gets the best of both: single-level until the index is large enough
//! that pinning it whole stops being cheap, then partitioned.
//!
//! Either way the index stays in memory until `finish` writes it: the
//! partitioned writer buffers its encoded partitions too. What bounds it is
//! the table: the writer reports the bytes it holds, and a run of tables
//! rotates to the next one before they reach its target size.

use crate::{
    CompressionType,
    checksum::ChecksummedWriter,
    encryption::EncryptionProvider,
    table::{
        index_block::KeyedBlockHandle,
        writer::index::{BlockIndexWriter, FullIndexWriter, PartitionedIndexWriter},
    },
};
use alloc::sync::Arc;
#[cfg(not(feature = "std"))]
use alloc::{boxed::Box, vec::Vec};

use crate::io::{Seek, Write};

/// Index-size threshold (bytes) at or below which the index is written
/// single-level. The estimate mirrors [`PartitionedIndexWriter`]'s own
/// per-entry accounting (`end_key.len() + size_of::<KeyedBlockHandle>()`).
///
/// A single-level index is reached by one lookup (versus two-level's two).
/// On hot levels the reader pins the index block; on cold levels it pages
/// the block through the shared block cache (evictable), so a multi-MiB
/// threshold does not pin unbounded RAM. 4 MiB keeps SSTs up to a
/// few-hundred-MB single-level (where single-level wins point reads) while
/// genuinely huge indexes still partition. Tunable per tree (see the
/// index-partition policy / runtime config).
pub const DEFAULT_SPILL_THRESHOLD: u64 = 4 * 1024 * 1024;

pub struct AdaptiveIndexWriter<W: Write + Seek + 'static> {
    // Forwarded config (applied to whichever inner writer is built).
    compression: CompressionType,
    restart_interval: u8,

    /// Whether zstd levels 19-22 run the `btultra2` two-pass seed, forwarded
    /// like the rest of this config. Defaults to `true`.
    #[cfg(zstd_any)]
    zstd_two_pass_seed: bool,
    /// Per-bottom-partition size, forwarded to the partitioned writer
    /// once spilled. Distinct from `spill_threshold` (which decides
    /// *whether* to partition at all).
    partition_size: u32,
    encryption: Option<Arc<dyn EncryptionProvider>>,
    table_id: crate::TableId,
    ecc: Option<crate::table::block::EccParams>,

    /// Total index size, in bytes, above which we spill to two-level.
    spill_threshold: u64,

    // Pre-spill state: buffered handles + running size estimate.
    buffer: Vec<KeyedBlockHandle>,
    buffered_bytes: u64,
    /// Bytes the buffered handles take encoded, bounded from above.
    buffered_encoded: u64,

    /// `Some` once spilled — every subsequent handle is forwarded here
    /// and `finish` delegates to it.
    spilled: Option<Box<dyn BlockIndexWriter<W>>>,
}

impl<W: Write + Seek + 'static> AdaptiveIndexWriter<W> {
    #[must_use]
    pub fn new(spill_threshold: u64) -> Self {
        Self {
            compression: CompressionType::None,
            restart_interval: 1,
            #[cfg(zstd_any)]
            zstd_two_pass_seed: true,
            partition_size: 4_096,
            encryption: None,
            table_id: 0,
            ecc: None,
            spill_threshold,
            buffer: Vec::new(),
            buffered_bytes: 0,
            buffered_encoded: 0,
            spilled: None,
        }
    }

    /// Applies the buffered config to a freshly-constructed inner writer
    /// (`Full` or `Partitioned`), so both layouts inherit identical
    /// compression / encryption / restart / ecc / table-id settings.
    fn configure(&self, inner: Box<dyn BlockIndexWriter<W>>) -> Box<dyn BlockIndexWriter<W>> {
        #[cfg(zstd_any)]
        let inner = inner.use_zstd_two_pass_seed(self.zstd_two_pass_seed);
        inner
            .use_compression(self.compression)
            .use_restart_interval(self.restart_interval)
            .use_partition_size(self.partition_size)
            .use_encryption(self.encryption.clone())
            .use_table_id(self.table_id)
            .use_ecc(self.ecc)
    }

    /// Bytes the single block the buffered handles encode to takes.
    fn buffered_block(&self) -> u64 {
        self.buffered_encoded + crate::table::block::TRAILER_LEN as u64
    }

    /// Transition to two-level: build a partitioned writer with the
    /// current config and replay the buffered handles into it.
    fn spill(&mut self) -> crate::Result<()> {
        let mut partitioned = self.configure(Box::new(PartitionedIndexWriter::new()));
        // Taking the buffer frees its allocation once replayed; draining it
        // would keep the capacity for the life of the table.
        for handle in core::mem::take(&mut self.buffer) {
            partitioned.register_data_block(handle)?;
        }
        self.buffered_bytes = 0;
        self.buffered_encoded = 0;
        self.spilled = Some(partitioned);
        Ok(())
    }
}

impl<W: Write + Seek + 'static> BlockIndexWriter<W> for AdaptiveIndexWriter<W> {
    fn register_data_block(&mut self, block_handle: KeyedBlockHandle) -> crate::Result<()> {
        if let Some(partitioned) = &mut self.spilled {
            return partitioned.register_data_block(block_handle);
        }

        // Mirror PartitionedIndexWriter's per-entry size estimate so the
        // threshold compares like-for-like against the partition logic.
        let entry_size =
            (block_handle.end_key().len() + core::mem::size_of::<KeyedBlockHandle>()) as u64;
        self.buffered_bytes += entry_size;
        self.buffered_encoded += block_handle.encoded_len_bound() as u64;
        self.buffer.push(block_handle);

        if self.buffered_bytes > self.spill_threshold {
            self.spill()?;
        }
        Ok(())
    }

    fn held_bytes(&self) -> u64 {
        match &self.spilled {
            Some(partitioned) => partitioned.held_bytes(),
            // The buffer's size counts bytes held in memory, so it fits in
            // `usize`.
            #[expect(clippy::cast_possible_truncation, reason = "bytes held in memory")]
            None => super::handles_held(
                self.buffered_bytes as usize,
                &self.buffer,
                self.buffer.capacity(),
            ) as u64,
        }
    }

    fn finish_scratch_bytes(&self) -> u64 {
        if let Some(partitioned) = &self.spilled {
            return partitioned.finish_scratch_bytes();
        }
        // The full writer takes the buffer over and encodes it into one
        // block, which the table keeps until it writes the tail mirror, and
        // frames it into a second when the block is transformed.
        let block = self.buffered_block();
        block
            + crate::table::block::transform_scratch_bound(
                block,
                crate::table::block::BlockType::Index,
                self.compression,
                self.encryption.as_deref(),
                self.ecc,
            )
    }

    fn finish_output_bytes(&self) -> u64 {
        match &self.spilled {
            Some(partitioned) => partitioned.finish_output_bytes(),
            None if self.buffer.is_empty() => 0,
            // One block, written twice: at the head and as the tail mirror.
            None => {
                2 * crate::table::block::framed_len_bound(
                    self.buffered_block(),
                    crate::table::block::BlockType::Index,
                    self.compression,
                    self.encryption.as_deref(),
                    self.ecc,
                )
            }
        }
    }

    fn finish(
        self: Box<Self>,
        file_writer: &mut crate::sfa::Writer<ChecksummedWriter<W>>,
    ) -> crate::Result<(usize, Vec<u8>)> {
        let mut this = *self;

        // Spilled → two-level index is already streaming; just finish it.
        if let Some(partitioned) = this.spilled.take() {
            return partitioned.finish(file_writer);
        }

        // Stayed small → single-level (Full) index.
        // The full writer takes the buffer over rather than a copy of it.
        // Both sizes count bytes held in memory, so they fit in `usize`.
        #[expect(clippy::cast_possible_truncation, reason = "bytes held in memory")]
        let full = FullIndexWriter::with_handles(
            core::mem::take(&mut this.buffer),
            this.buffered_bytes as usize,
            this.buffered_encoded as usize,
        );
        this.configure(Box::new(full)).finish(file_writer)
    }

    fn use_compression(
        mut self: Box<Self>,
        compression: CompressionType,
    ) -> Box<dyn BlockIndexWriter<W>> {
        self.compression = compression;
        self
    }

    #[cfg(zstd_any)]
    fn use_zstd_two_pass_seed(mut self: Box<Self>, enabled: bool) -> Box<dyn BlockIndexWriter<W>> {
        self.zstd_two_pass_seed = enabled;
        // A writer this one has already spilled into is past `configure`, so it
        // has to learn the setting directly.
        if let Some(spilled) = self.spilled.take() {
            self.spilled = Some(spilled.use_zstd_two_pass_seed(enabled));
        }
        self
    }

    fn use_restart_interval(mut self: Box<Self>, interval: u8) -> Box<dyn BlockIndexWriter<W>> {
        self.restart_interval = interval;
        self
    }

    fn use_partition_size(mut self: Box<Self>, size: u32) -> Box<dyn BlockIndexWriter<W>> {
        self.partition_size = size;
        self
    }

    fn use_encryption(
        mut self: Box<Self>,
        encryption: Option<Arc<dyn EncryptionProvider>>,
    ) -> Box<dyn BlockIndexWriter<W>> {
        self.encryption = encryption;
        self
    }

    fn use_table_id(mut self: Box<Self>, table_id: crate::TableId) -> Box<dyn BlockIndexWriter<W>> {
        self.table_id = table_id;
        self
    }

    fn use_ecc(
        mut self: Box<Self>,
        ecc: Option<crate::table::block::EccParams>,
    ) -> Box<dyn BlockIndexWriter<W>> {
        self.ecc = ecc;
        self
    }
}

#[cfg(test)]
mod tests;