coordinode-lsm-tree 7.0.0

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

#[cfg(zstd_any)]
use crate::compression::CompressionProvider as _;
#[cfg(not(feature = "std"))]
use alloc::boxed::Box;

use super::meta::{META_VERSION, Metadata};
use crate::io::BufWriter;
#[cfg(not(feature = "std"))]
use crate::io::Write;
use crate::io::{LittleEndian, WriteBytesExt};
use crate::path::{Path, PathBuf};
use crate::{
    Checksum, CompressionType, KeyRange, SeqNo, TreeId, UserKey,
    checksum::ChecksummedWriter,
    fs::{Fs, FsFile, FsOpenOptions, SyncMode},
    vlog::BlobFileId,
};
#[cfg(feature = "std")]
use std::io::Write;

/// Safety cap on blob value size (256 MiB).
///
/// Enforced on the write path to prevent producing blobs that are
/// unreasonably large. The reader applies its own copy of this limit.
///
/// The blob-format definition site: `reader` and `scanner` import this rather
/// than duplicate it. Kept SEPARATE from `table::block`'s cap (a `u32`) on
/// purpose, since blocks and blobs are independent storage formats that may
/// diverge.
pub(super) const MAX_DECOMPRESSION_SIZE: usize = 256 * 1024 * 1024;

/// Returns `Err(DecompressedSizeTooLarge)` if `len > MAX_DECOMPRESSION_SIZE`.
pub(super) fn check_size_cap(len: usize) -> crate::Result<()> {
    if len > MAX_DECOMPRESSION_SIZE {
        return Err(crate::Error::DecompressedSizeTooLarge {
            declared: len as u64,
            limit: MAX_DECOMPRESSION_SIZE as u64,
        });
    }
    Ok(())
}

/// `value` as a blob frame stores it under `compression`: compressed with
/// `dictionary` for a dictionary codec, as it is for none. The output is
/// checked against the size cap, which applies to blob data whatever its
/// encoding.
#[cfg_attr(
    not(any(feature = "lz4", zstd_any)),
    expect(
        clippy::unnecessary_wraps,
        reason = "without a codec built in, a value is stored as it is and cannot fail"
    )
)]
#[cfg_attr(
    not(zstd_any),
    expect(
        clippy::elidable_lifetime_names,
        reason = "named to stay valid in the zstd build, where the dictionary is a second reference"
    )
)]
pub(super) fn encode_value<'v>(
    compression: CompressionType,
    #[cfg(zstd_any)] dictionary: Option<&crate::compression::ZstdDictionary>,
    #[cfg(zstd_any)] two_pass_seed: bool,
    value: &'v [u8],
) -> crate::Result<alloc::borrow::Cow<'v, [u8]>> {
    Ok(match compression {
        CompressionType::None => alloc::borrow::Cow::Borrowed(value),

        #[cfg(feature = "lz4")]
        CompressionType::Lz4 => {
            let compressed = lz4_flex::block::compress(value);
            check_size_cap(compressed.len())?;
            alloc::borrow::Cow::Owned(compressed)
        }

        #[cfg(zstd_any)]
        CompressionType::Zstd(level) => {
            let compressed =
                crate::compression::ZstdBackend::compress(value, level, two_pass_seed)?;
            check_size_cap(compressed.len())?;
            alloc::borrow::Cow::Owned(compressed)
        }

        #[cfg(zstd_any)]
        CompressionType::ZstdDict { level, dict_id } => {
            let dict = dictionary.ok_or(crate::Error::ZstdDictMismatch {
                expected: dict_id,
                got: None,
            })?;
            if dict.id() != dict_id {
                return Err(crate::Error::ZstdDictMismatch {
                    expected: dict_id,
                    got: Some(dict.id()),
                });
            }
            let compressed = crate::compression::ZstdBackend::compress_with_dict(
                value,
                level,
                dict,
                two_pass_seed,
            )?;
            check_size_cap(compressed.len())?;
            alloc::borrow::Cow::Owned(compressed)
        }
    })
}

// Note: these constants are `pub` for crate-internal use but the parent
// `vlog` module is NOT exported from `lib.rs`, so they are not public API.

/// Blob frame magic. This is the ONLY readable frame shape: the engine
/// supports exactly one on-disk format (see [`crate::FormatVersion`]) — a
/// frame under any other magic (including the retired pre-V5 `b"BLOB"`
/// layout without a header CRC) is corruption, not a compat case.
pub const BLOB_HEADER_MAGIC: &[u8] = b"BLO4";

