coordinode-lsm-tree 5.12.3

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

use alloc::boxed::Box;
use alloc::sync::Arc;

use super::{Block, DataBlock};
use crate::io::BufReader;
use crate::path::Path;
use crate::{
    CompressionType, InternalValue, SeqNo,
    comparator::SharedComparator,
    encryption::EncryptionProvider,
    fs::{FileHint, Fs, FsFile, FsOpenOptions},
    table::{block::BlockType, iter::OwnedDataBlockIter},
};

/// Table reader that is optimized for consuming an entire table
pub struct Scanner {
    reader: BufReader<Box<dyn FsFile>>,
    iter: OwnedDataBlockIter,

    compression: CompressionType,
    block_count: usize,
    read_count: usize,

    global_seqno: SeqNo,

    encryption: Option<Arc<dyn EncryptionProvider>>,
    comparator: SharedComparator,

    /// Per-SST Page-ECC scheme from table metadata: data blocks omit the
    /// `block_flags` byte, so the scanner's block-read transform is told
    /// here the parity scheme (sizing + recovery). `None` = no parity.
    ecc: Option<crate::table::block::EccParams>,
    /// Per-SST per-KV-footer flag (`kv_checksum_algo.is_some()`): supplied to
    /// `from_loaded` so it strips the footer (data blocks omit the byte).
    has_kv_footer: bool,

    #[cfg(zstd_any)]
    zstd_dictionary: Option<Arc<crate::compression::ZstdDictionary>>,

    /// Table id of the SST being scanned; threaded through to
    /// per-block reads via `BlockIdentity`.
    table_id: crate::TableId,

    /// Whether this SST stores columnar (PAX) data blocks. When set, each
    /// fetched block is read as `BlockType::Columnar` and reconstructed into a
    /// row-major block so the scan iterator is unchanged.
    columnar: bool,
    /// Data-block restart interval, used to re-encode a reconstructed columnar
    /// block (matches the value the SST recorded).
    restart_interval: u8,

    /// Tight-space restriction lower bound of the view being scanned. The
    /// caller already positions the scan at the first live block; this filter
    /// drops the sub-bound entries of that STRADDLING block (their
    /// authoritative copies live in the superseding slice output).
    lower_bound: Option<crate::UserKey>,
    /// `true` while entries may still fall below `lower_bound`. Keys ascend,
    /// so once one entry reaches the bound the comparison is retired.
    filtering_below_bound: bool,

    /// The last key of the previous block, as `(user key, table-local seqno)`.
    /// A table's blocks hold disjoint, ascending key ranges, so each block's
    /// first key sorts after the previous block's last: a block read in
    /// another block's place, or a block read twice, breaks that order.
    prev_block_last: Option<(crate::UserKey, SeqNo)>,
    /// Where the current block stands in that check.
    block_order: BlockOrder,
}

/// The block-order check's state for the block being scanned.
#[derive(Clone, Copy, PartialEq, Eq)]
enum BlockOrder {
    /// The block's first entry has not been checked yet.
    FirstPending,
    /// The block's first entry sorted after the previous block's.
    Checked,
    /// A block failed the check; the scan yields nothing more.
    Failed,
}

