Skip to main content

turbo_persistence/
static_sorted_file_builder.rs

1use std::{
2    borrow::Cow,
3    collections::VecDeque,
4    io::{BufWriter, Write},
5    path::{Path, PathBuf},
6};
7
8use anyhow::{Context, Result};
9use byteorder::{BE, ByteOrder, WriteBytesExt};
10use either::Either;
11use fs_err::File;
12
13use crate::{
14    Compression,
15    compression::{Compressor, checksum_block},
16    constants::{MAX_INLINE_VALUE_SIZE, MAX_SMALL_VALUE_SIZE, MIN_SMALL_VALUE_BLOCK_SIZE},
17    meta_file::MetaEntryFlags,
18    static_sorted_file::{
19        BLOB_VALUE_REF_SIZE, BLOCK_TYPE_INDEX, FIXED_KEY_BLOCK_MIXED_VALUE_TYPE, FixedRegions,
20        KEY_BLOCK_ENTRY_TYPE_BLOB, KEY_BLOCK_ENTRY_TYPE_INLINE_MIN,
21        KEY_BLOCK_ENTRY_TYPE_KEY_DELETED, KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN,
22        KEY_BLOCK_ENTRY_TYPE_MEDIUM, KEY_BLOCK_ENTRY_TYPE_SMALL,
23        KEY_BLOCK_TABLE_ENTRY_SIZE_NO_HASH, KEY_DELETED_REF_SIZE, KeyBlockLayout,
24        MEDIUM_VALUE_REF_SIZE, SMALL_VALUE_REF_SIZE, key_block_table_stride,
25    },
26};
27
28/// Size of the per-block header on disk: 4 bytes uncompressed_size + 4 bytes CRC32 checksum.
29pub const BLOCK_HEADER_SIZE: usize = 8;
30
31/// The maximum number of entries that should go into a single key block
32const MAX_KEY_BLOCK_ENTRIES: usize = MAX_KEY_BLOCK_SIZE / KEY_BLOCK_ENTRY_META_OVERHEAD;
33/// The maximum bytes that should go into a single key block
34// Note this must fit into 3 bytes length
35const MAX_KEY_BLOCK_SIZE: usize = 16 * 1024;
36/// Overhead of bytes that should be counted for entries in a key block in addition to the key size.
37/// This covers the worst case (small values):
38/// - 1 byte type (key block header)
39/// - 3 bytes position (key block header)
40/// - 8 bytes hash (optional, but unknown at collection time)
41/// - 2 bytes block index
42/// - 2 bytes size
43/// - 4 bytes position in block
44const KEY_BLOCK_ENTRY_META_OVERHEAD: usize = 20;
45/// The aimed false positive rate for the AMQF
46const AMQF_FALSE_POSITIVE_RATE: f64 = 0.01;
47/// Assumed average small value size for pre-allocation estimates.
48/// Intentionally conservative (small values range from MAX_INLINE_VALUE_SIZE+1 to
49/// MAX_SMALL_VALUE_SIZE = 4096): a low estimate over-counts value blocks, which is
50/// preferable to under-allocating vectors.
51const AVG_SMALL_VALUE_SIZE: usize = 64;
52
53/// Safety margin for block index capacity estimation in
54/// [`StreamingSstWriter::has_block_index_capacity`]. Accounts for rounding in the entry-count and
55/// byte-size based estimates of pending key blocks.
56const BLOCK_INDEX_CAPACITY_BUFFER: usize = 16;
57
58/// Minimum key size (in bytes) for attempting LZ4 compression on key blocks.
59///
60/// For small keys (below this threshold), compression is unlikely to find enough to work with and
61/// only wastes CPU time, so we skip the attempt entirely.
62const MIN_KEY_SIZE_FOR_COMPRESSION: usize = 16;
63
64/// Maximum key length that can use fixed-size key block layout.
65///
66/// The on-disk fixed-key header stores the key size as a single byte, so keys longer than this
67/// fall back to variable-size layout.
68const MAX_FIXED_KEY_LEN: usize = u8::MAX as usize;
69
70/// Maximum value size that can use fixed-size key block layout.
71///
72/// Mixed-type fixed blocks store the value size in a single header byte, since it can no longer be
73/// derived from a single shared entry type.
74const MAX_FIXED_VAL_SIZE: usize = u8::MAX as usize;
75
76/// Newtype for the key block entry type byte.
77///
78/// This encodes what kind of value reference an entry has (small, medium, blob, deleted, or
79/// inline with embedded length). See `KEY_BLOCK_ENTRY_TYPE_*` constants.
80#[derive(Clone, Copy, PartialEq, Eq, Debug)]
81struct EntryType(u8);
82
83/// Tracks whether a key block's entries are uniform enough for fixed-size layout.
84///
85/// Fixed layout needs a uniform *stride*, which requires a uniform key length and a uniform value
86/// size. A uniform value *type* is a stronger condition that additionally lets the type be hoisted
87/// into the block header; when types differ but sizes agree, the type is stored per entry instead
88/// (1 byte, still cheaper than the 4-byte offset table entry a variable block would need).
89///
90/// State transitions:
91/// - `Unknown` → first entry → `Fixed`
92/// - `Fixed` + matching key_len and value type → stays `Fixed`
93/// - `Fixed` + matching key_len and value *size* → `Fixed` with `value_type: None`
94/// - `Fixed` + mismatched key_len or value size → `Variable`
95/// - `Variable` → stays `Variable`
96#[derive(Clone, Copy)]
97enum KeyBlockFormat {
98    /// No entries yet — format undetermined.
99    Unknown,
100    /// All entries so far have uniform key length and value size.
101    Fixed {
102        key_len: u8,
103        val_size: u8,
104        /// The shared entry type, or `None` if entries have differing types of the same size.
105        value_type: Option<EntryType>,
106    },
107    /// Entries have mixed key lengths or value sizes; must use offset table.
108    Variable,
109}
110
111impl KeyBlockFormat {
112    /// Updates the format after seeing an entry with the given key length and value type.
113    ///
114    /// A `Fixed` state is only reachable when all entries have matching key length and value size,
115    /// and the key length fits in a u8 (required by the on-disk header).
116    fn update(&mut self, key_len: usize, value_type: EntryType) {
117        let val_size = value_type_val_size(value_type);
118        *self = match *self {
119            KeyBlockFormat::Unknown => {
120                if key_len <= MAX_FIXED_KEY_LEN && val_size <= MAX_FIXED_VAL_SIZE {
121                    KeyBlockFormat::Fixed {
122                        key_len: key_len as u8,
123                        val_size: val_size as u8,
124                        value_type: Some(value_type),
125                    }
126                } else {
127                    KeyBlockFormat::Variable
128                }
129            }
130            KeyBlockFormat::Fixed {
131                key_len: k,
132                val_size: s,
133                value_type: v,
134            } if k as usize == key_len && s as usize == val_size => KeyBlockFormat::Fixed {
135                key_len: k,
136                val_size: s,
137                // Collapse to `None` as soon as two entries disagree on type.
138                value_type: v.filter(|v| *v == value_type),
139            },
140            KeyBlockFormat::Fixed { .. } | KeyBlockFormat::Variable => KeyBlockFormat::Variable,
141        };
142    }
143}
144
145/// Copy-able snapshot of the accumulator state needed by [`flush_key_block`].
146#[derive(Clone, Copy)]
147struct KeyBlockFlushInfo {
148    max_key_len: usize,
149    min_key_len: usize,
150    format: KeyBlockFormat,
151}
152
153impl KeyBlockFlushInfo {
154    /// The shared key length when every entry in the block has the same one, else `None`.
155    fn uniform_key_len(&self) -> Option<usize> {
156        (self.min_key_len == self.max_key_len).then_some(self.max_key_len)
157    }
158}
159
160/// Tracks the accumulated state of the current incomplete key block.
161///
162/// During streaming, this sits on [`StreamingSstWriter`] and tracks the tail of the resolved
163/// prefix. Entries are added one at a time; when [`should_flush`](Self::should_flush) returns
164/// `true`, the caller should flush the block and call [`reset`](Self::reset).
165struct KeyBlockAccumulator {
166    /// Accumulated byte size (keys + per-entry overhead) of entries in this block.
167    size: usize,
168    /// Number of entries accumulated so far.
169    entry_count: usize,
170    /// Maximum key length among accumulated entries (determines whether hashes are stored).
171    max_key_len: usize,
172    min_key_len: usize,
173    /// Hash of the most recently added entry (used to avoid splitting entries with equal hashes
174    /// across blocks).
175    last_hash: u64,
176    /// Whether the block qualifies for fixed-size layout.
177    format: KeyBlockFormat,
178}
179
180impl KeyBlockAccumulator {
181    fn new() -> Self {
182        Self {
183            size: 0,
184            entry_count: 0,
185            max_key_len: 0,
186            min_key_len: usize::MAX,
187            last_hash: 0,
188            format: KeyBlockFormat::Unknown,
189        }
190    }
191
192    /// Records a new entry in the accumulator.
193    fn add(&mut self, key_len: usize, key_hash: u64, value_type: EntryType) {
194        self.size += key_len + KEY_BLOCK_ENTRY_META_OVERHEAD;
195        self.max_key_len = self.max_key_len.max(key_len);
196        self.min_key_len = self.min_key_len.min(key_len);
197        self.entry_count += 1;
198        self.last_hash = key_hash;
199        self.format.update(key_len, value_type);
200    }
201
202    /// Snapshots the state needed by `flush_key_block`.
203    fn flush_info(&self) -> KeyBlockFlushInfo {
204        KeyBlockFlushInfo {
205            max_key_len: self.max_key_len,
206            min_key_len: self.min_key_len,
207            format: self.format,
208        }
209    }
210
211    /// Returns `true` if the block should be flushed before adding an entry with the given key
212    /// length and hash. Returns `false` for empty blocks and when the next entry shares its hash
213    /// with the current last entry (to avoid splitting equal-hash runs).
214    fn should_flush(&self, next_key_len: usize, next_key_hash: u64) -> bool {
215        if self.entry_count == 0 {
216            return false;
217        }
218        let would_exceed_size =
219            self.size + next_key_len + KEY_BLOCK_ENTRY_META_OVERHEAD > MAX_KEY_BLOCK_SIZE;
220        let would_exceed_entries = self.entry_count >= MAX_KEY_BLOCK_ENTRIES;
221        // Never split entries with the same hash across blocks.
222        (would_exceed_size || would_exceed_entries) && self.last_hash != next_key_hash
223    }
224
225    /// Resets the accumulator for a new key block.
226    fn reset(&mut self) {
227        self.size = 0;
228        self.entry_count = 0;
229        self.max_key_len = 0;
230        self.min_key_len = usize::MAX;
231        self.format = KeyBlockFormat::Unknown;
232        // last_hash is intentionally not reset -- it is overwritten on the next add() call.
233    }
234}
235
236/// Chooses a key block's layout from the longest key it holds.
237fn choose_layout(max_key_len: usize) -> KeyBlockLayout {
238    // Short keys are cheap enough to compare directly that storing an 8-byte hash per entry costs
239    // more space than the comparison saves, so those blocks omit it and reorder entries by key.
240    if max_key_len > 32 {
241        KeyBlockLayout::HashThenKey
242    } else {
243        KeyBlockLayout::KeyOnly
244    }
245}
246
247#[inline]
248fn be_key_u32(key: &[u8]) -> u32 {
249    u32::from_be_bytes(key.try_into().expect("4-byte key"))
250}
251
252#[inline]
253fn be_key_u64(key: &[u8]) -> u64 {
254    u64::from_be_bytes(key.try_into().expect("8-byte key"))
255}
256
257/// Trait for entries from that SST files can be created
258pub trait Entry {
259    /// Returns the hash of the key
260    fn key_hash(&self) -> u64;
261    /// Returns the length of the key
262    fn key_len(&self) -> usize;
263    /// Returns the key's bytes.
264    fn key_bytes(&self) -> &[u8];
265    /// Writes the key to a buffer
266    fn write_key_to(&self, buf: &mut Vec<u8>) {
267        buf.extend_from_slice(self.key_bytes());
268    }
269
270    /// Returns the value
271    fn value(&self) -> EntryValue<'_>;
272}
273
274impl<E: Entry> Entry for &E {
275    fn key_hash(&self) -> u64 {
276        (*self).key_hash()
277    }
278    fn key_len(&self) -> usize {
279        (*self).key_len()
280    }
281    fn key_bytes(&self) -> &[u8] {
282        (*self).key_bytes()
283    }
284    fn value(&self) -> EntryValue<'_> {
285        (*self).value()
286    }
287}
288
289/// Reference to a value
290#[derive(Copy, Clone)]
291pub enum EntryValue<'l> {
292    /// Inline value stored directly in the key block.
293    Inline { value: &'l [u8] },
294    /// Small-sized value. They are stored in shared value blocks.
295    Small { value: &'l [u8] },
296    /// Medium-sized value. They are stored in their own value block.
297    Medium { value: &'l [u8] },
298    /// Medium-sized value. They are stored in their own value block. In the raw form as on disk.
299    MediumRaw {
300        /// The uncompressed size of the block data. `0` means the block is stored uncompressed
301        /// (and thus the size is the `len` of the block)
302        uncompressed_size: u32,
303        /// CRC32 checksum of the on-disk block data (after compression).
304        checksum: u32,
305        block: &'l [u8],
306    },
307    /// Large-sized value. They are stored in a blob file.
308    Large { blob: u32 },
309    /// Tombstone. The value was removed.
310    KeyDeleted,
311    /// Key-value tombstone. Only the one carried value was removed; other values for the same key
312    /// survive. MultiValue families only. The value must be at most [`MAX_INLINE_VALUE_SIZE`]
313    /// bytes.
314    KeyValueDeleted { value: &'l [u8] },
315}
316
317#[derive(Debug, Clone)]
318pub struct StaticSortedFileBuilderMeta<'a> {
319    /// The minimum hash of the keys in the SST file
320    pub min_hash: u64,
321    /// The maximum hash of the keys in the SST file
322    pub max_hash: u64,
323    /// The AMQF data
324    pub amqf: Cow<'a, [u8]>,
325    /// The number of blocks in the SST file
326    pub block_count: u16,
327    /// The file size of the SST file
328    pub size: u64,
329    /// The status flags for this SST file
330    pub flags: MetaEntryFlags,
331    /// The number of entries in the SST file
332    pub entries: u64,
333}
334
335/// Writes an SST file from a pre-sorted slice of entries.
336///
337/// Entries must be sorted in (key-hash, key) order, the same contract as
338/// [`StreamingSstWriter::add`].
339///
340/// This is a convenience wrapper around [`StreamingSstWriter`] for callers that already have all
341/// entries in memory.
342// TODO: Consider adding a variant that takes ownership (Vec<E> or drain iterator)
343// to free entry memory as blocks are written.
344pub fn write_static_stored_file<E: Entry>(
345    entries: &[E],
346    file: &Path,
347    flags: MetaEntryFlags,
348    compression: Compression,
349) -> Result<(StaticSortedFileBuilderMeta<'static>, File)> {
350    debug_assert!(
351        entries
352            .iter()
353            .map(|e| (e.key_hash(), e.key_bytes()))
354            .is_sorted()
355    );
356    let mut writer = StreamingSstWriter::new(file, flags, entries.len() as u64, compression)?;
357    for entry in entries {
358        if let Err(err) = writer.add(entry) {
359            writer.cancel();
360            return Err(err);
361        }
362    }
363    writer.close()
364}
365
366// ---------------------------------------------------------------------------
367// Block I/O helpers (free functions for borrow-checker friendliness)
368// ---------------------------------------------------------------------------
369
370/// Writes a raw (already-formatted) block to the file. Returns the block index assigned.
371///
372/// `uncompressed_size` is the original uncompressed size of the block data, or `0` if the block
373/// is stored uncompressed.
374fn write_raw_block_to_file(
375    file: &mut BufWriter<File>,
376    block_offsets: &mut Vec<u32>,
377    uncompressed_size: u32,
378    checksum: u32,
379    block: &[u8],
380) -> Result<u16> {
381    let block_index: u16 = block_offsets
382        .len()
383        .try_into()
384        .expect("Block index overflow");
385
386    let len: u32 = (block.len() + BLOCK_HEADER_SIZE).try_into().unwrap();
387    let offset = block_offsets
388        .last()
389        .copied()
390        .unwrap_or_default()
391        .checked_add(len)
392        .expect("Block offset overflow");
393    block_offsets.push(offset);
394
395    file.write_u32::<BE>(uncompressed_size)
396        .context("Failed to write uncompressed size")?;
397    file.write_u32::<BE>(checksum)
398        .context("Failed to write checksum")?;
399    file.write_all(block)
400        .context("Failed to write block data")?;
401    Ok(block_index)
402}
403
404/// Writes a block to the file, optionally compressing it. Returns the block index assigned.
405fn write_block_to_file(
406    file: &mut BufWriter<File>,
407    compress_buffer: &mut Vec<u8>,
408    block_offsets: &mut Vec<u32>,
409    block: &[u8],
410    try_compress: bool,
411    compressor: &mut Compressor,
412) -> Result<u16> {
413    let (uncompressed_size, data_to_write): (u32, &[u8]) = if try_compress {
414        compressor.compress_into_buffer(block, compress_buffer)?;
415        // Same threshold as LevelDB/RocksDB: require at least 12.5% savings.
416        if compress_buffer.len() < block.len() - (block.len() / 8) {
417            (block.len().try_into().unwrap(), compress_buffer)
418        } else {
419            (0, block)
420        }
421    } else {
422        (0, block)
423    };
424
425    // Checksum is computed on the on-disk data (after compression).
426    let checksum = checksum_block(data_to_write);
427
428    write_raw_block_to_file(
429        file,
430        block_offsets,
431        uncompressed_size,
432        checksum,
433        data_to_write,
434    )
435}
436
437// ---------------------------------------------------------------------------
438// StreamingSstWriter
439// ---------------------------------------------------------------------------
440
441/// Where a key entry's value lives (or will live once the small block flushes).
442enum ValueRef {
443    /// Value in a known small value block (already flushed).
444    Small {
445        block_index: u16,
446        offset: u32,
447        size: u16,
448    },
449    /// Value is in a small value block that hasn't been written yet. Will be resolved in-place
450    /// to [`ValueRef::Small`] when the small block is flushed.
451    PendingSmall {
452        #[cfg(debug_assertions)]
453        small_block_id: u16,
454        offset: u32,
455        size: u16,
456    },
457    /// Medium value already written to its own block.
458    Medium { block_index: u16 },
459    /// Inline value (stored directly in the key block).
460    Inline {
461        data: [u8; MAX_INLINE_VALUE_SIZE],
462        len: u8,
463    },
464    /// Large blob stored externally.
465    Blob { blob_id: u32 },
466    /// Tombstone.
467    KeyDeleted,
468    /// Key-value tombstone: deletes only the carried value from the key's group. The value is
469    /// stored inline, exactly like [`ValueRef::Inline`].
470    KeyValueDeleted {
471        data: [u8; MAX_INLINE_VALUE_SIZE],
472        len: u8,
473    },
474}
475
476impl ValueRef {
477    /// Returns the key block entry type for this value reference.
478    fn entry_type(&self) -> EntryType {
479        EntryType(match self {
480            ValueRef::Small { .. } | ValueRef::PendingSmall { .. } => KEY_BLOCK_ENTRY_TYPE_SMALL,
481            ValueRef::Medium { .. } => KEY_BLOCK_ENTRY_TYPE_MEDIUM,
482            ValueRef::Inline { len, .. } => KEY_BLOCK_ENTRY_TYPE_INLINE_MIN + *len,
483            ValueRef::Blob { .. } => KEY_BLOCK_ENTRY_TYPE_BLOB,
484            ValueRef::KeyDeleted => KEY_BLOCK_ENTRY_TYPE_KEY_DELETED,
485            ValueRef::KeyValueDeleted { len, .. } => {
486                KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN + *len
487            }
488        })
489    }
490
491    /// Writes the value bytes for this reference to a buffer.
492    ///
493    /// This is the shared serialization logic used by both variable-size and fixed-size key block
494    /// builders.
495    fn write_value_to(&self, buffer: &mut Vec<u8>) {
496        match self {
497            ValueRef::Small {
498                block_index,
499                offset,
500                size,
501            } => {
502                let mut scratch = [0u8; 8];
503                BE::write_u16(&mut scratch, *block_index);
504                BE::write_u16(&mut scratch[2..], *size);
505                BE::write_u32(&mut scratch[4..], *offset);
506                buffer.extend(&scratch);
507            }
508            ValueRef::Medium { block_index } => {
509                let mut scratch = [0u8; 2];
510                BE::write_u16(&mut scratch, *block_index);
511                buffer.extend(scratch);
512            }
513            ValueRef::Inline { data, len } => {
514                buffer.extend(&data[..*len as usize]);
515            }
516            ValueRef::Blob { blob_id } => {
517                let mut scratch = [0u8; 4];
518                BE::write_u32(&mut scratch, *blob_id);
519                buffer.extend(scratch);
520            }
521            ValueRef::KeyDeleted => { /* no value bytes */ }
522            ValueRef::KeyValueDeleted { data, len } => {
523                buffer.extend(&data[..*len as usize]);
524            }
525            ValueRef::PendingSmall { .. } => {
526                unreachable!("PendingSmall should have been resolved");
527            }
528        }
529    }
530}
531
532struct PendingEntry<E> {
533    entry: E,
534    value_ref: ValueRef,
535}
536
537/// A streaming SST file writer that writes blocks to disk incrementally.
538///
539/// Instead of materializing all entries in memory and then writing all value blocks followed by all
540/// key blocks, this writer interleaves block writes as entries arrive. Medium values are written
541/// immediately, small values are accumulated into blocks, and key blocks are flushed as soon as
542/// their value references are all resolved.
543///
544/// The SST reader is block-index-addressed (not file-position-addressed), so interleaving block
545/// types is fully compatible.
546pub struct StreamingSstWriter<E: Entry> {
547    // File I/O. Wrapped in Option so close() can take ownership without a partial-move
548    // compile error (partial moves are forbidden when the type has a Drop impl).
549    file: Option<BufWriter<File>>,
550    file_path: PathBuf,
551    /// Whether close() successfully finalized the SST, so Drop must preserve the file.
552    preserve_file: bool,
553    compress_buffer: Vec<u8>,
554    block_offsets: Vec<u32>,
555    compressor: Compressor,
556
557    /// Pending key entries waiting to be flushed as key blocks.
558    ///
559    /// Entries are appended at the back and drained from the front once flushed.
560    ///
561    /// ```text
562    ///  Resolved entries              Unresolved entries
563    ///  (value block index known)     (PendingSmall references)
564    /// |------------------------------|--------------------------|
565    /// 0                     first_pending_small_index         len()
566    ///
567    ///  ^-- current_key_block tracks      ^-- these wait for
568    ///      the incomplete tail block         flush_small_value_block()
569    ///      within this region                to resolve them
570    /// ```
571    ///
572    /// [`advance_boundary_to`](Self::advance_boundary_to) scans the resolved prefix, flushes
573    /// complete key blocks from the front, and drains them. When a small value block is flushed,
574    /// all `PendingSmall` entries are resolved in-place and the boundary advances to `len()`.
575    ///
576    /// **Unbounded growth note:** If a small number of small values appear early, followed by
577    /// many medium/inline values, the queue grows because the front entries block on the
578    /// unflushed small value block while the back keeps accepting resolved entries.
579    pending_keys: VecDeque<PendingEntry<E>>,
580
581    /// Index into `pending_keys` of the first entry that has a `PendingSmall` reference for the
582    /// current (unflushed) small value block. All entries before this index are fully resolved
583    /// (their value block indices are known). Equals `pending_keys.len()` when no pending small
584    /// entries exist.
585    first_pending_small_index: usize,
586
587    /// The current small_block_id being accumulated into (debug-only consistency check).
588    #[cfg(debug_assertions)]
589    current_small_block_id: u16,
590
591    // Pending small value block buffer.
592    pending_small_value_block: Vec<u8>,
593
594    // Reusable buffer for building key blocks
595    key_buffer: Vec<u8>,
596
597    // Reusable buffer for the tail region of a key block: the key (when the search region holds
598    // hashes) and the value. Appended to `key_buffer` when the block is finished.
599    key_value_buffer: Vec<u8>,
600
601    // Collected key hashes truncated to u32 for deferred AMQF construction via sorted Builder
602    // in close(). Fingerprint size is always <32 bits, so the lower 32 bits suffice.
603    collected_fingerprints: Vec<u32>,
604
605    // Index block data: (first_hash, block_index) for each key block written
606    key_block_boundaries: Vec<(u64, u16)>,
607
608    // Metadata
609    min_hash: u64,
610    max_hash: u64,
611    entry_count: u64,
612    flags: MetaEntryFlags,
613
614    // Fullness tracking (for compaction callers)
615    total_key_size: usize,
616    total_value_size: usize,
617
618    /// Total byte size of keys in `pending_keys` (for block capacity estimation).
619    pending_key_total_size: usize,
620
621    /// State of the current incomplete key block at the tail of the resolved prefix.
622    current_key_block: KeyBlockAccumulator,
623
624    /// Set to `true` by `close()` or `cancel()` so the Drop guard can detect writers dropped
625    /// without completing their lifecycle.
626    #[cfg(debug_assertions)]
627    finished: bool,
628}
629
630impl<E: Entry> StreamingSstWriter<E> {
631    /// Creates a new streaming SST writer.
632    ///
633    /// `max_entry_count` is used to pre-allocate buffers and estimate block counts.
634    pub fn new(
635        file: &Path,
636        flags: MetaEntryFlags,
637        max_entry_count: u64,
638        compression: Compression,
639    ) -> Result<Self> {
640        let file_path = file.to_owned();
641        let file = BufWriter::new(File::create(file)?);
642        let compressor = Compressor::new(compression)?;
643
644        // Estimate number of key blocks based on max entry count.
645        // Each key block holds up to MAX_KEY_BLOCK_ENTRIES entries.
646        let estimated_key_blocks = (max_entry_count as usize)
647            .div_ceil(MAX_KEY_BLOCK_ENTRIES)
648            .max(1);
649        // Estimate value blocks assuming all entries are small values of average size.
650        // Each small value block holds ~MIN_SMALL_VALUE_BLOCK_SIZE / AVG_SMALL_VALUE_SIZE entries.
651        let entries_per_value_block = MIN_SMALL_VALUE_BLOCK_SIZE / AVG_SMALL_VALUE_SIZE;
652        let estimated_value_blocks = (max_entry_count as usize)
653            .div_ceil(entries_per_value_block)
654            .max(1);
655        let estimated_total_blocks = estimated_key_blocks + estimated_value_blocks + 1;
656
657        Ok(Self {
658            file: Some(file),
659            file_path,
660            preserve_file: false,
661            compress_buffer: Vec::with_capacity(MIN_SMALL_VALUE_BLOCK_SIZE + MAX_SMALL_VALUE_SIZE),
662            block_offsets: Vec::with_capacity(estimated_total_blocks),
663            compressor,
664            pending_keys: VecDeque::with_capacity(entries_per_value_block),
665            first_pending_small_index: 0,
666            #[cfg(debug_assertions)]
667            current_small_block_id: 0,
668            pending_small_value_block: Vec::with_capacity(
669                MIN_SMALL_VALUE_BLOCK_SIZE + MAX_SMALL_VALUE_SIZE,
670            ),
671            // `FixedKeyBlockBuilder::finish` appends the tail back into `key_buffer`, so it still
672            // holds a whole block. The tail buffer is only used by fixed-size blocks and is
673            // `reserve`d to the exact region size per block, so it starts empty.
674            key_buffer: Vec::with_capacity(MAX_KEY_BLOCK_SIZE),
675            key_value_buffer: Vec::new(),
676            collected_fingerprints: Vec::with_capacity(max_entry_count as usize),
677            key_block_boundaries: Vec::with_capacity(estimated_key_blocks),
678            min_hash: u64::MAX,
679            max_hash: 0,
680            entry_count: 0,
681            flags,
682            total_key_size: 0,
683            total_value_size: 0,
684            pending_key_total_size: 0,
685            current_key_block: KeyBlockAccumulator::new(),
686            #[cfg(debug_assertions)]
687            finished: false,
688        })
689    }
690
691    /// Returns true if the SST file has reached capacity limits.
692    ///
693    /// This is intended for compaction callers that need to split output across multiple SST files.
694    pub fn is_full(&self, max_entries: usize, max_data_size: usize) -> bool {
695        self.entry_count as usize >= max_entries
696            || self.total_key_size + self.total_value_size >= max_data_size
697            || !self.has_block_index_capacity()
698    }
699
700    /// Returns true if the SST file has room for more blocks without overflowing the `u16` block
701    /// index. Uses the exact count of blocks already written plus a conservative estimate of
702    /// blocks still needed for pending entries and the index.
703    fn has_block_index_capacity(&self) -> bool {
704        let blocks_written = self.block_offsets.len();
705        // Blocks still needed:
706        // - 1 pending small value block (if buffer is non-empty)
707        // - key blocks for pending entries (upper bound from both entry count and byte size)
708        // - 1 index block
709        let pending_small_block = usize::from(!self.pending_small_value_block.is_empty());
710        let pending_key_blocks = self
711            .pending_keys
712            .len()
713            .div_ceil(MAX_KEY_BLOCK_ENTRIES)
714            .max(self.pending_key_total_size.div_ceil(MAX_KEY_BLOCK_SIZE))
715            .max(1);
716        let index_block = 1;
717        let buffer = BLOCK_INDEX_CAPACITY_BUFFER;
718        blocks_written + pending_small_block + pending_key_blocks + index_block + buffer
719            < u16::MAX as usize
720    }
721
722    /// Adds an entry to the SST file. Entries must be added in (key-hash, key) order.
723    pub fn add(&mut self, entry: E) -> Result<()> {
724        let key_hash = entry.key_hash();
725        let key_len = entry.key_len();
726
727        // Update metadata
728        if self.entry_count == 0 {
729            self.min_hash = key_hash;
730        }
731        self.max_hash = key_hash;
732        self.entry_count += 1;
733
734        // Collect hash for deferred AMQF construction in close()
735        self.collected_fingerprints.push(key_hash as u32);
736
737        // Track key size for fullness and block capacity
738        self.total_key_size += key_len;
739        self.pending_key_total_size += key_len;
740
741        // Route value
742        let value_ref = match entry.value() {
743            EntryValue::Medium { value } => {
744                self.total_value_size += value.len();
745                let block_index = write_block_to_file(
746                    self.file.as_mut().unwrap(),
747                    &mut self.compress_buffer,
748                    &mut self.block_offsets,
749                    value,
750                    true,
751                    &mut self.compressor,
752                )
753                .context("Failed to write value block")?;
754                ValueRef::Medium { block_index }
755            }
756            EntryValue::MediumRaw {
757                uncompressed_size,
758                checksum,
759                block,
760            } => {
761                // Note: tracks compressed block size (not uncompressed) unlike EntryValue::Medium.
762                // Both are acceptable approximations of disk usage for is_full() thresholds.
763                self.total_value_size += block.len();
764                let block_index = write_raw_block_to_file(
765                    self.file.as_mut().unwrap(),
766                    &mut self.block_offsets,
767                    uncompressed_size,
768                    checksum,
769                    block,
770                )
771                .context("Failed to write compressed value block")?;
772                ValueRef::Medium { block_index }
773            }
774            EntryValue::Small { value } => {
775                self.total_value_size += value.len();
776
777                let offset = self.pending_small_value_block.len() as u32;
778                let size: u16 = value.len().try_into().unwrap();
779                self.pending_small_value_block.extend_from_slice(value);
780
781                // Track where the first PendingSmall entry is in the queue
782                if self.first_pending_small_index >= self.pending_keys.len() {
783                    self.first_pending_small_index = self.pending_keys.len();
784                }
785
786                let value_ref = ValueRef::PendingSmall {
787                    #[cfg(debug_assertions)]
788                    small_block_id: self.current_small_block_id,
789                    offset,
790                    size,
791                };
792
793                self.push_pending_key_entry(entry, value_ref);
794
795                // Eagerly flush the small block AFTER pushing the new entry. This resolves
796                // the just-pushed entry immediately via advance_boundary_to(), so key blocks
797                // can be flushed incrementally.
798                if self.pending_small_value_block.len() >= MIN_SMALL_VALUE_BLOCK_SIZE {
799                    self.flush_small_value_block()?;
800                }
801
802                return Ok(());
803            }
804            EntryValue::Inline { value } => {
805                debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
806                let mut data = [0u8; MAX_INLINE_VALUE_SIZE];
807                data[..value.len()].copy_from_slice(value);
808                ValueRef::Inline {
809                    data,
810                    len: value.len() as u8,
811                }
812            }
813            EntryValue::Large { blob } => ValueRef::Blob { blob_id: blob },
814            EntryValue::KeyDeleted => ValueRef::KeyDeleted,
815            EntryValue::KeyValueDeleted { value } => {
816                // Enforced by `WriteBatch::delete_value`, which rejects oversized values.
817                debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
818                let mut data = [0u8; MAX_INLINE_VALUE_SIZE];
819                data[..value.len()].copy_from_slice(value);
820                ValueRef::KeyValueDeleted {
821                    data,
822                    len: value.len() as u8,
823                }
824            }
825        };
826
827        self.push_pending_key_entry(entry, value_ref);
828        self.try_flush_key_blocks()
829    }
830
831    /// Abandons this writer without flushing buffered data or finalizing the SST file.
832    pub fn cancel(mut self) {
833        self.discard_partial_file();
834        #[cfg(debug_assertions)]
835        {
836            self.finished = true;
837        }
838    }
839
840    /// Closes the raw handle without flushing its buffer and best-effort removes the partial SST.
841    fn discard_partial_file(&mut self) {
842        if let Some(file) = self.file.take() {
843            // Unlike dropping BufWriter, into_parts() does not attempt to flush its buffer.
844            let (file, _) = file.into_parts();
845            drop(file);
846        }
847        // Startup recovery is the fallback if deletion itself fails (for example on Windows if
848        // another handle is still open). Never replace the error that caused cancellation.
849        let _ = fs_err::remove_file(&self.file_path);
850    }
851
852    /// Appends a new entry to the pending-keys queue.
853    fn push_pending_key_entry(&mut self, entry: E, value_ref: ValueRef) {
854        self.pending_keys
855            .push_back(PendingEntry { entry, value_ref });
856    }
857
858    /// Advances `first_pending_small_index` past the just-pushed entry if it is resolved and
859    /// sits right at the current boundary. Flushes complete key blocks incrementally.
860    ///
861    /// Must be called immediately after [`push_pending_key_entry`] with a resolved
862    /// (non-`PendingSmall`) entry.
863    fn try_flush_key_blocks(&mut self) -> Result<()> {
864        debug_assert!(!matches!(
865            self.pending_keys.back().unwrap().value_ref,
866            ValueRef::PendingSmall { .. }
867        ));
868        if self.first_pending_small_index != self.pending_keys.len() - 1 {
869            // Boundary is blocked by earlier unresolved PendingSmall entries.
870            return Ok(());
871        }
872        self.advance_boundary_to(self.pending_keys.len())
873    }
874
875    /// Advances the resolved boundary from its current position to `new_boundary`,
876    /// incrementally tracking key block sizes and flushing complete key blocks.
877    ///
878    /// All entries in `pending_keys[self.first_pending_small_index..new_boundary]`
879    /// must have resolved (non-`PendingSmall`) value references.
880    fn advance_boundary_to(&mut self, new_boundary: usize) -> Result<()> {
881        let mut last_flushed_end = 0usize;
882        // Cumulative key sizes of all entries visited so far, and the snapshot at the last
883        // flush point. The difference at the end gives the total key size of drained entries.
884        let mut cumulative_key_size = 0usize;
885        let mut flushed_key_size = 0usize;
886
887        for i in self.first_pending_small_index..new_boundary {
888            let entry = &self.pending_keys[i];
889            let key_len = entry.entry.key_len();
890            let key_hash = entry.entry.key_hash();
891            let value_type = entry.value_ref.entry_type();
892
893            if self.current_key_block.should_flush(key_len, key_hash) {
894                let block_end = last_flushed_end + self.current_key_block.entry_count;
895                let info = self.current_key_block.flush_info();
896                self.flush_key_block(last_flushed_end, block_end, info)?;
897                flushed_key_size = cumulative_key_size;
898                last_flushed_end = block_end;
899                self.current_key_block.reset();
900            }
901
902            cumulative_key_size += key_len;
903            self.current_key_block.add(key_len, key_hash, value_type);
904        }
905
906        if last_flushed_end > 0 {
907            self.pending_key_total_size -= flushed_key_size;
908            self.pending_keys.drain(..last_flushed_end);
909        }
910
911        self.first_pending_small_index = new_boundary - last_flushed_end;
912        Ok(())
913    }
914
915    /// Flushes the current pending small value block to disk and resolves all `PendingSmall`
916    /// entries in-place.
917    fn flush_small_value_block(&mut self) -> Result<()> {
918        // Early return if empty -- this simplifies trailing small value block handling in
919        // `close()` where we call this unconditionally.
920        if self.pending_small_value_block.is_empty() {
921            return Ok(());
922        }
923
924        let block_index = write_block_to_file(
925            self.file.as_mut().unwrap(),
926            &mut self.compress_buffer,
927            &mut self.block_offsets,
928            &self.pending_small_value_block,
929            true,
930            &mut self.compressor,
931        )
932        .context("Failed to write small value block")?;
933
934        // Resolve all PendingSmall entries for this block in-place.
935        // Only scan from first_pending_small_index -- entries before it are guaranteed
936        // already resolved (from previous flush calls).
937        #[cfg(debug_assertions)]
938        let flushed_id = self.current_small_block_id;
939        for i in self.first_pending_small_index..self.pending_keys.len() {
940            let entry = &mut self.pending_keys[i];
941            if let ValueRef::PendingSmall {
942                #[cfg(debug_assertions)]
943                small_block_id,
944                offset,
945                size,
946            } = entry.value_ref
947            {
948                #[cfg(debug_assertions)]
949                debug_assert_eq!(
950                    small_block_id, flushed_id,
951                    "all pending small entries must reference the small value block that was just \
952                     written"
953                );
954                entry.value_ref = ValueRef::Small {
955                    block_index,
956                    offset,
957                    size,
958                };
959            }
960        }
961
962        // All PendingSmall entries are now resolved. Advance the boundary through all of
963        // them, flushing key blocks incrementally as we go.
964        self.advance_boundary_to(self.pending_keys.len())?;
965
966        // Advance to next small block id (debug-only consistency check)
967        #[cfg(debug_assertions)]
968        {
969            self.current_small_block_id += 1;
970        }
971        self.pending_small_value_block.clear();
972
973        Ok(())
974    }
975
976    /// Flushes a single key block from `pending_keys[start..end]`.
977    ///
978    /// Potentially reorders the keys into key order if we are not storing hashes.
979    fn flush_key_block(&mut self, start: usize, end: usize, info: KeyBlockFlushInfo) -> Result<()> {
980        let entry_count = end - start;
981        let layout = choose_layout(info.max_key_len);
982        let try_compress = info.max_key_len >= MIN_KEY_SIZE_FOR_COMPRESSION;
983
984        // Read the boundary hash before reordering, which would move a different entry to `start`.
985        // The index block must keep routing by the block's lowest hash.
986        let first_hash = self.pending_keys[start].entry.key_hash();
987        // Split the borrow of `self` so the block builders can hold `&mut key_buffer` while the
988        // loops read `pending_keys`.
989        let Self {
990            key_buffer,
991            key_value_buffer,
992            pending_keys,
993            ..
994        } = self;
995        key_buffer.clear();
996        // The layout fixes the order entries must be written in, the same way it fixes their
997        // encoding, so deriving the order from the same `layout` the builders encode by keeps the
998        // two from disagreeing. `KeyOnly` blocks are searched by key and need a re-sorted copy;
999        // `HashThenKey` blocks are already in the caller's `(hash, key)` order and are yielded
1000        // straight from `pending_keys` with no allocation.
1001        let block_entries = |start: usize, end: usize| {
1002            if layout == KeyBlockLayout::HashThenKey {
1003                return Either::Left(pending_keys.range(start..end));
1004            }
1005            let mut entries: Vec<&PendingEntry<E>> = pending_keys.range(start..end).collect();
1006            // Stable sort is important to preserve relative order of tombstones
1007            match info.uniform_key_len() {
1008                Some(4) => entries.sort_by_key(|&e| be_key_u32(e.entry.key_bytes())),
1009                Some(8) => entries.sort_by_key(|&e| be_key_u64(e.entry.key_bytes())),
1010                _ => entries.sort_by_key(|&e| e.entry.key_bytes()),
1011            }
1012            Either::Right(entries.into_iter())
1013        };
1014
1015        if let KeyBlockFormat::Fixed {
1016            key_len: key_size,
1017            val_size,
1018            value_type,
1019        } = info.format
1020        {
1021            let mut builder = FixedKeyBlockBuilder::new(
1022                key_buffer,
1023                key_value_buffer,
1024                entry_count as u32,
1025                layout,
1026                key_size,
1027                val_size,
1028                value_type,
1029            );
1030            for pending in block_entries(start, end) {
1031                builder.put(&pending.entry, &pending.value_ref);
1032            }
1033            builder.finish();
1034        } else {
1035            let mut builder = KeyBlockBuilder::new(key_buffer, entry_count as u32, layout);
1036            for pending in block_entries(start, end) {
1037                builder.put(&pending.entry, &pending.value_ref);
1038            }
1039            builder.finish();
1040        }
1041
1042        let block_index = write_block_to_file(
1043            self.file.as_mut().unwrap(),
1044            &mut self.compress_buffer,
1045            &mut self.block_offsets,
1046            &self.key_buffer,
1047            try_compress,
1048            &mut self.compressor,
1049        )
1050        .context("Failed to write key block")?;
1051        self.key_block_boundaries.push((first_hash, block_index));
1052
1053        Ok(())
1054    }
1055
1056    /// Finishes writing the SST file. Flushes remaining blocks, writes the index, and returns
1057    /// metadata.
1058    pub fn close(mut self) -> Result<(StaticSortedFileBuilderMeta<'static>, File)> {
1059        #[cfg(debug_assertions)]
1060        {
1061            self.finished = true;
1062        }
1063
1064        // Flush remaining small value block (even if under MIN_SMALL_VALUE_BLOCK_SIZE).
1065        self.flush_small_value_block()?;
1066
1067        // Now all PendingSmall entries are resolved. Flush all remaining key blocks.
1068        self.flush_remaining_key_blocks()?;
1069
1070        assert!(
1071            !self.key_block_boundaries.is_empty(),
1072            "StreamingSstWriter::close() called with no entries"
1073        );
1074
1075        let mut file = self.file.take().unwrap();
1076
1077        // Write index block (never compressed). Buffer into a Vec first so we can
1078        // compute the checksum, then write via the standard block helper.
1079        let index_entry_count: u16 = (self.key_block_boundaries.len() - 1)
1080            .try_into()
1081            .expect("Index entries count overflow");
1082        let index_block_size: usize =
1083            INDEX_BLOCK_HEADER_SIZE + index_entry_count as usize * INDEX_BLOCK_ENTRY_SIZE;
1084        let mut index_buf = Vec::with_capacity(index_block_size);
1085        {
1086            let first_block = self.key_block_boundaries[0].1;
1087            let mut index_block = IndexBlockBuilder::new(&mut index_buf, first_block);
1088            for &(hash, block) in &self.key_block_boundaries[1..] {
1089                index_block.put(hash, block);
1090            }
1091        }
1092        let index_checksum = checksum_block(&index_buf);
1093        write_raw_block_to_file(
1094            &mut file,
1095            &mut self.block_offsets,
1096            0,
1097            index_checksum,
1098            &index_buf,
1099        )
1100        .context("Failed to write index block")?;
1101
1102        // Write block offset table
1103        for offset in &self.block_offsets {
1104            file.write_u32::<BE>(*offset)
1105                .context("Failed to write block offset")?;
1106        }
1107
1108        let block_count: u16 = self
1109            .block_offsets
1110            .len()
1111            .try_into()
1112            .expect("Block count overflow");
1113
1114        // Build AMQF from collected hashes using sorted Builder insertion.
1115        // Hashes are already sorted by key_hash (SST invariant), but fingerprints
1116        // (truncated hashes) may not be sorted, so we sort by `fingerprint & mask`.
1117        let actual_count = self.collected_fingerprints.len() as u64;
1118        let mut builder = qfilter::Builder::new(
1119            qfilter::Filter::new(actual_count.max(1), AMQF_FALSE_POSITIVE_RATE)
1120                .expect("Filter can't be constructed"),
1121        );
1122
1123        let fp_size = builder.fingerprint_size();
1124        assert!(fp_size < 32, "fp_size {fp_size} exceeds u32");
1125        let fp_mask = (1u32 << fp_size) - 1;
1126        // Mask in-place to fingerprint size and sort.
1127        self.collected_fingerprints
1128            .sort_unstable_by_key(|&h| h & fp_mask);
1129        for &h in &self.collected_fingerprints {
1130            builder
1131                .insert_fingerprint(false, h as u64)
1132                .expect("AMQF insert failed");
1133        }
1134        let filter = builder.into_filter();
1135
1136        // Serialize AMQF using postcard for zero-copy deserialization via FilterRef
1137        let amqf = postcard::to_allocvec(&filter).expect("AMQF serialization failed");
1138
1139        // Compute file size from block offsets rather than calling stream_position()
1140        // (which requires a flush + seek).
1141        let last_block_end = self.block_offsets.last().copied().unwrap_or_default() as u64;
1142        let offset_table_size = block_count as u64 * size_of::<u32>() as u64;
1143        let file_size = last_block_end + offset_table_size;
1144
1145        let meta = StaticSortedFileBuilderMeta {
1146            min_hash: self.min_hash,
1147            max_hash: self.max_hash,
1148            amqf: Cow::Owned(amqf),
1149            block_count,
1150            size: file_size,
1151            flags: self.flags,
1152            entries: self.entry_count,
1153        };
1154
1155        let file = file.into_inner()?;
1156        self.preserve_file = true;
1157        Ok((meta, file))
1158    }
1159
1160    /// Flushes all remaining entries as key blocks. Called from `close()` after all small value
1161    /// blocks have been flushed, so all PendingSmall entries are resolved.
1162    ///
1163    /// This loop mirrors [`advance_boundary_to`], but uses a local accumulator (since the
1164    /// `self.current_key_block` state is stale) and flushes the final incomplete block
1165    /// (unlike `advance_boundary_to`, which keeps it for more entries during streaming).
1166    fn flush_remaining_key_blocks(&mut self) -> Result<()> {
1167        if self.pending_keys.is_empty() {
1168            return Ok(());
1169        }
1170
1171        // After flush_small_value_block() in close(), no PendingSmall entries should remain.
1172        // first_pending_small_index may be non-zero (when all entries are medium/inline/etc
1173        // and advance_boundary_to was never called), but it must equal pending_keys.len(),
1174        // meaning no entries after the boundary exist.
1175        debug_assert_eq!(
1176            self.first_pending_small_index,
1177            self.pending_keys.len(),
1178            "expected no unresolved PendingSmall entries after flush_small_value_block"
1179        );
1180
1181        let total = self.pending_keys.len();
1182        let mut block_start = 0;
1183        let mut acc = KeyBlockAccumulator::new();
1184
1185        for i in 0..total {
1186            let entry = &self.pending_keys[i];
1187            let key_len = entry.entry.key_len();
1188            let key_hash = entry.entry.key_hash();
1189            let value_type = entry.value_ref.entry_type();
1190
1191            if acc.should_flush(key_len, key_hash) {
1192                self.flush_key_block(block_start, i, acc.flush_info())?;
1193                block_start = i;
1194                acc.reset();
1195            }
1196
1197            acc.add(key_len, key_hash, value_type);
1198        }
1199
1200        // Flush the final block
1201        if block_start < total {
1202            self.flush_key_block(block_start, total, acc.flush_info())?;
1203        }
1204
1205        // Free VecDeque memory. Numeric fields are not reset because close() consumes self.
1206        self.pending_keys.clear();
1207        Ok(())
1208    }
1209}
1210
1211impl<E: Entry> Drop for StreamingSstWriter<E> {
1212    fn drop(&mut self) {
1213        if !self.preserve_file {
1214            self.discard_partial_file();
1215        }
1216
1217        // Skip assertion during panic unwinding to avoid a double-panic (which would abort).
1218        #[cfg(debug_assertions)]
1219        if !std::thread::panicking() {
1220            assert!(
1221                self.finished || self.entry_count == 0,
1222                "StreamingSstWriter dropped without calling close() or cancel()"
1223            );
1224        }
1225    }
1226}
1227
1228// ---------------------------------------------------------------------------
1229// KeyBlockBuilder
1230// ---------------------------------------------------------------------------
1231
1232/// Builder for a single key block.
1233///
1234/// Entries are added via [`Self::put`], which writes key data and value references into the buffer.
1235/// The block format uses a fixed-size header table followed by variable-length entry data.
1236struct KeyBlockBuilder<'l> {
1237    current_entry: usize,
1238    header_size: usize,
1239    /// Whether entries hoist their hash into the table slot. Chosen at construction and consulted
1240    /// by [`Self::put`], so a caller cannot pair a block with the wrong entry encoding.
1241    layout: KeyBlockLayout,
1242    /// Bytes per offset table entry, which is wider when the block stores hashes.
1243    table_stride: usize,
1244    buffer: &'l mut Vec<u8>,
1245}
1246
1247/// The size of the key block header (block type + entry count).
1248const KEY_BLOCK_HEADER_SIZE: usize = 4;
1249
1250impl<'l> KeyBlockBuilder<'l> {
1251    /// Creates a new key block builder for the number of entries.
1252    fn new(buffer: &'l mut Vec<u8>, entry_count: u32, layout: KeyBlockLayout) -> Self {
1253        debug_assert!(entry_count < (1 << 24));
1254
1255        const ESTIMATED_KEY_SIZE: usize = 16;
1256        let table_stride = key_block_table_stride(layout.hash_len());
1257        buffer.reserve(entry_count as usize * (ESTIMATED_KEY_SIZE + table_stride));
1258        let block_type = layout.block_type(false);
1259        buffer.write_u8(block_type).unwrap();
1260        buffer.write_u24::<BE>(entry_count).unwrap();
1261        // Reserve the offset table; each entry's slot is filled in as it is written.
1262        buffer.resize(buffer.len() + entry_count as usize * table_stride, 0);
1263        Self {
1264            current_entry: 0,
1265            header_size: buffer.len(),
1266            layout,
1267            table_stride,
1268            buffer,
1269        }
1270    }
1271
1272    /// Writes the type and payload position into the current entry's table slot.
1273    ///
1274    /// The word sits at the end of the slot, after the hash for a `HashThenKey` block.
1275    fn write_entry_header(&mut self, entry_type: EntryType) {
1276        let pos = self.buffer.len() - self.header_size;
1277        let slot = KEY_BLOCK_HEADER_SIZE + self.current_entry * self.table_stride;
1278        let word_offset = slot + self.table_stride - KEY_BLOCK_TABLE_ENTRY_SIZE_NO_HASH;
1279        let header = (pos as u32) | ((entry_type.0 as u32) << 24);
1280        BE::write_u32(&mut self.buffer[word_offset..word_offset + 4], header);
1281    }
1282
1283    /// Writes a single entry (table slot + maybe hash? + key + value data) to the block.
1284    fn put<E: Entry>(&mut self, entry: &E, value_ref: &ValueRef) {
1285        self.write_entry_header(value_ref.entry_type());
1286        if self.layout == KeyBlockLayout::HashThenKey {
1287            let slot = KEY_BLOCK_HEADER_SIZE + self.current_entry * self.table_stride;
1288            self.buffer[slot..slot + size_of::<u64>()]
1289                .copy_from_slice(&entry.key_hash().to_be_bytes());
1290        }
1291        entry.write_key_to(self.buffer);
1292        value_ref.write_value_to(self.buffer);
1293        self.current_entry += 1;
1294    }
1295    /// Returns the key block buffer.
1296    fn finish(self) -> &'l mut Vec<u8> {
1297        self.buffer
1298    }
1299}
1300
1301// ---------------------------------------------------------------------------
1302// FixedKeyBlockBuilder
1303// ---------------------------------------------------------------------------
1304
1305/// The size of the fixed-size key block header (block type + entry count + key size + value type).
1306/// Mixed-type blocks append one more byte for the value size.
1307const FIXED_KEY_BLOCK_HEADER_SIZE: usize = 6;
1308
1309/// Builder for a fixed-size key block where all entries share the same key size and value size.
1310///
1311/// No offset table is written — entry positions are computed arithmetically from the stride. When
1312/// entries share a value size but not a value type, the header records
1313/// [`FIXED_KEY_BLOCK_MIXED_VALUE_TYPE`] and each entry carries its own type byte before its value.
1314///
1315/// Entries are written as two regions rather than interleaved, so that the bytes a lookup's binary
1316/// search probes are contiguous: the search region holds only what the lookup compares first (the
1317/// hash for `HashThenKey`, the key for `KeyOnly`), and everything else follows in the tail region,
1318/// addressed by the same entry index. [`FixedRegions`] derives that geometry for both this builder
1319/// and the reader; see [`KEY_BLOCK_TABLE_ENTRY_SIZE_WITH_HASH`] for why the compared bytes are
1320/// hoisted out of the payload.
1321struct FixedKeyBlockBuilder<'l> {
1322    /// Receives the header and then the search region.
1323    buffer: &'l mut Vec<u8>,
1324    /// Accumulates the tail region, appended to `buffer` by [`Self::finish`].
1325    tail: &'l mut Vec<u8>,
1326    /// Whether each entry writes its own type byte (set for mixed-type blocks).
1327    per_entry_type: bool,
1328    /// Which of the two regions the key goes in: the search region for `KeyOnly`, the tail for
1329    /// `HashThenKey`. Also checks that callers pair the layout with the matching `put` method.
1330    layout: KeyBlockLayout,
1331}
1332
1333impl<'l> FixedKeyBlockBuilder<'l> {
1334    fn new(
1335        buffer: &'l mut Vec<u8>,
1336        tail: &'l mut Vec<u8>,
1337        entry_count: u32,
1338        layout: KeyBlockLayout,
1339        key_size: u8,
1340        val_size: u8,
1341        value_type: Option<EntryType>,
1342    ) -> Self {
1343        let per_entry_type = value_type.is_none();
1344        // The two regions partition the entry bytes: the search region takes the bytes compared
1345        // first, the tail takes the rest. `FixedRegions` owns that split for reader and writer
1346        // alike, so the geometry is derived in one place. Its `val_size` includes the per-entry
1347        // type byte, which the block header keeps separate from the value size.
1348        let FixedRegions {
1349            search_stride,
1350            tail_stride,
1351            ..
1352        } = FixedRegions::new(
1353            entry_count as usize,
1354            layout,
1355            key_size as usize,
1356            val_size as usize + usize::from(per_entry_type),
1357        );
1358        // `finish` appends the tail back into `buffer`, so reserve room for the whole block here
1359        // and the append never reallocates.
1360        buffer.reserve(
1361            FIXED_KEY_BLOCK_HEADER_SIZE + entry_count as usize * (search_stride + tail_stride),
1362        );
1363        tail.clear();
1364        tail.reserve(entry_count as usize * tail_stride);
1365
1366        let block_type = layout.block_type(true);
1367        buffer.extend_from_slice(&[
1368            block_type,
1369            (entry_count >> 16) as u8,
1370            (entry_count >> 8) as u8,
1371            entry_count as u8,
1372            key_size,
1373            value_type.map_or(FIXED_KEY_BLOCK_MIXED_VALUE_TYPE, |ty| ty.0),
1374        ]);
1375        // Mixed-type blocks cannot derive the value size from the header's type byte, so it is
1376        // written explicitly.
1377        if per_entry_type {
1378            buffer.push(val_size);
1379        }
1380
1381        Self {
1382            buffer,
1383            tail,
1384            per_entry_type,
1385            layout,
1386        }
1387    }
1388
1389    /// Writes a single entry, splitting it between the two regions according to the block's
1390    /// layout: `HashThenKey` puts the hash in the search region and the key in the tail, `KeyOnly`
1391    /// puts the key itself in the search region. The layout decides that, not the caller.
1392    fn put<E: Entry>(&mut self, entry: &E, value_ref: &ValueRef) {
1393        match self.layout {
1394            KeyBlockLayout::HashThenKey => {
1395                self.buffer
1396                    .extend_from_slice(&entry.key_hash().to_be_bytes());
1397                entry.write_key_to(self.tail);
1398            }
1399            KeyBlockLayout::KeyOnly => entry.write_key_to(self.buffer),
1400        }
1401        self.put_tail(value_ref);
1402    }
1403
1404    /// Appends the parts of an entry that the search never reads.
1405    fn put_tail(&mut self, value_ref: &ValueRef) {
1406        if self.per_entry_type {
1407            self.tail.push(value_ref.entry_type().0);
1408        }
1409        value_ref.write_value_to(self.tail);
1410    }
1411
1412    fn finish(self) -> &'l mut Vec<u8> {
1413        self.buffer.extend_from_slice(self.tail);
1414        self.buffer
1415    }
1416}
1417
1418/// Returns the value size for a given entry type (builder-side, infallible).
1419///
1420/// This mirrors `entry_val_size` in the reader but panics on invalid types since the builder
1421/// only produces valid types.
1422fn value_type_val_size(ty: EntryType) -> usize {
1423    match ty.0 {
1424        KEY_BLOCK_ENTRY_TYPE_SMALL => SMALL_VALUE_REF_SIZE,
1425        KEY_BLOCK_ENTRY_TYPE_MEDIUM => MEDIUM_VALUE_REF_SIZE,
1426        KEY_BLOCK_ENTRY_TYPE_BLOB => BLOB_VALUE_REF_SIZE,
1427        KEY_BLOCK_ENTRY_TYPE_KEY_DELETED => KEY_DELETED_REF_SIZE,
1428        // Must precede the inline arm: both are open-ended and the tombstone range sits above it.
1429        ty if ty >= KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN => {
1430            (ty - KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN) as usize
1431        }
1432        ty if ty >= KEY_BLOCK_ENTRY_TYPE_INLINE_MIN => {
1433            (ty - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize
1434        }
1435        _ => panic!("Invalid key block entry type: {:?}", ty),
1436    }
1437}
1438
1439// ---------------------------------------------------------------------------
1440// IndexBlockBuilder
1441// ---------------------------------------------------------------------------
1442
1443/// Builder for a single index block.
1444struct IndexBlockBuilder<W: Write> {
1445    writer: W,
1446}
1447
1448/// Size of a single index block entry (u64 hash + u16 block index).
1449pub(crate) const INDEX_BLOCK_ENTRY_SIZE: usize = size_of::<u64>() + size_of::<u16>();
1450
1451/// Size of the index block header (u8 type + u16 first_block).
1452pub(crate) const INDEX_BLOCK_HEADER_SIZE: usize = size_of::<u8>() + size_of::<u16>();
1453
1454impl<W: Write> IndexBlockBuilder<W> {
1455    /// Creates a new builder for an index block with the specified number of entries and a pointer
1456    /// to the first block.
1457    fn new(mut writer: W, first_block: u16) -> Self {
1458        writer.write_u8(BLOCK_TYPE_INDEX).unwrap();
1459        writer.write_u16::<BE>(first_block).unwrap();
1460        Self { writer }
1461    }
1462
1463    /// Adds a hash boundary to the index block.
1464    fn put(&mut self, hash: u64, block: u16) {
1465        self.writer.write_u64::<BE>(hash).unwrap();
1466        self.writer.write_u16::<BE>(block).unwrap();
1467    }
1468}
1469
1470#[cfg(test)]
1471mod tests {
1472    use super::*;
1473    use crate::{
1474        key::hash_key,
1475        lookup_entry::LookupValue,
1476        static_sorted_file::{
1477            BlockCache, SstLookupResult, StaticSortedFile, StaticSortedFileMetaData,
1478        },
1479    };
1480
1481    fn make_cache() -> BlockCache {
1482        BlockCache::with(
1483            100,
1484            4 * 1024 * 1024,
1485            Default::default(),
1486            Default::default(),
1487            Default::default(),
1488        )
1489    }
1490
1491    /// A simple entry type for testing with configurable value type.
1492    struct TestEntry {
1493        key: Vec<u8>,
1494        hash: u64,
1495        value_kind: TestValueKind,
1496    }
1497
1498    enum TestValueKind {
1499        Inline(Vec<u8>),
1500        Small(Vec<u8>),
1501        Medium(Vec<u8>),
1502        /// Already-formatted block with `uncompressed_size = 0` (stored as-is).
1503        MediumRaw(Vec<u8>),
1504        Blob(u32),
1505        KeyDeleted,
1506        KeyValueDeleted(Vec<u8>),
1507    }
1508
1509    impl TestEntry {
1510        fn new(key: &[u8], value_kind: TestValueKind) -> Self {
1511            let key = key.to_vec();
1512            let hash = hash_key(&key);
1513            Self {
1514                key,
1515                hash,
1516                value_kind,
1517            }
1518        }
1519
1520        fn small(key: &[u8], value: &[u8]) -> Self {
1521            Self::new(key, TestValueKind::Small(value.to_vec()))
1522        }
1523
1524        fn inline(key: &[u8], value: &[u8]) -> Self {
1525            debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
1526            Self::new(key, TestValueKind::Inline(value.to_vec()))
1527        }
1528
1529        fn medium(key: &[u8], value: &[u8]) -> Self {
1530            Self::new(key, TestValueKind::Medium(value.to_vec()))
1531        }
1532
1533        fn blob(key: &[u8], blob_id: u32) -> Self {
1534            Self::new(key, TestValueKind::Blob(blob_id))
1535        }
1536
1537        fn deleted(key: &[u8]) -> Self {
1538            Self::new(key, TestValueKind::KeyDeleted)
1539        }
1540
1541        fn medium_raw(key: &[u8], value: &[u8]) -> Self {
1542            Self::new(key, TestValueKind::MediumRaw(value.to_vec()))
1543        }
1544
1545        fn expected_value(&self) -> Option<&[u8]> {
1546            match &self.value_kind {
1547                TestValueKind::Inline(v)
1548                | TestValueKind::Small(v)
1549                | TestValueKind::Medium(v)
1550                | TestValueKind::MediumRaw(v) => Some(v),
1551                _ => None,
1552            }
1553        }
1554    }
1555
1556    impl Entry for TestEntry {
1557        fn key_hash(&self) -> u64 {
1558            self.hash
1559        }
1560
1561        fn key_len(&self) -> usize {
1562            self.key.len()
1563        }
1564
1565        fn key_bytes(&self) -> &[u8] {
1566            &self.key
1567        }
1568
1569        fn value(&self) -> EntryValue<'_> {
1570            match &self.value_kind {
1571                TestValueKind::Inline(v) => EntryValue::Inline { value: v },
1572                TestValueKind::Small(v) => EntryValue::Small { value: v },
1573                TestValueKind::Medium(v) => EntryValue::Medium { value: v },
1574                TestValueKind::MediumRaw(v) => EntryValue::MediumRaw {
1575                    // uncompressed_size = 0 means the block is stored as-is (no compression).
1576                    uncompressed_size: 0,
1577                    checksum: checksum_block(v),
1578                    block: v,
1579                },
1580                TestValueKind::Blob(id) => EntryValue::Large { blob: *id },
1581                TestValueKind::KeyDeleted => EntryValue::KeyDeleted,
1582                TestValueKind::KeyValueDeleted(v) => EntryValue::KeyValueDeleted { value: v },
1583            }
1584        }
1585    }
1586
1587    /// Sort entries by (hash, key) (required by SST writer).
1588    fn sort_entries(entries: &mut [TestEntry]) {
1589        entries.sort_by(|a, b| a.hash.cmp(&b.hash).then_with(|| a.key.cmp(&b.key)));
1590    }
1591
1592    /// Open an SST file for lookup given a path and metadata.
1593    fn open_sst(
1594        dir: &Path,
1595        seq: u32,
1596        meta: &StaticSortedFileBuilderMeta<'_>,
1597    ) -> Result<StaticSortedFile> {
1598        StaticSortedFile::open(
1599            dir,
1600            StaticSortedFileMetaData {
1601                sequence_number: seq,
1602                block_count: meta.block_count,
1603            },
1604            Compression::Lz4,
1605            crate::mmap_access_mode(),
1606        )
1607    }
1608
1609    /// Helper: write entries via StreamingSstWriter, return meta.
1610    fn write_sst(
1611        dir: &Path,
1612        seq: u32,
1613        entries: &[TestEntry],
1614        flags: MetaEntryFlags,
1615    ) -> Result<StaticSortedFileBuilderMeta<'static>> {
1616        let sst_path = dir.join(format!("{seq:08}.sst"));
1617        let mut writer =
1618            StreamingSstWriter::new(&sst_path, flags, entries.len() as u64, Compression::Lz4)?;
1619        for entry in entries {
1620            if let Err(err) = writer.add(entry) {
1621                writer.cancel();
1622                return Err(err);
1623            }
1624        }
1625        let (meta, _file) = writer.close()?;
1626        Ok(meta)
1627    }
1628
1629    /// Lookup a key in an SST file and assert it matches the expected value kind.
1630    fn assert_lookup(
1631        sst: &StaticSortedFile,
1632        entry: &TestEntry,
1633        kc: &BlockCache,
1634        vc: &BlockCache,
1635    ) -> Result<()> {
1636        let result = sst.lookup::<_, false>(entry.hash, &entry.key, kc, vc)?;
1637        match (&entry.value_kind, result) {
1638            (_, SstLookupResult::Found(values))
1639                if values.len() == 1 && matches!(values[0], LookupValue::Slice { .. }) =>
1640            {
1641                let LookupValue::Slice { value } = &values[0] else {
1642                    unreachable!()
1643                };
1644                let expected = entry
1645                    .expected_value()
1646                    .expect("Got Slice but entry has no value");
1647                assert_eq!(
1648                    value.as_ref(),
1649                    expected,
1650                    "value mismatch for key {:?}",
1651                    std::str::from_utf8(&entry.key)
1652                );
1653            }
1654            (TestValueKind::Blob(expected_id), SstLookupResult::Found(values))
1655                if values.len() == 1 && matches!(values[0], LookupValue::Blob { .. }) =>
1656            {
1657                let LookupValue::Blob { sequence_number } = &values[0] else {
1658                    unreachable!()
1659                };
1660                assert_eq!(*sequence_number, *expected_id);
1661            }
1662            (TestValueKind::KeyDeleted, SstLookupResult::Found(values))
1663                if values.len() == 1 && matches!(values[0], LookupValue::KeyDeleted) => {}
1664            (TestValueKind::KeyValueDeleted(expected), SstLookupResult::Found(values))
1665                if values.len() == 1
1666                    && matches!(values[0], LookupValue::KeyValueDeleted { .. }) =>
1667            {
1668                let LookupValue::KeyValueDeleted { value } = &values[0] else {
1669                    unreachable!()
1670                };
1671                assert_eq!(
1672                    value.as_ref(),
1673                    expected.as_slice(),
1674                    "tombstone value mismatch for key {:?}",
1675                    std::str::from_utf8(&entry.key)
1676                );
1677            }
1678            _ => {
1679                panic!(
1680                    "Unexpected lookup result for key {:?}",
1681                    std::str::from_utf8(&entry.key)
1682                );
1683            }
1684        }
1685        Ok(())
1686    }
1687
1688    #[test]
1689    fn single_inline_entry() -> Result<()> {
1690        let dir = tempfile::tempdir()?;
1691        let mut entries = vec![TestEntry::inline(b"key1", b"val1")];
1692        sort_entries(&mut entries);
1693
1694        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1695        assert_eq!(meta.entries, 1);
1696
1697        let sst = open_sst(dir.path(), 1, &meta)?;
1698        let kc = make_cache();
1699        let vc = make_cache();
1700        assert_lookup(&sst, &entries[0], &kc, &vc)?;
1701        Ok(())
1702    }
1703
1704    #[test]
1705    fn single_small_entry() -> Result<()> {
1706        let dir = tempfile::tempdir()?;
1707        let value = vec![0xAB; 100]; // > MAX_INLINE_VALUE_SIZE, <= MAX_SMALL_VALUE_SIZE
1708        let mut entries = vec![TestEntry::small(b"skey", &value)];
1709        sort_entries(&mut entries);
1710
1711        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1712        assert_eq!(meta.entries, 1);
1713
1714        let sst = open_sst(dir.path(), 1, &meta)?;
1715        let kc = make_cache();
1716        let vc = make_cache();
1717        assert_lookup(&sst, &entries[0], &kc, &vc)?;
1718        Ok(())
1719    }
1720
1721    #[test]
1722    fn single_medium_entry() -> Result<()> {
1723        let dir = tempfile::tempdir()?;
1724        let value = vec![0xCD; 8192]; // > MAX_SMALL_VALUE_SIZE
1725        let mut entries = vec![TestEntry::medium(b"mkey", &value)];
1726        sort_entries(&mut entries);
1727
1728        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1729        assert_eq!(meta.entries, 1);
1730
1731        let sst = open_sst(dir.path(), 1, &meta)?;
1732        let kc = make_cache();
1733        let vc = make_cache();
1734        assert_lookup(&sst, &entries[0], &kc, &vc)?;
1735        Ok(())
1736    }
1737
1738    #[test]
1739    fn single_blob_entry() -> Result<()> {
1740        let dir = tempfile::tempdir()?;
1741        let mut entries = vec![TestEntry::blob(b"bkey", 42)];
1742        sort_entries(&mut entries);
1743
1744        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1745        assert_eq!(meta.entries, 1);
1746
1747        let sst = open_sst(dir.path(), 1, &meta)?;
1748        let kc = make_cache();
1749        let vc = make_cache();
1750        assert_lookup(&sst, &entries[0], &kc, &vc)?;
1751        Ok(())
1752    }
1753
1754    #[test]
1755    fn single_deleted_entry() -> Result<()> {
1756        let dir = tempfile::tempdir()?;
1757        let mut entries = vec![TestEntry::deleted(b"dkey")];
1758        sort_entries(&mut entries);
1759
1760        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1761        assert_eq!(meta.entries, 1);
1762
1763        let sst = open_sst(dir.path(), 1, &meta)?;
1764        let kc = make_cache();
1765        let vc = make_cache();
1766        assert_lookup(&sst, &entries[0], &kc, &vc)?;
1767        Ok(())
1768    }
1769
1770    #[test]
1771    fn many_small_values() -> Result<()> {
1772        let dir = tempfile::tempdir()?;
1773        // Create enough small entries to trigger multiple small value block flushes.
1774        // MIN_SMALL_VALUE_BLOCK_SIZE = 8KB, each value is 200 bytes -> ~40 entries per block.
1775        let count = 200;
1776        let mut entries: Vec<TestEntry> = (0..count)
1777            .map(|i| {
1778                let key = format!("key-{i:04}");
1779                let value = vec![(i & 0xFF) as u8; 200];
1780                TestEntry::small(key.as_bytes(), &value)
1781            })
1782            .collect();
1783        sort_entries(&mut entries);
1784
1785        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1786        assert_eq!(meta.entries, count as u64);
1787
1788        let sst = open_sst(dir.path(), 1, &meta)?;
1789        let kc = make_cache();
1790        let vc = make_cache();
1791
1792        for entry in &entries {
1793            assert_lookup(&sst, entry, &kc, &vc)?;
1794        }
1795        Ok(())
1796    }
1797
1798    #[test]
1799    fn many_medium_values() -> Result<()> {
1800        let dir = tempfile::tempdir()?;
1801        let count = 50;
1802        let mut entries: Vec<TestEntry> = (0..count)
1803            .map(|i| {
1804                let key = format!("mkey-{i:04}");
1805                let value = vec![(i & 0xFF) as u8; 8192];
1806                TestEntry::medium(key.as_bytes(), &value)
1807            })
1808            .collect();
1809        sort_entries(&mut entries);
1810
1811        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1812        assert_eq!(meta.entries, count as u64);
1813
1814        let sst = open_sst(dir.path(), 1, &meta)?;
1815        let kc = make_cache();
1816        let vc = make_cache();
1817
1818        for entry in &entries {
1819            assert_lookup(&sst, entry, &kc, &vc)?;
1820        }
1821        Ok(())
1822    }
1823
1824    #[test]
1825    fn mixed_value_types() -> Result<()> {
1826        let dir = tempfile::tempdir()?;
1827        let mut entries = vec![
1828            TestEntry::inline(b"a-inline", b"tiny"),
1829            TestEntry::small(b"b-small", &[0x11; 200]),
1830            TestEntry::medium(b"c-medium", &[0x22; 8192]),
1831            TestEntry::blob(b"d-blob", 99),
1832            TestEntry::deleted(b"e-deleted"),
1833            TestEntry::small(b"f-small2", &[0x33; 300]),
1834            TestEntry::inline(b"g-inline2", b"mini"),
1835            TestEntry::medium(b"h-medium2", &[0x44; 16384]),
1836        ];
1837        sort_entries(&mut entries);
1838
1839        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1840        assert_eq!(meta.entries, 8);
1841
1842        let sst = open_sst(dir.path(), 1, &meta)?;
1843        let kc = make_cache();
1844        let vc = make_cache();
1845
1846        for entry in &entries {
1847            assert_lookup(&sst, entry, &kc, &vc)?;
1848        }
1849        Ok(())
1850    }
1851
1852    #[test]
1853    #[cfg(debug_assertions)]
1854    fn cancel_after_failed_add_does_not_trigger_drop_assertion() {
1855        let dir = tempfile::tempdir().unwrap();
1856        let sst_path = dir.path().join("test.sst");
1857        let mut writer =
1858            StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 1, Compression::Lz4)
1859                .unwrap();
1860
1861        // Replace the output with an unbuffered, read-only handle to force a deterministic write
1862        // error without relying on the filesystem being full.
1863        drop(writer.file.take());
1864        writer.file = Some(BufWriter::with_capacity(0, File::open(&sst_path).unwrap()));
1865        let error = writer
1866            .add(TestEntry::medium(b"key", &[0; 8192]))
1867            .unwrap_err();
1868        assert!(format!("{error:#}").contains("Failed to write value block"));
1869
1870        let cancel_result =
1871            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| writer.cancel()));
1872        assert!(
1873            cancel_result.is_ok(),
1874            "cancelling a writer after an add error must not trigger the lifecycle assertion"
1875        );
1876        assert!(!sst_path.exists(), "cancel should remove the partial SST");
1877    }
1878
1879    #[test]
1880    fn failed_close_removes_partial_file() {
1881        let dir = tempfile::tempdir().unwrap();
1882        let sst_path = dir.path().join("test.sst");
1883        let mut writer =
1884            StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 1, Compression::Lz4)
1885                .unwrap();
1886        writer.add(TestEntry::inline(b"key", b"value")).unwrap();
1887
1888        // Force close() to fail while flushing the pending key block.
1889        drop(writer.file.take());
1890        writer.file = Some(BufWriter::with_capacity(0, File::open(&sst_path).unwrap()));
1891        let error = writer.close().unwrap_err();
1892
1893        assert!(format!("{error:#}").contains("Failed to write key block"));
1894        assert!(
1895            !sst_path.exists(),
1896            "a failed close should remove the partial SST"
1897        );
1898    }
1899
1900    #[test]
1901    fn is_full_entry_count_limit() {
1902        let dir = tempfile::tempdir().unwrap();
1903        let sst_path = dir.path().join("test.sst");
1904        let mut writer =
1905            StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 100, Compression::Lz4)
1906                .unwrap();
1907
1908        let max_entries = 50;
1909        for i in 0..max_entries {
1910            let key = format!("k{i:06}");
1911            let entry = TestEntry::inline(key.as_bytes(), &[0; 4]);
1912            writer.add(entry).unwrap();
1913        }
1914
1915        assert_eq!(writer.entry_count, max_entries as u64);
1916        assert!(
1917            writer.is_full(max_entries, usize::MAX),
1918            "Should be full when entry count reaches max_entries"
1919        );
1920        assert!(
1921            !writer.is_full(max_entries + 1, usize::MAX),
1922            "Should not be full when limit is higher"
1923        );
1924        writer.close().unwrap();
1925    }
1926
1927    #[test]
1928    fn is_full_data_size_limit() {
1929        let dir = tempfile::tempdir().unwrap();
1930        let sst_path = dir.path().join("test.sst");
1931        let mut writer =
1932            StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 100, Compression::Lz4)
1933                .unwrap();
1934
1935        let value = vec![0u8; 1000];
1936        for i in 0..10 {
1937            let key = format!("k{i:06}");
1938            let entry = TestEntry::small(key.as_bytes(), &value);
1939            writer.add(entry).unwrap();
1940        }
1941
1942        let total = writer.total_key_size + writer.total_value_size;
1943        assert!(total > 10_000, "total data should exceed 10KB");
1944        assert!(writer.is_full(usize::MAX, total - 1));
1945        assert!(!writer.is_full(usize::MAX, total + 1));
1946        writer.close().unwrap();
1947    }
1948
1949    #[test]
1950    fn write_static_stored_file_matches_streaming() -> Result<()> {
1951        let dir = tempfile::tempdir()?;
1952
1953        let mut entries: Vec<TestEntry> = (0..100)
1954            .map(|i| {
1955                let key = format!("rkey-{i:04}");
1956                if i % 3 == 0 {
1957                    TestEntry::inline(key.as_bytes(), &[(i & 0xFF) as u8; 4])
1958                } else if i % 3 == 1 {
1959                    TestEntry::small(key.as_bytes(), &[(i & 0xFF) as u8; 200])
1960                } else {
1961                    TestEntry::medium(key.as_bytes(), &[(i & 0xFF) as u8; 8192])
1962                }
1963            })
1964            .collect();
1965        sort_entries(&mut entries);
1966
1967        // Write via convenience function
1968        let batch_path = dir.path().join("00000001.sst");
1969        let (meta1, _) = write_static_stored_file(
1970            &entries,
1971            &batch_path,
1972            MetaEntryFlags::default(),
1973            Compression::Lz4,
1974        )?;
1975
1976        // Write via streaming API
1977        let streaming_path = dir.path().join("00000002.sst");
1978        let mut writer = StreamingSstWriter::new(
1979            &streaming_path,
1980            MetaEntryFlags::default(),
1981            entries.len() as u64,
1982            Compression::Lz4,
1983        )?;
1984        for entry in &entries {
1985            if let Err(err) = writer.add(entry) {
1986                writer.cancel();
1987                return Err(err);
1988            }
1989        }
1990        let (meta2, _) = writer.close()?;
1991
1992        // Metadata should match
1993        assert_eq!(meta1.entries, meta2.entries);
1994        assert_eq!(meta1.min_hash, meta2.min_hash);
1995        assert_eq!(meta1.max_hash, meta2.max_hash);
1996        assert_eq!(meta1.block_count, meta2.block_count);
1997
1998        // Both files should produce the same lookup results
1999        let sst1 = StaticSortedFile::open(
2000            dir.path(),
2001            StaticSortedFileMetaData {
2002                sequence_number: 1,
2003                block_count: meta1.block_count,
2004            },
2005            Compression::Lz4,
2006            crate::mmap_access_mode(),
2007        )?;
2008        let sst2 = StaticSortedFile::open(
2009            dir.path(),
2010            StaticSortedFileMetaData {
2011                sequence_number: 2,
2012                block_count: meta2.block_count,
2013            },
2014            Compression::Lz4,
2015            crate::mmap_access_mode(),
2016        )?;
2017        let kc = make_cache();
2018        let vc = make_cache();
2019
2020        for entry in &entries {
2021            let r1 = sst1.lookup::<_, false>(entry.hash, &entry.key, &kc, &vc)?;
2022            let r2 = sst2.lookup::<_, false>(entry.hash, &entry.key, &kc, &vc)?;
2023            match (&r1, &r2) {
2024                (SstLookupResult::Found(v1), SstLookupResult::Found(v2))
2025                    if v1.len() == 1 && v2.len() == 1 =>
2026                {
2027                    match (&v1[0], &v2[0]) {
2028                        (
2029                            LookupValue::Slice { value: val1 },
2030                            LookupValue::Slice { value: val2 },
2031                        ) => {
2032                            assert_eq!(
2033                                val1.as_ref(),
2034                                val2.as_ref(),
2035                                "Value mismatch for key {:?}",
2036                                std::str::from_utf8(&entry.key)
2037                            );
2038                        }
2039                        (LookupValue::KeyDeleted, LookupValue::KeyDeleted) => {}
2040                        (
2041                            LookupValue::Blob {
2042                                sequence_number: s1,
2043                            },
2044                            LookupValue::Blob {
2045                                sequence_number: s2,
2046                            },
2047                        ) => {
2048                            assert_eq!(s1, s2);
2049                        }
2050                        _ => panic!(
2051                            "Mismatched results for key {:?}",
2052                            std::str::from_utf8(&entry.key)
2053                        ),
2054                    }
2055                }
2056                _ => panic!(
2057                    "Mismatched results for key {:?}",
2058                    std::str::from_utf8(&entry.key)
2059                ),
2060            }
2061        }
2062        Ok(())
2063    }
2064
2065    #[test]
2066    #[should_panic(expected = "StreamingSstWriter::close() called with no entries")]
2067    fn close_empty_writer_panics() {
2068        let dir = tempfile::tempdir().unwrap();
2069        let sst_path = dir.path().join("empty.sst");
2070        let writer = StreamingSstWriter::<TestEntry>::new(
2071            &sst_path,
2072            MetaEntryFlags::default(),
2073            0,
2074            Compression::Lz4,
2075        )
2076        .unwrap();
2077        writer.close().unwrap();
2078    }
2079
2080    #[test]
2081    fn key_block_boundary_at_max_entries() -> Result<()> {
2082        let dir = tempfile::tempdir()?;
2083        let count = MAX_KEY_BLOCK_ENTRIES + 1;
2084        let mut entries: Vec<TestEntry> = (0..count)
2085            .map(|i| {
2086                let key = format!("boundary-{i:06}");
2087                TestEntry::inline(key.as_bytes(), &[0u8; 4])
2088            })
2089            .collect();
2090        sort_entries(&mut entries);
2091
2092        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2093        assert_eq!(meta.entries, count as u64);
2094        // count > MAX_KEY_BLOCK_ENTRIES so we need at least 2 key blocks plus 1 index block
2095        assert!(
2096            meta.block_count >= 3,
2097            "expected at least 2 key blocks + 1 index block"
2098        );
2099
2100        let sst = open_sst(dir.path(), 1, &meta)?;
2101        let kc = make_cache();
2102        let vc = make_cache();
2103        for entry in &entries {
2104            assert_lookup(&sst, entry, &kc, &vc)?;
2105        }
2106        Ok(())
2107    }
2108
2109    /// Reads the first key block of an SST, returning its raw (uncompressed) bytes.
2110    ///
2111    /// Block 0 is always a key block; the block offset table sits at the end of the file.
2112    fn read_first_block(dir: &Path, seq: u32, block_count: u16) -> Result<Vec<u8>> {
2113        let data = fs_err::read(dir.join(format!("{seq:08}.sst")))?;
2114        let offsets_start = data.len() - block_count as usize * size_of::<u32>();
2115        let end = BE::read_u32(&data[offsets_start..]) as usize;
2116        let raw = &data[..end];
2117        // Each block is prefixed by BLOCK_HEADER_SIZE bytes: 4B uncompressed size + 4B checksum.
2118        // An uncompressed size of 0 means the block is stored as-is.
2119        let uncompressed_size = BE::read_u32(raw) as usize;
2120        let body = &raw[BLOCK_HEADER_SIZE..];
2121        Ok(if uncompressed_size == 0 {
2122            body.to_vec()
2123        } else {
2124            let mut out = vec![0u8; uncompressed_size];
2125            lz4_flex::block::decompress_into(body, &mut out)?;
2126            out
2127        })
2128    }
2129
2130    /// A tombstone and a value of the same size keep the block in fixed layout.
2131    ///
2132    /// This is what makes tombstones cheap for uniform-key families like the task cache: without
2133    /// the mixed-type layout, one tombstone would demote its whole block to the variable format
2134    /// and add a 4-byte offset table entry for every entry in it.
2135    #[test]
2136    fn fixed_layout_survives_mixed_value_types_of_equal_size() -> Result<()> {
2137        let dir = tempfile::tempdir()?;
2138
2139        // Uniform 8-byte keys, uniform 4-byte values, but two different entry types.
2140        let mut entries: Vec<TestEntry> = (0..64u64)
2141            .map(|i| {
2142                let key = format!("k-{i:06}");
2143                if i % 4 == 0 {
2144                    TestEntry::new(
2145                        key.as_bytes(),
2146                        TestValueKind::KeyValueDeleted(vec![0xAAu8; 4]),
2147                    )
2148                } else {
2149                    TestEntry::inline(key.as_bytes(), &[0xBBu8; 4])
2150                }
2151            })
2152            .collect();
2153        sort_entries(&mut entries);
2154
2155        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2156        let block = read_first_block(dir.path(), 1, meta.block_count)?;
2157
2158        assert_eq!(
2159            KeyBlockLayout::from_block_type(block[0]).map(|(_, fixed)| fixed),
2160            Some(true),
2161            "mixed value types of equal size should stay in fixed layout, got block type {}",
2162            block[0]
2163        );
2164        assert_eq!(
2165            block[5], FIXED_KEY_BLOCK_MIXED_VALUE_TYPE,
2166            "block should be marked mixed-type"
2167        );
2168        assert_eq!(block[6], 4, "value size should be recorded in the header");
2169
2170        // The layout is only useful if it still reads back correctly.
2171        let sst = open_sst(dir.path(), 1, &meta)?;
2172        let kc = make_cache();
2173        let vc = make_cache();
2174        for entry in &entries {
2175            assert_lookup(&sst, entry, &kc, &vc)?;
2176        }
2177        Ok(())
2178    }
2179
2180    /// Differing value *sizes* cannot share a stride, so the block must fall back to variable
2181    /// layout rather than silently misreading entries.
2182    #[test]
2183    fn mixed_value_sizes_fall_back_to_variable_layout() -> Result<()> {
2184        let dir = tempfile::tempdir()?;
2185
2186        let mut entries: Vec<TestEntry> = (0..64u64)
2187            .map(|i| {
2188                let key = format!("k-{i:06}");
2189                let len = if i % 4 == 0 { 2 } else { 4 };
2190                TestEntry::inline(key.as_bytes(), &vec![0xCCu8; len])
2191            })
2192            .collect();
2193        sort_entries(&mut entries);
2194
2195        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2196        let block = read_first_block(dir.path(), 1, meta.block_count)?;
2197        assert_eq!(
2198            KeyBlockLayout::from_block_type(block[0]).map(|(_, fixed)| fixed),
2199            Some(false),
2200            "differing value sizes should use variable layout, got block type {}",
2201            block[0]
2202        );
2203
2204        let sst = open_sst(dir.path(), 1, &meta)?;
2205        let kc = make_cache();
2206        let vc = make_cache();
2207        for entry in &entries {
2208            assert_lookup(&sst, entry, &kc, &vc)?;
2209        }
2210        Ok(())
2211    }
2212
2213    #[test]
2214    fn single_medium_raw_entry() -> Result<()> {
2215        let dir = tempfile::tempdir()?;
2216        let value = vec![0xBE; 8192];
2217        let mut entries = vec![TestEntry::medium_raw(b"rkey", &value)];
2218        sort_entries(&mut entries);
2219
2220        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2221        assert_eq!(meta.entries, 1);
2222
2223        let sst = open_sst(dir.path(), 1, &meta)?;
2224        let kc = make_cache();
2225        let vc = make_cache();
2226        assert_lookup(&sst, &entries[0], &kc, &vc)?;
2227        Ok(())
2228    }
2229
2230    /// Flip a single byte in an SST file at the given position.
2231    fn corrupt_sst_byte(dir: &Path, seq: u32, pos: u64) {
2232        use std::io::{Seek, SeekFrom, Write as _};
2233
2234        let sst_path = dir.join(format!("{seq:08}.sst"));
2235        let file_bytes = std::fs::read(&sst_path).unwrap();
2236        let original = file_bytes[pos as usize];
2237        let mut file = std::fs::OpenOptions::new()
2238            .write(true)
2239            .open(&sst_path)
2240            .unwrap();
2241        file.seek(SeekFrom::Start(pos)).unwrap();
2242        file.write_all(&[original ^ 0xFF]).unwrap();
2243        file.sync_all().unwrap();
2244    }
2245
2246    /// Assert that looking up the first entry in a corrupted SST returns a corruption error.
2247    fn assert_corruption_detected(
2248        dir: &Path,
2249        seq: u32,
2250        meta: &StaticSortedFileBuilderMeta<'_>,
2251        entries: &[TestEntry],
2252    ) {
2253        let sst = open_sst(dir, seq, meta).unwrap();
2254        let kc = make_cache();
2255        let vc = make_cache();
2256        match sst.lookup::<_, false>(entries[0].hash, &entries[0].key, &kc, &vc) {
2257            Err(err) => {
2258                let msg = format!("{err}");
2259                assert!(
2260                    msg.contains("corruption"),
2261                    "Expected corruption error, got: {msg}"
2262                );
2263            }
2264            Ok(_) => panic!("Expected checksum error, but lookup succeeded"),
2265        }
2266    }
2267
2268    #[test]
2269    fn checksum_detects_corrupted_compressed_block() {
2270        let dir = tempfile::tempdir().unwrap();
2271        // Medium value is large enough to get its own value block, which will be compressed
2272        let value = vec![0xCD; 8192];
2273        let entries = vec![TestEntry::medium(b"mkey", &value)];
2274
2275        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default()).unwrap();
2276
2277        // Corrupt the stored checksum of the first block (bytes 4..8).
2278        // This guarantees a mismatch regardless of whether LZ4 decompression succeeds.
2279        corrupt_sst_byte(dir.path(), 1, 4);
2280        assert_corruption_detected(dir.path(), 1, &meta, &entries);
2281    }
2282
2283    #[test]
2284    fn checksum_detects_corrupted_uncompressed_block() {
2285        let dir = tempfile::tempdir().unwrap();
2286        // Single inline entry - the key block will be small and likely stored uncompressed
2287        let entries = vec![TestEntry::inline(b"key1", b"val1")];
2288
2289        let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default()).unwrap();
2290
2291        // Corrupt a byte in the first block's data (after the 8-byte header)
2292        corrupt_sst_byte(dir.path(), 1, BLOCK_HEADER_SIZE as u64 + 1);
2293        assert_corruption_detected(dir.path(), 1, &meta, &entries);
2294    }
2295
2296    #[test]
2297    fn be_key_order_matches_byte_order() {
2298        let keys4: Vec<[u8; 4]> = vec![
2299            [0, 0, 0, 0],
2300            [0, 0, 0, 1],
2301            [0, 0, 1, 0],
2302            [0x7f, 0xff, 0xff, 0xff],
2303            [0x80, 0, 0, 0],
2304            [0xff, 0xfe, 0, 0],
2305            [0xff, 0xff, 0xff, 0xff],
2306        ];
2307        for a in &keys4 {
2308            for b in &keys4 {
2309                assert_eq!(
2310                    be_key_u32(a).cmp(&be_key_u32(b)),
2311                    a[..].cmp(&b[..]),
2312                    "u32 order disagrees with byte order for {a:?} vs {b:?}"
2313                );
2314            }
2315        }
2316        let keys8: Vec<[u8; 8]> = vec![
2317            [0; 8],
2318            [0, 0, 0, 0, 0, 0, 0, 1],
2319            [0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
2320            [0x80, 0, 0, 0, 0, 0, 0, 0],
2321            [0xff; 8],
2322        ];
2323        for a in &keys8 {
2324            for b in &keys8 {
2325                assert_eq!(
2326                    be_key_u64(a).cmp(&be_key_u64(b)),
2327                    a[..].cmp(&b[..]),
2328                    "u64 order disagrees with byte order for {a:?} vs {b:?}"
2329                );
2330            }
2331        }
2332    }
2333
2334    /// `uniform_key_len` must only report a length when the block's keys really are all that long,
2335    /// since the specialized sorts are unsound otherwise.
2336    #[test]
2337    fn uniform_key_len_requires_equal_lengths() {
2338        let mut acc = KeyBlockAccumulator::new();
2339        assert_eq!(acc.flush_info().uniform_key_len(), None, "empty block");
2340
2341        let ty = EntryType(KEY_BLOCK_ENTRY_TYPE_INLINE_MIN);
2342        acc.add(8, 1, ty);
2343        acc.add(8, 2, ty);
2344        assert_eq!(acc.flush_info().uniform_key_len(), Some(8));
2345
2346        acc.add(4, 3, ty);
2347        assert_eq!(
2348            acc.flush_info().uniform_key_len(),
2349            None,
2350            "mixed lengths must not report a uniform length"
2351        );
2352
2353        acc.reset();
2354        acc.add(4, 4, ty);
2355        assert_eq!(
2356            acc.flush_info().uniform_key_len(),
2357            Some(4),
2358            "reset clears min"
2359        );
2360    }
2361}