/// Blob frame header length: 42 bytes, `header_crc` included. The fields sum to
/// 42: 4 (magic), 16 (u128 checksum), 8 (seqno), 2 (key len), 4 (real val len),
/// 4 (on-disk val len), 4 (header crc). Pinned by the reader tests
/// (`assert_eq!(BLOB_HEADER_LEN, 42)`).
pub const BLOB_HEADER_LEN: usize = BLOB_HEADER_MAGIC.len()
    + core::mem::size_of::<u128>() // Checksum
    + core::mem::size_of::<u64>() // SeqNo
    + core::mem::size_of::<u16>() // Key length
    + core::mem::size_of::<u32>() // Real value length
    + core::mem::size_of::<u32>() // On-disk value length
    + core::mem::size_of::<u32>(); // Header CRC

/// Compute the frame header CRC from header fields.
/// Returns a 4-byte truncated xxh3 hash.
#[expect(
    clippy::cast_possible_truncation,
    reason = "intentionally truncated to 4-byte CRC"
)]
pub(super) fn compute_header_crc(
    seqno: u64,
    key_len: u16,
    real_val_len: u32,
    on_disk_val_len: u32,
) -> u32 {
    let mut hasher = xxhash_rust::xxh3::Xxh3::default();
    hasher.update(&seqno.to_le_bytes());
    hasher.update(&key_len.to_le_bytes());
    hasher.update(&real_val_len.to_le_bytes());
    hasher.update(&on_disk_val_len.to_le_bytes());
    hasher.digest() as u32
}

/// Validate the frame header CRC: recompute from header fields and compare
/// against the stored value.
pub(super) fn validate_header_crc(
    seqno: u64,
    key_len: u16,
    real_val_len: u32,
    on_disk_val_len: u32,
    stored_crc: u32,
) -> crate::Result<()> {
    let recomputed_crc = compute_header_crc(seqno, key_len, real_val_len, on_disk_val_len);

    if stored_crc != recomputed_crc {
        return Err(crate::Error::HeaderCrcMismatch {
            recomputed: recomputed_crc,
            stored: stored_crc,
        });
    }

    Ok(())
}

/// Blob file writer
pub struct Writer {
    pub(crate) tree_id: TreeId,
    pub path: PathBuf,
    pub(crate) blob_file_id: BlobFileId,

    #[expect(clippy::struct_field_names)]
    writer: crate::sfa::Writer<ChecksummedWriter<BufWriter<Box<dyn FsFile>>>>,

    offset: u64,

    pub(crate) item_count: u64,
    pub(crate) written_blob_bytes: u64,
    pub(crate) uncompressed_bytes: u64,

    pub(crate) first_key: Option<UserKey>,
    pub(crate) last_key: Option<UserKey>,

    pub(crate) compression: CompressionType,

    /// Compression to RECORD in the blob file's metadata, when it must differ
    /// from [`Self::compression`]. Set for relocation / GC passthrough: the
    /// caller hands `write_raw` bytes that are ALREADY compressed (so
    /// `compression` stays `None` to skip re-compressing), but the file's
    /// metadata must still declare the real codec so a reopened reader
    /// decompresses correctly. `None` here means "record `compression`".
    pub(crate) metadata_compression_override: Option<CompressionType>,

    /// The lifetime class this file holds, recorded in its metadata.
    pub(crate) lifetime_class: u8,

    /// The age to record instead of the time the file is finished, for a file
    /// rebuilt from another store.
    pub(crate) created_at: Option<u128>,

    /// The tree's clock, read for the age when `created_at` does not set it;
    /// `None` reads the process clock.
    clock: Option<alloc::sync::Arc<dyn crate::Clock>>,

    /// Durability level for the final blob-file fsync. Default
    /// [`SyncMode::Normal`]; wired from `Config::sync_mode` via
    /// [`Self::use_sync_mode`].
    pub(crate) sync_mode: SyncMode,

    /// Bytes gathered between writebacks the writer starts while writing; `0`
    /// starts none. Set via [`Self::use_writeback_bytes`].
    writeback_bytes: u64,

    /// Where the last writeback ended: every byte before it was handed over.
    written_back: u64,

    /// Dictionary for `ZstdDict` compression.  Must be supplied when
    /// `compression` is [`CompressionType::ZstdDict`].
    #[cfg(zstd_any)]
    pub(crate) zstd_dictionary: Option<alloc::sync::Arc<crate::compression::ZstdDictionary>>,

    /// Whether zstd levels 19-22 run the `btultra2` two-pass seed. Defaults to
    /// `true`, the codec's own behaviour; set from the live runtime config via
    /// [`Self::use_zstd_two_pass_seed`].
    #[cfg(zstd_any)]
    pub(crate) zstd_two_pass_seed: bool,
}