impl Scanner {
    #[expect(
        clippy::too_many_arguments,
        reason = "scanner ctor takes one local per piece of state it needs to thread \
                  through fetch_next_block; collapsing them into a config struct would \
                  add an indirection without removing any per-call decision the caller \
                  makes about the values"
    )]
    pub fn new(
        fs: &Arc<dyn Fs>,
        path: &Path,
        block_count: usize,
        compression: CompressionType,
        global_seqno: SeqNo,
        encryption: Option<Arc<dyn EncryptionProvider>>,
        ecc: Option<crate::table::block::EccParams>,
        has_kv_footer: bool,
        #[cfg(zstd_any)] zstd_dictionary: Option<Arc<crate::compression::ZstdDictionary>>,
        comparator: SharedComparator,
        table_id: crate::TableId,
        columnar: bool,
        restart_interval: u8,
        start_offset: u64,
        lower_bound: Option<crate::UserKey>,
    ) -> crate::Result<Self> {
        // 2 MiB buffer matches RocksDB's `compaction_readahead_size`
        // default and is large enough that the kernel can fold the
        // sequential-scan readahead heuristic into a single big
        // userspace fill per ~500 typical (4 KiB) data blocks instead
        // of one syscall every 8 blocks. Picked to dominate any
        // pre-#133-Phase1c micro-cost: the loss is at most a few MB
        // of allocator overhead per concurrent compaction, and
        // compaction concurrency is bounded by the scheduler. HDD-
        // tuning beyond this would benefit from a configurable knob,
        // tracked as the configurable-readahead follow-up under #133.
        const SCANNER_READAHEAD_BYTES: usize = 2 * 1024 * 1024;
        let mut file = fs.open(path, &FsOpenOptions::new().read(true))?;
        // The scanner walks every block in order — tell the kernel so
        // it can ramp readahead aggressively and evict already-read
        // pages instead of pinning them. Best-effort: hint() is
        // advisory and any failure here would just leave the kernel
        // on the default heuristic, so drop the error rather than
        // failing the open.
        let _ = file.hint(FileHint::Sequential);
        // A tight-space restricted view starts at its first LIVE block: the
        // blocks below it are hole-punched (they read as zeros and cannot
        // decode). `block_count` was reduced by the caller to match.
        if start_offset > 0 {
            #[cfg(not(feature = "std"))]
            use crate::io::{Seek, SeekFrom};
            #[cfg(feature = "std")]
            use std::io::{Seek, SeekFrom};
            file.seek(SeekFrom::Start(start_offset))?;
        }
        let mut reader = BufReader::with_capacity(SCANNER_READAHEAD_BYTES, file);

        let block = Self::fetch_next_block(
            &mut reader,
            table_id,
            compression,
            encryption.as_deref(),
            ecc,
            has_kv_footer,
            columnar,
            restart_interval,
            #[cfg(zstd_any)]
            zstd_dictionary.as_deref(),
        )?;
        let cmp = comparator.clone();
        let iter = OwnedDataBlockIter::try_new(block, |b| b.try_iter(cmp))?;

        Ok(Self {
            reader,
            iter,

            compression,
            block_count,
            read_count: 1,

            global_seqno,
            encryption,
            comparator,

            ecc,
            has_kv_footer,

            #[cfg(zstd_any)]
            zstd_dictionary,

            table_id,
            columnar,
            restart_interval,

            filtering_below_bound: lower_bound.is_some(),
            lower_bound,

            prev_block_last: None,
            block_order: BlockOrder::FirstPending,
        })
    }

    /// Checks the first entry of a block against the previous block's last.
    ///
    /// # Errors
    ///
    /// Returns [`crate::Error::InvalidHeader`] when the block does not sort
    /// after the previous one.
    fn check_block_order(&self, first: &InternalValue) -> crate::Result<()> {
        if let Some((prev_key, prev_seqno)) = &self.prev_block_last {
            let ascends = match self.comparator.compare(prev_key, &first.key.user_key) {
                core::cmp::Ordering::Less => true,
                // Versions of one key descend by seqno across a block boundary,
                // with one legitimate tie: merge operands of one write batch
                // share its seqno and can straddle the boundary.
                core::cmp::Ordering::Equal => *prev_seqno >= first.key.seqno,
                core::cmp::Ordering::Greater => false,
            };
            if !ascends {
                return Err(crate::Error::InvalidHeader(
                    "data block does not sort after the block before it",
                ));
            }
        }
        Ok(())
    }

    /// The last key of the block being scanned, recorded before the scan
    /// moves to the next block.
    ///
    /// # Errors
    ///
    /// Propagates a block-decode failure.
    fn current_block_last(&self) -> crate::Result<Option<(crate::UserKey, SeqNo)>> {
        use crate::table::block::ParsedItem as _;
        let block = self.iter.borrow_owner();
        Ok(block
            .try_iter(self.comparator.clone())?
            .next_back()
            .map(|item| {
                let key = item.materialize(block.as_slice()).key;
                (key.user_key, key.seqno)
            }))
    }

    #[expect(
        clippy::too_many_arguments,
        reason = "per-block read threads each SST-level read parameter through; a config struct would add indirection without removing a caller decision"
    )]
    fn fetch_next_block(
        reader: &mut BufReader<Box<dyn FsFile>>,
        table_id: crate::TableId,
        compression: CompressionType,
        encryption: Option<&dyn EncryptionProvider>,
        ecc: Option<crate::table::block::EccParams>,
        has_kv_footer: bool,
        columnar: bool,
        restart_interval: u8,
        #[cfg(zstd_any)] zstd_dict: Option<&crate::compression::ZstdDictionary>,
    ) -> crate::Result<DataBlock> {
        // A columnar SST's blocks are tagged (and AAD-bound) as Columnar; read
        // them under that identity so verification / decryption matches.
        let block_type = if columnar {
            BlockType::Columnar
        } else {
            BlockType::Data
        };
        let block = Block::from_reader(
            reader,
            crate::table::block::BlockIdentity {
                table_id,
                block_type,
                dict_id: compression.dict_id(),
                window_log: 0,
            },
            &{
                // SST blocks omit the block_flags byte, so ECC presence is a
                // per-SST property: upgrade the transform to its `*Ecc`
                // variant when this table was written with Page ECC, so the
                // reader expects the parity trailer. Identity without the
                // `page_ecc` feature.
                let t = crate::table::block::BlockTransform::from_parts(
                    compression,
                    encryption,
                    #[cfg(zstd_any)]
                    zstd_dict,
                )?;
                if let Some(ecc) = ecc {
                    t.with_ecc(ecc)
                } else {
                    t
                }
            },
        );

        match block {
            Ok(block) => {
                // Columnar blocks are reconstructed into row blocks; a row block
                // must be BlockType::Data. `from_loaded` strips the per-KV
                // checksum footer when this SST carries one (per-SST
                // `has_kv_footer`, since data blocks omit the byte).
                if columnar {
                    return Self::reconstruct_columnar(&block, restart_interval);
                }
                if block.header.block_type != BlockType::Data {
                    return Err(crate::Error::InvalidTag((
                        "BlockType",
                        block.header.block_type.into(),
                    )));
                }

                DataBlock::from_loaded(block, has_kv_footer)
            }
            Err(e) => Err(e),
        }
    }

    /// Reconstructs a row-major [`DataBlock`] from a loaded columnar block. With
    /// the `columnar` feature this re-encodes the PAX block; without the feature
    /// a columnar SST cannot be read.
    fn reconstruct_columnar(block: &Block, restart_interval: u8) -> crate::Result<DataBlock> {
        #[cfg(feature = "columnar")]
        {
            DataBlock::from_columnar_block(&block.data, restart_interval)
        }
        #[cfg(not(feature = "columnar"))]
        {
            let _ = (block, restart_interval);
            Err(crate::Error::FeatureUnsupported("columnar"))
        }
    }
}

