Skip to main content

turbo_persistence/
static_sorted_file.rs

1use std::{
2    cmp::Ordering,
3    hash::BuildHasherDefault,
4    path::Path,
5    rc::Rc,
6    sync::{
7        Arc,
8        atomic::{AtomicU64, Ordering as AtomicOrdering},
9    },
10};
11
12use anyhow::{Context, Result, bail, ensure};
13use fs_err::File;
14use memmap2::Mmap;
15use quick_cache::{Lifecycle, sync::GuardResult};
16use rustc_hash::FxHasher;
17use smallvec::SmallVec;
18
19use crate::{
20    QueryKey,
21    arc_bytes::ArcBytes,
22    be,
23    compression::checksum_block,
24    constants::MAX_INLINE_VALUE_SIZE,
25    lookup_entry::{IterValue, LookupEntry, LookupValue},
26    mmap_helper::advise_mmap_for_persistence,
27    rc_bytes::RcBytes,
28    shared_bytes::SharedBytes,
29    static_sorted_file_builder::{BLOCK_HEADER_SIZE, INDEX_BLOCK_ENTRY_SIZE},
30};
31
32/// The block header for an index block.
33pub const BLOCK_TYPE_INDEX: u8 = 0;
34/// The block header for a key block with 8-byte hash per entry.
35pub const BLOCK_TYPE_KEY_WITH_HASH: u8 = 1;
36/// The block header for a key block without hash.
37pub const BLOCK_TYPE_KEY_NO_HASH: u8 = 2;
38/// The block header for a fixed-size key block with 8-byte hash per entry.
39pub const BLOCK_TYPE_FIXED_KEY_WITH_HASH: u8 = 3;
40/// The block header for a fixed-size key block without hash.
41pub const BLOCK_TYPE_FIXED_KEY_NO_HASH: u8 = 4;
42
43/// The tag for a small-sized value.
44pub const KEY_BLOCK_ENTRY_TYPE_SMALL: u8 = 0;
45/// The tag for the blob value.
46pub const KEY_BLOCK_ENTRY_TYPE_BLOB: u8 = 1;
47/// The tag for the deleted value.
48pub const KEY_BLOCK_ENTRY_TYPE_DELETED: u8 = 2;
49/// The tag for a medium-sized value.
50pub const KEY_BLOCK_ENTRY_TYPE_MEDIUM: u8 = 3;
51/// The minimum tag for inline values. The actual size is (tag - INLINE_MIN).
52pub const KEY_BLOCK_ENTRY_TYPE_INLINE_MIN: u8 = 8;
53
54/// Encoded size of a small value reference: 2B block index + 2B size + 4B offset.
55pub(crate) const SMALL_VALUE_REF_SIZE: usize = 8;
56/// Encoded size of a medium value reference: 2B block index.
57pub(crate) const MEDIUM_VALUE_REF_SIZE: usize = 2;
58/// Encoded size of a blob value reference: 4B blob id.
59pub(crate) const BLOB_VALUE_REF_SIZE: usize = 4;
60/// Encoded size of a deleted (tombstone) value reference.
61pub(crate) const DELETED_VALUE_REF_SIZE: usize = 0;
62
63// Static assertion: MAX_INLINE_VALUE_SIZE must fit in the key type encoding.
64// Key types 8-255 encode inline values of size 0-247, so max is 255 - 8 = 247.
65const _: () = assert!(
66    MAX_INLINE_VALUE_SIZE <= (u8::MAX - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize,
67    "MAX_INLINE_VALUE_SIZE exceeds what can be encoded in key type byte"
68);
69
70/// The result of a lookup operation.
71pub enum SstLookupResult {
72    /// One or more values were found.
73    Found(SmallVec<[LookupValue; 1]>),
74    /// The key was not found.
75    NotFound,
76}
77
78impl From<LookupValue> for SstLookupResult {
79    fn from(value: LookupValue) -> Self {
80        SstLookupResult::Found(smallvec::smallvec![value])
81    }
82}
83
84#[derive(Clone, Default)]
85pub struct BlockWeighter;
86
87impl quick_cache::Weighter<(u32, u16), ArcBytes> for BlockWeighter {
88    fn weight(&self, _key: &(u32, u16), val: &ArcBytes) -> u64 {
89        if val.is_mmap_backed() {
90            // Mmap-backed blocks bypass the cache (served directly from mmap),
91            // so this branch should never be reached.
92            debug_assert!(
93                !val.is_mmap_backed(),
94                "mmap-backed block should not be inserted into BlockCache"
95            );
96            64
97        } else {
98            val.len() as u64 + 8
99        }
100    }
101}
102
103/// Lifecycle hooks for the block cache that prevent eviction of entries
104/// still referenced outside the cache (i.e., with `Arc` strong count > 1).
105#[derive(Clone, Default)]
106pub struct BlockCacheLifecycle;
107
108impl Lifecycle<(u32, u16), ArcBytes> for BlockCacheLifecycle {
109    type RequestState = ();
110
111    #[inline]
112    fn is_pinned(&self, _key: &(u32, u16), val: &ArcBytes) -> bool {
113        val.is_shared_arc()
114    }
115
116    #[inline]
117    fn begin_request(&self) -> Self::RequestState {}
118
119    #[inline]
120    fn on_evict(&self, _state: &mut Self::RequestState, _key: (u32, u16), _val: ArcBytes) {}
121}
122
123pub type BlockCache = quick_cache::sync::Cache<
124    (u32, u16),
125    ArcBytes,
126    BlockWeighter,
127    BuildHasherDefault<FxHasher>,
128    BlockCacheLifecycle,
129>;
130
131/// Trait abstracting value block reading for `handle_key_match_generic`.
132///
133/// Provides cached reads (small value blocks) and uncached reads (medium value
134/// blocks). Generic over the byte type so it works for both the lookup path
135/// (`ArcBytes` with `BlockCache`) and the iteration path (`RcBytes` with a
136/// single-entry `Option` cache).
137trait ValueBlockCache<B: SharedBytes> {
138    fn get_or_read(
139        self,
140        mmap: &B::MmapHandle,
141        meta: &StaticSortedFileMetaData,
142        block_index: u16,
143    ) -> Result<B>;
144}
145
146/// Bundles the shared block cache with the per-file CRC-verified bitmap,
147/// used on the lookup path.
148#[derive(Clone, Copy)]
149struct ArcBlockCacheReader<'a> {
150    cache: &'a BlockCache,
151    verified_blocks: &'a [AtomicU64],
152}
153
154/// Lookup-path: concurrent `BlockCache` with uncompressed-bypass and
155/// once-per-block CRC verification via `verified_blocks` bitmap.
156impl ValueBlockCache<ArcBytes> for ArcBlockCacheReader<'_> {
157    fn get_or_read(
158        self,
159        mmap: &Arc<Mmap>,
160        meta: &StaticSortedFileMetaData,
161        block_index: u16,
162    ) -> Result<ArcBytes> {
163        get_or_cache_block(mmap, meta, block_index, self.cache, self.verified_blocks)
164    }
165}
166
167/// Iteration-path: lightweight single-entry cache for sequential reads.
168impl ValueBlockCache<RcBytes> for &mut Option<(u16, RcBytes)> {
169    fn get_or_read(
170        self,
171        mmap: &Rc<Mmap>,
172        meta: &StaticSortedFileMetaData,
173        block_index: u16,
174    ) -> Result<RcBytes> {
175        if let Some((idx, block)) = self.as_ref()
176            && *idx == block_index
177        {
178            return Ok(block.clone());
179        }
180        let block: RcBytes = read_block_generic(mmap, meta, block_index)?;
181        *self = Some((block_index, block.clone()));
182        Ok(block)
183    }
184}
185
186#[derive(Clone, Copy, Debug)]
187pub struct StaticSortedFileMetaData {
188    /// The sequence number of this file.
189    pub sequence_number: u32,
190    /// The number of blocks in the SST file.
191    pub block_count: u16,
192}
193
194impl StaticSortedFileMetaData {
195    pub fn block_offsets_start(&self, sst_len: usize) -> usize {
196        let bc: usize = self.block_count.into();
197        sst_len - (bc * size_of::<u32>())
198    }
199}
200
201/// A memory mapped SST file.
202pub struct StaticSortedFile {
203    /// The meta file of this file.
204    meta: StaticSortedFileMetaData,
205    /// The memory mapped file.
206    /// We store as an Arc so we can hand out references (via ArcBytes) that can outlive this
207    /// struct (not that we expect them to outlive it by very much)
208    mmap: Arc<Mmap>,
209    /// One bit per block, set when that block's CRC has been verified at least once.
210    /// Uncompressed (mmap-backed) blocks bypass the `BlockCache`, so without this
211    /// bitmap the CRC would be re-computed on every access. `Relaxed` ordering
212    /// suffices: racing first-time verifications are idempotent.
213    verified_blocks: Box<[AtomicU64]>,
214}
215
216impl StaticSortedFile {
217    /// Opens an SST file at the given path. This memory maps the file, but does not read it yet.
218    /// It's lazy read on demand.
219    pub fn open(db_path: &Path, meta: StaticSortedFileMetaData) -> Result<Self> {
220        let filename = format!("{:08}.sst", meta.sequence_number);
221        let path = db_path.join(&filename);
222        let file = File::open(&path)?;
223        let mmap = unsafe { Mmap::map(file.file()) }.with_context(|| {
224            format!(
225                "Failed to mmap SST file {} ({} bytes)",
226                path.display(),
227                file.metadata().map(|m| m.len()).unwrap_or(0)
228            )
229        })?;
230        #[cfg(unix)]
231        {
232            mmap.advise(memmap2::Advice::Random)?;
233            let offset = meta.block_offsets_start(mmap.len());
234            let _ = mmap.advise_range(memmap2::Advice::Sequential, offset, mmap.len() - offset);
235        }
236        advise_mmap_for_persistence(&mmap)?;
237        let bitmap_words = (meta.block_count as usize).div_ceil(u64::BITS as usize);
238        let verified_blocks = (0..bitmap_words)
239            .map(|_| AtomicU64::new(0))
240            .collect::<Box<[_]>>();
241        Ok(Self {
242            meta,
243            mmap: Arc::new(mmap),
244            verified_blocks,
245        })
246    }
247
248    /// Looks up a key in this file.
249    ///
250    /// If `FIND_ALL` is false, returns after finding the first match.
251    /// If `FIND_ALL` is true, returns all entries with the same key (useful for
252    /// keyspaces where keys are hashes and collisions are possible).
253    pub fn lookup<K: QueryKey, const FIND_ALL: bool>(
254        &self,
255        key_hash: u64,
256        key: &K,
257        key_block_cache: &BlockCache,
258        value_block_cache: &BlockCache,
259    ) -> Result<SstLookupResult> {
260        // There is exactly one index block per file (always the last block).
261        // Read it first, then dispatch directly to the key block it points to.
262        let index_block_index = self.meta.block_count - 1;
263        let index_block = get_or_cache_block(
264            &self.mmap,
265            &self.meta,
266            index_block_index,
267            key_block_cache,
268            &self.verified_blocks,
269        )?;
270        let key_block_index = self.lookup_index_block(&index_block, key_hash)?;
271
272        let key_block_arc = get_or_cache_block(
273            &self.mmap,
274            &self.meta,
275            key_block_index,
276            key_block_cache,
277            &self.verified_blocks,
278        )?;
279        let reader = ArcBlockCacheReader {
280            cache: value_block_cache,
281            verified_blocks: &self.verified_blocks,
282        };
283        let block_type = be::read_u8(&key_block_arc);
284        match block_type {
285            BLOCK_TYPE_KEY_WITH_HASH | BLOCK_TYPE_KEY_NO_HASH => {
286                let has_hash = block_type == BLOCK_TYPE_KEY_WITH_HASH;
287                self.lookup_key_block::<K, FIND_ALL>(key_block_arc, key_hash, key, has_hash, reader)
288            }
289
290            BLOCK_TYPE_FIXED_KEY_WITH_HASH | BLOCK_TYPE_FIXED_KEY_NO_HASH => {
291                let has_hash = block_type == BLOCK_TYPE_FIXED_KEY_WITH_HASH;
292                self.lookup_fixed_key_block::<K, FIND_ALL>(
293                    key_block_arc,
294                    key_hash,
295                    key,
296                    has_hash,
297                    reader,
298                )
299            }
300            _ => {
301                bail!("Invalid block type");
302            }
303        }
304    }
305
306    /// Looks up a hash in a index block.
307    fn lookup_index_block(&self, block: &[u8], hash: u64) -> Result<u16> {
308        ensure!(block.len() >= 3, "index block too short");
309        debug_assert!(
310            be::read_u8(block) == BLOCK_TYPE_INDEX,
311            "expected index block as last block"
312        );
313        let first_block = be::read_u16(&block[1..]);
314        let (entries, remainder) = block[3..].as_chunks::<INDEX_BLOCK_ENTRY_SIZE>();
315        if entries.is_empty() {
316            return Ok(first_block);
317        }
318        if !remainder.is_empty() {
319            bail!("invalid index block, {} extra bytes", remainder.len())
320        }
321        match entries.binary_search_by(|entry| be::read_u64(entry).cmp(&hash)) {
322            Ok(i) => Ok(be::read_u16(&entries[i][8..])),
323            Err(0) => Ok(first_block),
324            Err(i) => Ok(be::read_u16(&entries[i - 1][8..])),
325        }
326    }
327
328    /// Looks up a key in a key block and the value in a value block.
329    ///
330    /// If `FIND_ALL` is false, returns after finding the first match.
331    /// If `FIND_ALL` is true, collects all entries with the same key.
332    fn lookup_key_block<K: QueryKey, const FIND_ALL: bool>(
333        &self,
334        block: ArcBytes,
335        key_hash: u64,
336        key: &K,
337        has_hash: bool,
338        reader: ArcBlockCacheReader<'_>,
339    ) -> Result<SstLookupResult> {
340        let hash_len: u8 = if has_hash { 8 } else { 0 };
341        ensure!(block.len() >= 4, "key block too short");
342        let entry_count = be::read_u24(&block[1..]) as usize;
343        let data = &block[4..];
344        ensure!(
345            data.len() >= entry_count * 4,
346            "key block too short for {entry_count} entries"
347        );
348        let offsets = &data[..entry_count * 4];
349        let entries = &data[entry_count * 4..];
350
351        self.lookup_block_inner::<K, FIND_ALL>(&block, entry_count, key_hash, key, reader, |i| {
352            get_key_entry(offsets, entries, entry_count, i, hash_len)
353        })
354    }
355
356    /// Looks up a key in a fixed-size key block.
357    ///
358    /// Fixed-size key blocks store entries at predictable offsets (no offset table),
359    /// enabling direct indexing during binary search.
360    fn lookup_fixed_key_block<K: QueryKey, const FIND_ALL: bool>(
361        &self,
362        block: ArcBytes,
363        key_hash: u64,
364        key: &K,
365        has_hash: bool,
366        reader: ArcBlockCacheReader<'_>,
367    ) -> Result<SstLookupResult> {
368        let hash_len: u8 = if has_hash { 8 } else { 0 };
369        ensure!(block.len() >= 6, "fixed key block too short");
370        let entry_count = be::read_u24(&block[1..]) as usize;
371        let key_size = be::read_u8(&block[4..]) as usize;
372        let value_type = be::read_u8(&block[5..]);
373        let val_size = entry_val_size(value_type)?;
374        let stride = hash_len as usize + key_size + val_size;
375        let entries = &block[6..];
376        ensure!(
377            entries.len() == entry_count * stride,
378            "fixed key block for {entry_count} entries must is the wrong size"
379        );
380
381        self.lookup_block_inner::<K, FIND_ALL>(&block, entry_count, key_hash, key, reader, |i| {
382            Ok(get_fixed_key_entry(
383                entries, i, hash_len, key_size, value_type, stride,
384            ))
385        })
386    }
387
388    /// Shared binary search + collection logic for both key block variants.
389    ///
390    /// The `get_entry` closure abstracts over the difference between variable-size
391    /// key blocks (offset table lookup) and fixed-size key blocks (stride-based indexing).
392    fn lookup_block_inner<'a, K: QueryKey, const FIND_ALL: bool>(
393        &self,
394        block: &ArcBytes,
395        entry_count: usize,
396        key_hash: u64,
397        key: &K,
398        reader: ArcBlockCacheReader<'_>,
399        get_entry: impl Fn(usize) -> Result<GetKeyEntryResult<'a>>,
400    ) -> Result<SstLookupResult> {
401        let mut l = 0;
402        let mut r = entry_count;
403        // binary search for a matching key
404        while l < r {
405            let m = (l + r) / 2;
406            let GetKeyEntryResult {
407                hash: mid_hash,
408                key: mid_key,
409                ty,
410                val,
411            } = get_entry(m)?;
412
413            let comparison = compare_hash_key(mid_hash, mid_key, key_hash, key);
414
415            match comparison {
416                Ordering::Less => r = m,
417                Ordering::Equal => {
418                    if !FIND_ALL {
419                        // SingleValue mode: each key has exactly one entry
420                        // this is enforced when writing
421                        let result = self.handle_key_match(ty, val, block, reader)?;
422                        return Ok(SstLookupResult::Found(SmallVec::from_buf([result])));
423                    }
424                    // FIND_ALL (MultiValue) mode: collect all values for this key.
425                    // Tombstones (Deleted) sort last within each key group, so we
426                    // scan backward to find the start of the key group, then forward
427                    // to collect all entries. The tombstone, if present, will be the
428                    // last entry in the results.
429                    let mut results = SmallVec::new();
430                    for i in (l..m).rev() {
431                        let GetKeyEntryResult {
432                            hash,
433                            key: entry_key,
434                            ty,
435                            val,
436                        } = get_entry(i)?;
437                        if !entry_matches_key(hash, entry_key, key_hash, key) {
438                            break;
439                        }
440                        results.push(self.handle_key_match(ty, val, block, reader)?);
441                    }
442                    // Technically we could `.reverse()` the items collected by the backwards
443                    // scan, but the only ordering constraint we need to maintain for single
444                    // sst multivalue reads is that a deleted token, if it exists comes last.
445                    // Because all the backwards scan items are strictly before the found item
446                    // we know they don't contain the _last_ item. So we don't care about
447                    // their order.
448
449                    // Add the entry at `m`
450                    results.push(self.handle_key_match(ty, val, block, reader)?);
451                    for i in (m + 1)..r {
452                        let GetKeyEntryResult {
453                            hash,
454                            key: entry_key,
455                            ty,
456                            val,
457                        } = get_entry(i)?;
458                        if !entry_matches_key(hash, entry_key, key_hash, key) {
459                            break;
460                        }
461                        results.push(self.handle_key_match(ty, val, block, reader)?);
462                    }
463                    return Ok(SstLookupResult::Found(results));
464                }
465                Ordering::Greater => l = m + 1,
466            }
467        }
468
469        Ok(SstLookupResult::NotFound)
470    }
471
472    /// Handles a key match by looking up the value.
473    fn handle_key_match(
474        &self,
475        ty: u8,
476        val: &[u8],
477        key_block_arc: &ArcBytes,
478        reader: ArcBlockCacheReader<'_>,
479    ) -> Result<LookupValue> {
480        handle_key_match_generic(&self.mmap, &self.meta, ty, val, key_block_arc, reader)
481    }
482}
483
484/// Gets a block from the cache, or reads it from the mmap and inserts it.
485///
486/// Reads the block header exactly once via `get_raw_block_slice` (which
487/// includes all `strict_checks` bounds guards). Uncompressed blocks bypass
488/// the cache — an mmap-backed `ArcBytes` is cheaper than a cache lookup.
489/// Their CRC is verified at most once per file open, tracked by
490/// `verified_blocks`. Compressed blocks are looked up in `cache`; on a
491/// miss they are decompressed, CRC-verified, and inserted.
492fn get_or_cache_block(
493    mmap: &Arc<Mmap>,
494    meta: &StaticSortedFileMetaData,
495    block_index: u16,
496    cache: &BlockCache,
497    verified_blocks: &[AtomicU64],
498) -> Result<ArcBytes> {
499    let (uncompressed_length, checksum, block_data) = get_raw_block_slice(mmap, meta, block_index)
500        .with_context(|| {
501            format!(
502                "Failed to read raw block {} from {:08}.sst",
503                block_index, meta.sequence_number
504            )
505        })?;
506
507    if uncompressed_length == 0 {
508        // Uncompressed: serve directly from mmap. Verify CRC only once per file open.
509        verify_checksum_once(meta, block_data, checksum, block_index, verified_blocks)?;
510        // SAFETY: block_data points into the mmap backing `mmap`.
511        return Ok(unsafe { ArcBytes::from_mmap(mmap, block_data) });
512    }
513
514    // Compressed: check cache; decompress and insert on miss.
515    Ok(
516        match cache.get_value_or_guard(&(meta.sequence_number, block_index), None) {
517            GuardResult::Value(block) => block,
518            GuardResult::Guard(guard) => {
519                // A cached block may have been evicted, so re-reading still
520                // benefits from the bitmap to skip redundant CRC verification.
521                verify_checksum_once(meta, block_data, checksum, block_index, verified_blocks)?;
522                let block = ArcBytes::from_decompressed(uncompressed_length, block_data)
523                    .with_context(|| {
524                        format!(
525                            "Failed to decompress block {} from {:08}.sst ({} bytes uncompressed)",
526                            block_index, meta.sequence_number, uncompressed_length
527                        )
528                    })?;
529                let _ = guard.insert(block.clone());
530                block
531            }
532            GuardResult::Timeout => unreachable!(),
533        },
534    )
535}
536
537/// Gets the raw block slice directly from a memory-mapped file.
538/// Returns `(uncompressed_length, checksum, block_data)`.
539fn get_raw_block_slice<'a>(
540    mmap: &'a Mmap,
541    meta: &StaticSortedFileMetaData,
542    block_index: u16,
543) -> Result<(u32, u32, &'a [u8])> {
544    #[cfg(feature = "strict_checks")]
545    if block_index >= meta.block_count {
546        bail!(
547            "Corrupted file seq:{} block:{} > number of blocks {} (block_offsets: {:x})",
548            meta.sequence_number,
549            block_index,
550            meta.block_count,
551            meta.block_offsets_start(mmap.len()),
552        );
553    }
554    let offset = meta.block_offsets_start(mmap.len()) + block_index as usize * 4;
555    #[cfg(feature = "strict_checks")]
556    if offset + 4 > mmap.len() {
557        bail!(
558            "Corrupted file seq:{} block:{} block offset locations {} + 4 bytes > file end {} \
559             (block_offsets: {:x})",
560            meta.sequence_number,
561            block_index,
562            offset,
563            mmap.len(),
564            meta.block_offsets_start(mmap.len()),
565        );
566    }
567    let block_start = if block_index == 0 {
568        0
569    } else {
570        be::read_u32(&mmap[offset - 4..]) as usize
571    };
572    let block_end = be::read_u32(&mmap[offset..]) as usize;
573    #[cfg(feature = "strict_checks")]
574    if block_end > mmap.len() || block_start > mmap.len() {
575        bail!(
576            "Corrupted file seq:{} block:{} block {} - {} > file end {} (block_offsets: {:x})",
577            meta.sequence_number,
578            block_index,
579            block_start,
580            block_end,
581            mmap.len(),
582            meta.block_offsets_start(mmap.len()),
583        );
584    }
585    ensure!(
586        block_start + BLOCK_HEADER_SIZE <= block_end,
587        "block {} header truncated in {:08}.sst",
588        block_index,
589        meta.sequence_number
590    );
591    let uncompressed_length = be::read_u32(&mmap[block_start..]);
592    let checksum = be::read_u32(&mmap[block_start + 4..]);
593    let block = &mmap[block_start + BLOCK_HEADER_SIZE..block_end];
594    Ok((uncompressed_length, checksum, block))
595}
596
597/// Verifies the CRC32 checksum of on-disk block data. Returns an error on mismatch.
598fn verify_checksum(
599    meta: &StaticSortedFileMetaData,
600    data: &[u8],
601    expected: u32,
602    block_index: u16,
603) -> Result<()> {
604    let actual = checksum_block(data);
605    if actual != expected {
606        bail!(
607            "Cache corruption detected: checksum mismatch in block {} of {:08}.sst (expected \
608             {:08x}, got {:08x})",
609            block_index,
610            meta.sequence_number,
611            expected,
612            actual
613        );
614    }
615    Ok(())
616}
617
618/// Verifies a block's CRC using the `verified_blocks` bitmap to avoid redundant
619/// work. In practice each block is verified once, but concurrent first-time
620/// accesses may race and verify the same block more than once — this is harmless
621/// since the check is deterministic and idempotent. Verification failures are
622/// *not* recorded in the bitmap, so a corrupted block will be re-checked (and
623/// fail again) on every access.
624fn verify_checksum_once(
625    meta: &StaticSortedFileMetaData,
626    data: &[u8],
627    expected: u32,
628    block_index: u16,
629    verified_blocks: &[AtomicU64],
630) -> Result<()> {
631    let word_idx = block_index as usize / u64::BITS as usize;
632    let bit = 1u64 << (block_index as usize % u64::BITS as usize);
633    if verified_blocks[word_idx].load(AtomicOrdering::Relaxed) & bit != 0 {
634        return Ok(());
635    }
636    verify_checksum(meta, data, expected, block_index)?;
637    verified_blocks[word_idx].fetch_or(bit, AtomicOrdering::Relaxed);
638    Ok(())
639}
640
641/// Returns `(uncompressed_length, checksum, block)` wrapping the raw on-disk
642/// data as the given byte type. Generic over `ArcBytes`/`RcBytes`.
643fn get_raw_block_generic<B: SharedBytes>(
644    mmap: &B::MmapHandle,
645    meta: &StaticSortedFileMetaData,
646    block_index: u16,
647) -> Result<(u32, u32, B)> {
648    let (uncompressed_length, checksum, block) = get_raw_block_slice(mmap, meta, block_index)?;
649    // SAFETY: block points into mmap which backs the MmapHandle.
650    Ok((uncompressed_length, checksum, unsafe {
651        B::from_mmap(mmap, block)
652    }))
653}
654
655/// Reads a block, decompresses if needed, and verifies its checksum.
656/// Generic over the byte type (`ArcBytes` or `RcBytes`).
657#[tracing::instrument(level = "info", name = "reading database block", skip_all)]
658fn read_block_generic<B: SharedBytes>(
659    mmap: &B::MmapHandle,
660    meta: &StaticSortedFileMetaData,
661    block_index: u16,
662) -> Result<B> {
663    let (uncompressed_length, expected_checksum, block) =
664        get_raw_block_slice(mmap, meta, block_index).with_context(|| {
665            format!(
666                "Failed to read raw block {} from {:08}.sst",
667                block_index, meta.sequence_number
668            )
669        })?;
670
671    verify_checksum(meta, block, expected_checksum, block_index)?;
672
673    if uncompressed_length == 0 {
674        // SAFETY: callers guarantee block points into the mmap.
675        return Ok(unsafe { B::from_mmap(mmap, block) });
676    }
677
678    let buffer = B::from_decompressed(uncompressed_length, block).with_context(|| {
679        format!(
680            "Failed to decompress block {} from {:08}.sst ({} bytes uncompressed)",
681            block_index, meta.sequence_number, uncompressed_length
682        )
683    })?;
684    Ok(buffer)
685}
686
687/// Handles a key match by resolving the value reference. Generic over byte type.
688fn handle_key_match_generic<B: SharedBytes>(
689    mmap: &B::MmapHandle,
690    meta: &StaticSortedFileMetaData,
691    ty: u8,
692    val: &[u8],
693    key_block: &B,
694    reader: impl ValueBlockCache<B>,
695) -> Result<LookupValue<B>> {
696    Ok(match ty {
697        KEY_BLOCK_ENTRY_TYPE_SMALL => {
698            let block = be::read_u16(val);
699            let size = be::read_u16(&val[2..]) as usize;
700            let position = be::read_u32(&val[4..]) as usize;
701            let value = reader
702                .get_or_read(mmap, meta, block)?
703                .slice(position..position + size);
704            LookupValue::Slice { value }
705        }
706        KEY_BLOCK_ENTRY_TYPE_MEDIUM => {
707            let block = be::read_u16(val);
708            let value = read_block_generic(mmap, meta, block)?;
709            LookupValue::Slice { value }
710        }
711        KEY_BLOCK_ENTRY_TYPE_BLOB => {
712            let sequence_number = be::read_u32(val);
713            LookupValue::Blob { sequence_number }
714        }
715        KEY_BLOCK_ENTRY_TYPE_DELETED => LookupValue::Deleted,
716        _ => {
717            // Inline value — val is already the correct slice
718            // SAFETY: val points into key_block's data
719            let value = unsafe { key_block.slice_from_subslice(val) };
720            LookupValue::Slice { value }
721        }
722    })
723}
724
725/// An iterator over all entries in a SST file in sorted order.
726pub struct StaticSortedFileIter {
727    /// The memory-mapped file, wrapped in `Rc` for non-atomic refcounting.
728    /// All `RcBytes` slices produced during iteration share this `Rc`.
729    mmap: Rc<Mmap>,
730    /// Metadata (sequence number, block count) needed for block access.
731    meta: StaticSortedFileMetaData,
732
733    /// The root index block entries (body bytes starting after the type byte).
734    /// SST files have exactly one index level.
735    index_entries: RcBytes,
736    /// Total key block references in the index block (first_child + boundary entries).
737    num_index_entries: usize,
738    /// Next index entry to read from the index block.
739    index_pos: usize,
740    current_key_block: CurrentKeyBlock,
741    /// Single-entry value block cache. Within a key block, entries reference
742    /// value blocks sequentially and don't revisit earlier blocks, so caching
743    /// just the current one avoids redundant decompression.
744    value_block_cache: Option<(u16, RcBytes)>,
745}
746
747enum CurrentKeyBlockKind {
748    /// Variable-size entries with an offset table for random access.
749    Variable { offsets: RcBytes, hash_len: u8 },
750    /// Fixed-size entries with uniform key size and value type (no offset table).
751    Fixed {
752        hash_len: u8,
753        key_size: usize,
754        value_type: u8,
755        stride: usize,
756    },
757}
758
759struct CurrentKeyBlock {
760    kind: CurrentKeyBlockKind,
761    entries: RcBytes,
762    /// Number of entries in this key block (max ~819 per 16 KiB block).
763    entry_count: u32,
764    /// Current position within the key block.
765    index: u32,
766}
767
768impl Iterator for StaticSortedFileIter {
769    type Item = Result<LookupEntry>;
770
771    fn next(&mut self) -> Option<Self::Item> {
772        self.next_internal().transpose()
773    }
774}
775
776impl StaticSortedFileIter {
777    /// Opens an SST file for sequential iteration. Uses `MADV_SEQUENTIAL` for
778    /// read-ahead and wraps the mmap in `Rc<Mmap>` directly (no `Arc`),
779    /// eliminating all atomic refcounting during iteration.
780    pub fn open(db_path: &Path, meta: StaticSortedFileMetaData) -> Result<Self> {
781        let filename = format!("{:08}.sst", meta.sequence_number);
782        let path = db_path.join(&filename);
783        let file = File::open(&path)?;
784        let mmap = unsafe { Mmap::map(file.file()) }.with_context(|| {
785            format!(
786                "Failed to mmap SST file {} ({} bytes)",
787                path.display(),
788                file.metadata().map(|m| m.len()).unwrap_or(0)
789            )
790        })?;
791        #[cfg(unix)]
792        mmap.advise(memmap2::Advice::Sequential)?;
793        advise_mmap_for_persistence(&mmap)?;
794        Self::new(Rc::new(mmap), meta)
795            .with_context(|| format!("Unable to open static sorted file {filename}"))
796    }
797
798    fn new(mmap: Rc<Mmap>, meta: StaticSortedFileMetaData) -> Result<Self> {
799        let root_block_index = meta.block_count - 1;
800        let block: RcBytes = read_block_generic(&mmap, &meta, root_block_index)?;
801        let block_type = block[0];
802
803        // The builder always writes an index block as the root block.
804        if block_type != BLOCK_TYPE_INDEX {
805            bail!("Root block must be an index block");
806        }
807        let block_len = block.len();
808        ensure!(block_len >= 3, "index block too short");
809        let index_entries = block.slice(1..block_len);
810        let first_child = be::read_u16(&index_entries);
811        // Index block body layout: [first_child: u16] [hash: u64, block: u16]*
812        // Compute total key block references (first_child + N boundary entries)
813        // using ceil division: (body_len - sizeof(first_child) + ENTRY_SIZE - 1) / ENTRY_SIZE + 1
814        // simplified to (body_len + ENTRY_SIZE - 2) / ENTRY_SIZE
815        let num_index_entries: usize = (index_entries.len() + INDEX_BLOCK_ENTRY_SIZE
816            - size_of::<u16>())
817            / INDEX_BLOCK_ENTRY_SIZE;
818
819        let current_key_block = Self::parse_key_block(&mmap, &meta, first_child)?;
820        Ok(StaticSortedFileIter {
821            mmap,
822            meta,
823            index_entries,
824            num_index_entries,
825            index_pos: 1,
826            current_key_block,
827            value_block_cache: None,
828        })
829    }
830
831    /// Parses a key block at the given index, returning `RcBytes`-backed data.
832    fn parse_key_block(
833        mmap: &Rc<Mmap>,
834        meta: &StaticSortedFileMetaData,
835        block_index: u16,
836    ) -> Result<CurrentKeyBlock> {
837        let block: RcBytes = read_block_generic(mmap, meta, block_index)?;
838        let data = &*block;
839        ensure!(data.len() >= 4, "key block too short");
840        let block_type = data[0];
841        let entry_count = be::read_u24(&data[1..]);
842        match block_type {
843            BLOCK_TYPE_KEY_WITH_HASH | BLOCK_TYPE_KEY_NO_HASH => {
844                let hash_len = if block_type == BLOCK_TYPE_KEY_WITH_HASH {
845                    8
846                } else {
847                    0
848                };
849                let n = entry_count as usize;
850                let offsets_range = 4..4 + n * 4;
851                let entries_range = 4 + n * 4..block.len();
852                let offsets = block.clone().slice(offsets_range);
853                let entries = block.slice(entries_range);
854                Ok(CurrentKeyBlock {
855                    kind: CurrentKeyBlockKind::Variable { offsets, hash_len },
856                    entries,
857                    entry_count,
858                    index: 0,
859                })
860            }
861            BLOCK_TYPE_FIXED_KEY_WITH_HASH | BLOCK_TYPE_FIXED_KEY_NO_HASH => {
862                let hash_len = if block_type == BLOCK_TYPE_FIXED_KEY_WITH_HASH {
863                    8
864                } else {
865                    0
866                };
867                let key_size = data[4] as usize;
868                let value_type = data[5];
869                let val_size = entry_val_size(value_type)?;
870                let stride = hash_len as usize + key_size + val_size;
871                // Header is 6 bytes for fixed-size blocks
872                let entries_range = 6..block.len();
873                let entries = block.slice(entries_range);
874                Ok(CurrentKeyBlock {
875                    kind: CurrentKeyBlockKind::Fixed {
876                        hash_len,
877                        key_size,
878                        value_type,
879                        stride,
880                    },
881                    entries,
882                    entry_count,
883                    index: 0,
884                })
885            }
886            _ => {
887                bail!("Invalid key block type: {block_type}");
888            }
889        }
890    }
891
892    /// Gets the next entry in the file and moves the cursor.
893    fn next_internal(&mut self) -> Result<Option<LookupEntry>> {
894        loop {
895            let kb = &mut self.current_key_block;
896            if kb.index < kb.entry_count {
897                let index = kb.index as usize;
898                let entry_count = kb.entry_count as usize;
899                let GetKeyEntryResult { hash, key, ty, val } = match &kb.kind {
900                    CurrentKeyBlockKind::Variable { offsets, hash_len } => {
901                        get_key_entry(offsets, &kb.entries, entry_count, index, *hash_len)?
902                    }
903                    CurrentKeyBlockKind::Fixed {
904                        hash_len,
905                        key_size,
906                        value_type,
907                        stride,
908                    } => get_fixed_key_entry(
909                        &kb.entries,
910                        index,
911                        *hash_len,
912                        *key_size,
913                        *value_type,
914                        *stride,
915                    ),
916                };
917                let full_hash = if hash.is_empty() {
918                    crate::key::hash_key(&key)
919                } else {
920                    be::read_u64(hash)
921                };
922                let value = if ty == KEY_BLOCK_ENTRY_TYPE_MEDIUM {
923                    let block = be::read_u16(val);
924                    let (uncompressed_size, checksum, block) =
925                        get_raw_block_generic(&self.mmap, &self.meta, block)?;
926                    IterValue::Medium {
927                        uncompressed_size,
928                        checksum,
929                        block,
930                    }
931                } else {
932                    handle_key_match_generic(
933                        &self.mmap,
934                        &self.meta,
935                        ty,
936                        val,
937                        &kb.entries,
938                        &mut self.value_block_cache,
939                    )?
940                    .into()
941                };
942                let entry = LookupEntry {
943                    hash: full_hash,
944                    key: unsafe { kb.entries.slice_from_subslice(key) },
945                    value,
946                };
947                kb.index += 1;
948                return Ok(Some(entry));
949            }
950            if self.index_pos < self.num_index_entries {
951                let base = self.index_pos * INDEX_BLOCK_ENTRY_SIZE;
952                let block_index = be::read_u16(&self.index_entries[base..]);
953                self.index_pos += 1;
954                self.current_key_block =
955                    Self::parse_key_block(&self.mmap, &self.meta, block_index)?;
956            } else {
957                return Ok(None);
958            }
959        }
960    }
961}
962
963struct GetKeyEntryResult<'l> {
964    hash: &'l [u8],
965    key: &'l [u8],
966    ty: u8,
967    val: &'l [u8],
968}
969
970/// Compares a query (full_hash + query_key) against an entry (entry_hash + entry_key).
971/// Returns the ordering of query relative to entry.
972/// When entry_hash is empty, computes full hash from entry_key.
973fn compare_hash_key<K: QueryKey>(
974    entry_hash: &[u8],
975    entry_key: &[u8],
976    full_hash: u64,
977    query_key: &K,
978) -> Ordering {
979    if entry_hash.is_empty() {
980        // No hash stored - compute full hash from entry key
981        let entry_full_hash = crate::key::hash_key(&entry_key);
982        match full_hash.cmp(&entry_full_hash) {
983            Ordering::Equal => query_key.cmp(entry_key),
984            ord => ord,
985        }
986    } else {
987        // Full 8-byte hash stored - compare hashes first
988        let full_hash_bytes = full_hash.to_be_bytes();
989        match full_hash_bytes[..].cmp(entry_hash) {
990            Ordering::Equal => query_key.cmp(entry_key),
991            ord => ord,
992        }
993    }
994}
995
996/// Checks if a query key equals an entry key, optionally comparing stored hashes first.
997/// When a hash is stored (8 bytes), compares hashes before keys for speed.
998/// When no hash is stored, compares keys directly (avoiding hash recomputation).
999fn entry_matches_key<K: QueryKey>(
1000    entry_hash: &[u8],
1001    entry_key: &[u8],
1002    full_hash: u64,
1003    query_key: &K,
1004) -> bool {
1005    if entry_hash.is_empty() {
1006        // No hash stored - compare keys directly instead of recomputing hash
1007        query_key.cmp(entry_key) == Ordering::Equal
1008    } else {
1009        // Hash stored - cheap 8-byte comparison first, then key comparison
1010        full_hash.to_be_bytes()[..] == *entry_hash && query_key.cmp(entry_key) == Ordering::Equal
1011    }
1012}
1013
1014/// Returns the byte size of the value portion for a given key block entry type.
1015fn entry_val_size(ty: u8) -> Result<usize> {
1016    match ty {
1017        KEY_BLOCK_ENTRY_TYPE_SMALL => Ok(SMALL_VALUE_REF_SIZE),
1018        KEY_BLOCK_ENTRY_TYPE_MEDIUM => Ok(MEDIUM_VALUE_REF_SIZE),
1019        KEY_BLOCK_ENTRY_TYPE_BLOB => Ok(BLOB_VALUE_REF_SIZE),
1020        KEY_BLOCK_ENTRY_TYPE_DELETED => Ok(DELETED_VALUE_REF_SIZE),
1021        ty if ty >= KEY_BLOCK_ENTRY_TYPE_INLINE_MIN => {
1022            Ok((ty - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize)
1023        }
1024        _ => bail!("Invalid key block entry type: {ty}"),
1025    }
1026}
1027
1028/// Reads the type and start offset from an offset table entry.
1029/// Each entry is 4 bytes: 1 byte type + 3 bytes BE offset.
1030#[inline(always)]
1031fn read_offset_entry(offsets: &[u8], index: usize) -> (u8, usize) {
1032    let base = index * 4;
1033    let word = be::read_u32(&offsets[base..]);
1034    let ty = (word >> 24) as u8;
1035    let offset = (word & 0x00FF_FFFF) as usize;
1036    (ty, offset)
1037}
1038
1039/// Reads a key entry from a key block.
1040fn get_key_entry<'l>(
1041    offsets: &[u8],
1042    entries: &'l [u8],
1043    entry_count: usize,
1044    index: usize,
1045    hash_len: u8,
1046) -> Result<GetKeyEntryResult<'l>> {
1047    let hash_len_usize = hash_len as usize;
1048    let (ty, start) = read_offset_entry(offsets, index);
1049    let end = if index == entry_count - 1 {
1050        entries.len()
1051    } else {
1052        let (_, next_start) = read_offset_entry(offsets, index + 1);
1053        next_start
1054    };
1055    // Return the raw hash bytes slice (0-8 bytes depending on hash_len)
1056    let hash = &entries[start..start + hash_len_usize];
1057    let val_size = entry_val_size(ty)?;
1058    Ok(GetKeyEntryResult {
1059        hash,
1060        key: &entries[start + hash_len_usize..end - val_size],
1061        ty,
1062        val: &entries[end - val_size..end],
1063    })
1064}
1065
1066/// Reads a key entry from a fixed-size key block by direct indexing.
1067///
1068/// All entries have the same key size and value type, so positions are computed
1069/// arithmetically with no offset table indirection.
1070fn get_fixed_key_entry<'l>(
1071    entries: &'l [u8],
1072    index: usize,
1073    hash_len: u8,
1074    key_size: usize,
1075    value_type: u8,
1076    stride: usize,
1077) -> GetKeyEntryResult<'l> {
1078    let hash_len_usize = hash_len as usize;
1079    let start = index * stride;
1080    GetKeyEntryResult {
1081        hash: &entries[start..start + hash_len_usize],
1082        key: &entries[start + hash_len_usize..start + hash_len_usize + key_size],
1083        ty: value_type,
1084        val: &entries[start + hash_len_usize + key_size..(index + 1) * stride],
1085    }
1086}