impl Writer {
    /// Initializes a new blob file writer.
    ///
    /// Uses `create_new` (not `create+truncate`) because blob file IDs are
    /// monotonically unique — a path collision indicates a bug, not a retry.
    /// Orphaned files from crashes are cleaned up during recovery.
    ///
    /// # Errors
    ///
    /// Will return `Err` if an IO error occurs.
    #[doc(hidden)]
    pub fn new<P: AsRef<Path>>(
        path: P,
        blob_file_id: BlobFileId,
        tree_id: TreeId,
        fs: &dyn Fs,
    ) -> crate::Result<Self> {
        let path = path.as_ref();

        let file = fs.open(path, &FsOpenOptions::new().write(true).create_new(true))?;
        // From here the file exists and is OURS (`create_new` proved it), so
        // this constructor owns cleanup: an init failure below must not leak a
        // partial file, and a CALLER must never remove `path` on a constructor
        // error — an open failure above created nothing, so caller-side
        // cleanup would race a concurrent creator and delete a file it does
        // not own.
        let mut writer =
            crate::sfa::Writer::from_writer(ChecksummedWriter::new(BufWriter::new(file)));
        if let Err(e) = writer.start("data") {
            drop(writer);
            // Best-effort cleanup of the partial file we own, but do not swallow
            // a removal failure silently: log it with context so a leaked
            // partial blob file is diagnosable. The original start error is still
            // what the caller sees.
            if let Err(remove_err) = fs.remove_file(path) {
                log::warn!(
                    "failed to remove partial blob file {} after a writer-init error: {remove_err}",
                    path.display(),
                );
            }
            return Err(e.into());
        }

        Ok(Self {
            tree_id,
            path: path.into(),
            blob_file_id,

            writer,

            offset: 0,
            item_count: 0,
            written_blob_bytes: 0,
            uncompressed_bytes: 0,

            first_key: None,
            last_key: None,

            compression: CompressionType::None,
            metadata_compression_override: None,
            lifetime_class: 0,
            created_at: None,
            clock: None,
            sync_mode: SyncMode::Normal,
            writeback_bytes: 0,
            written_back: 0,

            #[cfg(zstd_any)]
            zstd_dictionary: None,
            #[cfg(zstd_any)]
            zstd_two_pass_seed: true,
        })
    }

    pub fn use_compression(mut self, compressor: CompressionType) -> Self {
        self.compression = compressor;
        self
    }

    /// Wires the tree's `Config::sync_mode` into the final blob-file fsync.
    #[must_use]
    pub fn use_sync_mode(mut self, sync_mode: SyncMode) -> Self {
        self.sync_mode = sync_mode;
        self
    }

    /// Reads the time on `clock`, the tree's own, instead of the process clock.
    #[must_use]
    pub(crate) fn use_clock(mut self, clock: Option<alloc::sync::Arc<dyn crate::Clock>>) -> Self {
        self.clock = clock;
        self
    }

    /// Starts writing back every `bytes` the file gathers while it is written
    /// ([`FsFile::start_writeback`]), so its final sync is short; `0` starts
    /// none.
    #[must_use]
    pub fn use_writeback_bytes(mut self, bytes: u64) -> Self {
        self.writeback_bytes = bytes;
        self
    }

    /// Records `class` as the lifetime class of this file and tells the device
    /// how long its data is expected to live, when `lifetime` says.
    ///
    /// The device hint is advisory: a filesystem or kernel that refuses it
    /// leaves the file written exactly as without it.
    pub(crate) fn set_lifetime_class(
        &mut self,
        class: u8,
        lifetime: Option<crate::fs::WriteLifetime>,
    ) {
        self.lifetime_class = class;
        if let Some(lifetime) = lifetime {
            let file = self.writer.get_mut().inner_mut().get_ref();
            if let Err(error) = file.set_write_lifetime(lifetime) {
                log::debug!(
                    "write lifetime hint refused for {}: {error}",
                    self.path.display()
                );
            }
        }
    }

    /// Starts writing back what the file wrote since the last writeback, once
    /// at least `writeback_bytes` of it gathered.
    fn writeback_if_due(&mut self) -> crate::Result<()> {
        // `written_back` only ever takes a value of `offset`, which grows.
        let pending = self.offset - self.written_back;
        if self.writeback_bytes == 0 || pending < self.writeback_bytes {
            return Ok(());
        }
        #[cfg(not(feature = "std"))]
        use crate::io::Write;
        #[cfg(feature = "std")]
        use std::io::Write;
        // The range is in the buffer until it is flushed to the file.
        let buffered = self.writer.get_mut().inner_mut();
        buffered.flush()?;
        crate::fs::hint_writeback(
            &**buffered.get_ref(),
            self.written_back,
            pending,
            &mut self.writeback_bytes,
        );
        self.written_back = self.offset;
        Ok(())
    }