impl Iterator for Scanner {
    type Item = crate::Result<InternalValue>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.block_order == BlockOrder::Failed {
            return None;
        }
        loop {
            if let Some(mut item) = self.iter.next() {
                if self.block_order == BlockOrder::FirstPending {
                    if let Err(e) = self.check_block_order(&item) {
                        // Poisoned: nothing after a misordered block is served.
                        self.block_order = BlockOrder::Failed;
                        return Some(Err(e));
                    }
                    self.block_order = BlockOrder::Checked;
                }
                // Sub-bound entries of the straddling first block belong to
                // the slice output that superseded the punched prefix; keys
                // ascend, so the comparison retires at the first live entry.
                if self.filtering_below_bound {
                    if let Some(bound) = &self.lower_bound
                        && self.comparator.compare(&item.key.user_key, bound.as_ref())
                            == core::cmp::Ordering::Less
                    {
                        continue;
                    }
                    self.filtering_below_bound = false;
                }
                item.key.seqno += self.global_seqno;
                return Some(Ok(item));
            }

            if self.read_count >= self.block_count {
                return None;
            }

            match self.current_block_last() {
                Ok(Some(last)) => self.prev_block_last = Some(last),
                Ok(None) => {}
                Err(e) => {
                    self.block_order = BlockOrder::Failed;
                    return Some(Err(e));
                }
            }

            // Init new block
            let block = match Self::fetch_next_block(
                &mut self.reader,
                self.table_id,
                self.compression,
                self.encryption.as_deref(),
                self.ecc,
                self.has_kv_footer,
                self.columnar,
                self.restart_interval,
                #[cfg(zstd_any)]
                self.zstd_dictionary.as_deref(),
            ) {
                Ok(block) => block,
                Err(e) => {
                    self.read_count = self.block_count;
                    return Some(Err(e));
                }
            };
            let cmp = self.comparator.clone();
            match OwnedDataBlockIter::try_new(block, |b| b.try_iter(cmp)) {
                Ok(iter) => {
                    self.iter = iter;
                    self.read_count += 1;
                    self.block_order = BlockOrder::FirstPending;
                }
                Err(e) => {
                    // Poison the scanner so callers cannot silently skip
                    // the corrupt block and resume on later blocks.
                    self.read_count = self.block_count;
                    return Some(Err(e));
                }
            }
        }
    }
}