    /// Selects whether zstd levels 19-22 run the `btultra2` two-pass seed.
    ///
    /// See [`crate::runtime_config::RuntimeConfig::zstd_two_pass_seed`]. Has no
    /// effect below level 19, where another strategy applies.
    #[cfg(zstd_any)]
    #[must_use]
    pub fn use_zstd_two_pass_seed(mut self, enabled: bool) -> Self {
        self.zstd_two_pass_seed = enabled;
        self
    }

    /// Provides the zstd dictionary for [`CompressionType::ZstdDict`] writes.
    ///
    /// Must be called before writing any blobs when the compression type is
    /// `ZstdDict`.  Passing `None` clears a previously set dictionary.
    #[cfg(zstd_any)]
    #[must_use]
    pub fn use_zstd_dictionary(
        mut self,
        dict: Option<alloc::sync::Arc<crate::compression::ZstdDictionary>>,
    ) -> Self {
        self.zstd_dictionary = dict;
        self
    }

    /// Returns the current offset in the file.
    #[must_use]
    pub(crate) fn offset(&self) -> u64 {
        self.offset
    }

    /// Returns the blob file ID.
    #[must_use]
    pub(crate) fn blob_file_id(&self) -> BlobFileId {
        self.blob_file_id
    }

    /// Skips `len` bytes of the data section, so the next value lands `len`
    /// bytes later: for a file rebuilt from one whose consumed prefix was
    /// reclaimed, whose live values keep the offsets their handles name. The
    /// gap is never written: a filesystem with sparse files leaves it
    /// unallocated, and the file is restricted past it. Its bytes are not part
    /// of the writer's checksum, which a restricted file does not use.
    #[cfg(feature = "std")]
    pub(crate) fn write_filler(&mut self, len: u64) -> crate::Result<()> {
        use std::io::Seek as _;
        let skip =
            i64::try_from(len).map_err(|_| crate::Error::InvalidHeader("blob file frontier"))?;
        self.writer
            .get_mut()
            .seek(std::io::SeekFrom::Current(skip))?;
        self.offset += len;
        Ok(())
    }

    pub(crate) fn write_raw(
        &mut self,
        key: &[u8],
        seqno: SeqNo,
        value: &[u8],
        uncompressed_len: u32,
    ) -> crate::Result<u32> {
        assert!(!key.is_empty());
        assert!(u16::try_from(key.len()).is_ok());
        assert!(u32::try_from(value.len()).is_ok());

        check_size_cap(uncompressed_len as usize)?;
        check_size_cap(value.len())?;

        // Perform compression before mutating writer state, so an error
        // leaves the writer consistent.
        let value = encode_value(
            self.compression,
            #[cfg(zstd_any)]
            self.zstd_dictionary.as_deref(),
            #[cfg(zstd_any)]
            self.zstd_two_pass_seed,
            value,
        )?;
        self.append(key, seqno, &value, uncompressed_len)
    }

    /// Writes `encoded`, a value already encoded under the codec this file
    /// records, without encoding it again. `uncompressed_len` is its decoded
    /// length.
    pub(crate) fn write_encoded(
        &mut self,
        key: &[u8],
        seqno: SeqNo,
        encoded: &[u8],
        uncompressed_len: u32,
    ) -> crate::Result<u32> {
        assert!(!key.is_empty());
        assert!(u16::try_from(key.len()).is_ok());

        check_size_cap(uncompressed_len as usize)?;
        check_size_cap(encoded.len())?;
        self.append(key, seqno, encoded, uncompressed_len)
    }

    /// Appends the frame of `value`, already in its on-disk form.
    fn append(
        &mut self,
        key: &[u8],
        seqno: SeqNo,
        value: &[u8],
        uncompressed_len: u32,
    ) -> crate::Result<u32> {
        // Ensure the compressed value length fits in u32 before we write it
        // to disk as a 32-bit length. This prevents truncation if compression
        // expands the payload (possible for incompressible data near u32 boundary).
        let compressed_len_u32 = u32::try_from(value.len())
            .map_err(|_| crate::io::Error::other("compressed value length exceeds u32::MAX"))?;

        if self.first_key.is_none() {
            self.first_key = Some(key.into());
        }
        self.last_key = Some(key.into());

        self.uncompressed_bytes += u64::from(uncompressed_len);

        // NOTE:
        // BLOB HEADER LAYOUT
        //
        // [MAGIC_BYTES; 4B]    - BLOB_HEADER_MAGIC (b"BLO4")
        // [Checksum; 16B]      - xxh3_128(key + value + header_crc_le)
        // [Seqno; 8B]
        // [key len; 2B]
        // [real val len; 4B]
        // [on-disk val len; 4B]
        // [header_crc; 4B]     - truncated xxh3(seqno + key_len + real_val_len + on_disk_val_len)
        // [...key; ?]
        // [...val; ?]

        #[expect(clippy::cast_possible_truncation, reason = "keys are u16 length max")]
        let header_crc = compute_header_crc(
            seqno,
            key.len() as u16,
            uncompressed_len,
            compressed_len_u32,
        );

        // Data checksum includes header_crc bytes so that changes to header
        // fields without correspondingly updating the data checksum will be
        // detected as an inconsistency between header and data.
        let checksum = {
            let mut hasher = xxhash_rust::xxh3::Xxh3::default();
            hasher.update(key);
            hasher.update(value);
            hasher.update(&header_crc.to_le_bytes());
            hasher.digest128()
        };

        // Write header
        self.writer.write_all(BLOB_HEADER_MAGIC)?;

        // Write data checksum
        self.writer.write_u128::<LittleEndian>(checksum)?;

        // Write seqno
        self.writer.write_u64::<LittleEndian>(seqno)?;

        #[expect(clippy::cast_possible_truncation, reason = "keys are u16 length max")]
        self.writer.write_u16::<LittleEndian>(key.len() as u16)?;

        // Write uncompressed value length
        self.writer.write_u32::<LittleEndian>(uncompressed_len)?;

        // Write compressed (on-disk) value length
        self.writer.write_u32::<LittleEndian>(compressed_len_u32)?;

        // Write header CRC
        self.writer.write_u32::<LittleEndian>(header_crc)?;

        self.writer.write_all(key)?;
        self.writer.write_all(value)?;

        // Update offset
        self.offset += BLOB_HEADER_LEN as u64;
        self.offset += key.len() as u64;
        self.offset += value.len() as u64;

        // Update metadata
        self.written_blob_bytes += value.len() as u64;
        self.item_count += 1;
        self.writeback_if_due()?;

        // TODO: if we store the offset before writing, we can return a vhandle here
        // instead of needing to call offset() and blob_file_id() before write()

        Ok(compressed_len_u32)
    }

    /// Writes an item into the file.
    ///
    /// Items need to be written in key order.
    ///
    /// # Errors
    ///
    /// Will return `Err` if an IO error occurs.
    ///
    /// Will return `Err(Error::DecompressedSizeTooLarge { .. })` if the value exceeds the 256 MiB limit.
    ///
    /// # Panics
    ///
    /// Panics if the key length is empty or greater than 2^16, or the value length is greater than 2^32.
    pub fn write(&mut self, key: &[u8], seqno: SeqNo, value: &[u8]) -> crate::Result<u32> {
        #[expect(clippy::cast_possible_truncation, reason = "values are u32 max")]
        self.write_raw(key, seqno, value, value.len() as u32)
    }

    pub(crate) fn finish(mut self) -> crate::Result<(Metadata, Checksum)> {
        self.writer.start("meta")?;

        // Write metadata
        let metadata = Metadata {
            id: self.blob_file_id,
            version: META_VERSION,
            created_at: self
                .created_at
                .unwrap_or_else(|| crate::time::now_on(self.clock.as_ref()).as_nanos()),
            item_count: self.item_count,
            total_compressed_bytes: self.written_blob_bytes,
            total_uncompressed_bytes: self.uncompressed_bytes,
            #[expect(clippy::expect_used, reason = "should have written at least 1 item")]
            key_range: KeyRange::new((
                self.first_key
                    .clone()
                    .expect("should have written at least 1 item"),
                self.last_key
                    .clone()
                    .expect("should have written at least 1 item"),
            )),
            // Record the override (relocation passthrough) when set, so a
            // reopened reader sees the real codec of the already-compressed
            // bytes; otherwise record this writer's own compression.
            compression: self
                .metadata_compression_override
                .unwrap_or(self.compression),
            lifetime_class: self.lifetime_class,
        };
        metadata.encode_into(&mut self.writer)?;

        let mut checksum = self.writer.into_inner()?;
        // Data and size are all a read of the file needs, so `sync_data`.
        FsFile::sync_data_with(&**checksum.inner_mut().get_mut(), self.sync_mode)?;
        let checksum = checksum.checksum();

        Ok((metadata, checksum))
    }
}

#[cfg(test)]
mod tests;