Skip to main content

turbo_persistence/
meta_file.rs

1use std::{
2    cmp::Ordering,
3    fmt::Display,
4    path::{Path, PathBuf},
5    sync::OnceLock,
6};
7
8use anyhow::{Context, Result, bail};
9use bitfield::bitfield;
10use byteorder::{BE, ReadBytesExt};
11use fs_err::File;
12use memmap2::{Mmap, MmapOptions};
13use smallvec::SmallVec;
14use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref, big_endian as be};
15
16use crate::{
17    QueryKey,
18    lookup_entry::LookupValue,
19    mmap_helper::advise_mmap_for_persistence,
20    static_sorted_file::{BlockCache, SstLookupResult, StaticSortedFile, StaticSortedFileMetaData},
21};
22
23bitfield! {
24    #[derive(Clone, Copy, Default)]
25    pub struct MetaEntryFlags(u32);
26    impl Debug;
27    impl From<u32>;
28    /// The SST file was compacted and none of the entries have been accessed recently.
29    pub cold, set_cold: 0;
30    /// The SST file was freshly written and has not been compacted yet.
31    pub fresh, set_fresh: 1;
32}
33
34impl MetaEntryFlags {
35    pub const FRESH: MetaEntryFlags = MetaEntryFlags(0b10);
36    pub const COLD: MetaEntryFlags = MetaEntryFlags(0b01);
37    pub const WARM: MetaEntryFlags = MetaEntryFlags(0b00);
38}
39
40impl Display for MetaEntryFlags {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        if self.fresh() {
43            f.pad_integral(true, "", "fresh")
44        } else if self.cold() {
45            f.pad_integral(true, "", "cold")
46        } else {
47            f.pad_integral(true, "", "warm")
48        }
49    }
50}
51
52/// On-disk layout of a single entry header in the `.meta` file.
53///
54/// Fields are big-endian to match the existing wire format written by [`MetaFileBuilder`].
55#[repr(C, packed)]
56#[derive(FromBytes, IntoBytes, Immutable, KnownLayout, Clone, Copy)]
57pub(crate) struct EntryHeader {
58    sequence_number: be::U32,
59    block_count: be::U16,
60    min_hash: be::U64,
61    max_hash: be::U64,
62    size: be::U64,
63    flags: be::U32,
64    amqf_end_offset: be::U32,
65}
66
67impl EntryHeader {
68    pub(crate) fn new(
69        sequence_number: u32,
70        block_count: u16,
71        min_hash: u64,
72        max_hash: u64,
73        size: u64,
74        flags: MetaEntryFlags,
75        amqf_end_offset: u32,
76    ) -> Self {
77        Self {
78            sequence_number: be::U32::new(sequence_number),
79            block_count: be::U16::new(block_count),
80            min_hash: be::U64::new(min_hash),
81            max_hash: be::U64::new(max_hash),
82            size: be::U64::new(size),
83            flags: be::U32::new(flags.0),
84            amqf_end_offset: be::U32::new(amqf_end_offset),
85        }
86    }
87}
88
89/// # Safety
90///
91/// `MetaEntry` stores a `FilterRef<'static>` with a transmuted lifetime that actually borrows
92/// from the parent [`MetaFile`]'s mmap. This is safe because entries are only accessed by
93/// reference through `MetaFile` and are never moved out.
94///
95/// For this reason this type should not implement Clone or Copy.
96pub struct MetaEntry {
97    /// The metadata for the static sorted file.
98    sst_data: StaticSortedFileMetaData,
99    /// The key family of the SST file.
100    family: u32,
101    /// The minimum hash value of the keys in the SST file.
102    min_hash: u64,
103    /// The maximum hash value of the keys in the SST file.
104    max_hash: u64,
105    /// The size of the SST file in bytes.
106    size: u64,
107    /// The status flags for this entry.
108    flags: MetaEntryFlags,
109    /// Byte offset range of the raw AMQF data within the mmap, used for carrying forward
110    /// serialized bytes during compaction without re-serializing.
111    amqf_data_offset: std::ops::Range<u32>,
112    /// The AMQF filter for this file, eagerly deserialized as a zero-copy [`qfilter::FilterRef`]
113    /// that borrows directly from the parent [`MetaFile`]'s memory-mapped file.
114    ///
115    /// The `'static` lifetime is transmuted — the actual borrow is from `MetaFile::mmap`.
116    amqf: qfilter::FilterRef<'static>,
117    /// The static sorted file that is lazily loaded
118    sst: OnceLock<StaticSortedFile>,
119}
120
121// Safety: FilterRef is a read-only view into the mmap which is Send+Sync.
122unsafe impl Send for MetaEntry {}
123unsafe impl Sync for MetaEntry {}
124
125impl MetaEntry {
126    pub fn sequence_number(&self) -> u32 {
127        self.sst_data.sequence_number
128    }
129
130    pub fn size(&self) -> u64 {
131        self.size
132    }
133
134    pub fn flags(&self) -> MetaEntryFlags {
135        self.flags
136    }
137
138    pub fn amqf_size(&self) -> u32 {
139        self.amqf_data_offset.end - self.amqf_data_offset.start
140    }
141
142    /// Returns the raw serialized AMQF bytes from the mmap.
143    pub fn raw_amqf<'l>(&self, amqf_data: &'l [u8]) -> &'l [u8] {
144        &amqf_data[self.amqf_data_offset.start as usize..self.amqf_data_offset.end as usize]
145    }
146
147    fn sst(&self, meta: &MetaFile) -> Result<&StaticSortedFile> {
148        self.sst.get_or_try_init(|| {
149            StaticSortedFile::open(&meta.db_path, self.sst_data).with_context(|| {
150                format!(
151                    "Unable to open static sorted file referenced from {:08}.meta",
152                    meta.sequence_number()
153                )
154            })
155        })
156    }
157
158    /// Returns the key family and hash range of this file.
159    pub fn range(&self) -> StaticSortedFileRange {
160        StaticSortedFileRange {
161            family: self.family,
162            min_hash: self.min_hash,
163            max_hash: self.max_hash,
164        }
165    }
166
167    pub fn min_hash(&self) -> u64 {
168        self.min_hash
169    }
170
171    pub fn max_hash(&self) -> u64 {
172        self.max_hash
173    }
174
175    pub fn block_count(&self) -> u16 {
176        self.sst_data.block_count
177    }
178
179    /// Returns the SST metadata needed to open the file independently.
180    /// Used during compaction to avoid caching mmaps on the MetaEntry.
181    pub fn sst_metadata(&self) -> StaticSortedFileMetaData {
182        self.sst_data
183    }
184}
185
186/// The result of a lookup operation.
187pub enum MetaLookupResult {
188    /// The key was not found because it is from a different key family.
189    FamilyMiss,
190    /// The key was not found because it is out of the range of this SST file. But it was the
191    /// correct key family.
192    RangeMiss,
193    /// The key was not found because it was not in the AMQF filter. But it was in the range.
194    QuickFilterMiss,
195    /// The key was looked up in the SST file. It was in the AMQF filter.
196    SstLookup(SstLookupResult),
197}
198
199/// The result of a batch lookup operation.
200#[derive(Default)]
201pub struct MetaBatchLookupResult {
202    /// The key was not found because it is from a different key family.
203    #[cfg(feature = "stats")]
204    pub family_miss: bool,
205    /// The key was not found because it is out of the range of this SST file. But it was the
206    /// correct key family.
207    #[cfg(feature = "stats")]
208    pub range_misses: usize,
209    /// The key was not found because it was not in the AMQF filter. But it was in the range.
210    #[cfg(feature = "stats")]
211    pub quick_filter_misses: usize,
212    /// The key was unsuccessfully looked up in the SST file. It was in the AMQF filter.
213    #[cfg(feature = "stats")]
214    pub sst_misses: usize,
215    /// The key was found in the SST file.
216    #[cfg(feature = "stats")]
217    pub hits: usize,
218}
219
220/// The key family and hash range of an SST file.
221#[derive(Clone, Copy)]
222pub struct StaticSortedFileRange {
223    pub family: u32,
224    pub min_hash: u64,
225    pub max_hash: u64,
226}
227
228/// # Safety
229///
230/// `entries` **must** be declared before `mmap` so that Rust's field drop order (declaration
231/// order) drops all `FilterRef`s before the mmap is unmapped.  Reordering these fields would
232/// be unsound.
233pub struct MetaFile {
234    /// The database path
235    db_path: PathBuf,
236    /// The sequence number of this file.
237    sequence_number: u32,
238    /// The key family of the SST files in this meta file.
239    family: u32,
240    /// The entries of the file. Dropped before `mmap` (field declaration order).
241    entries: Vec<MetaEntry>,
242    /// The entries that have been marked as obsolete.
243    obsolete_entries: Vec<u32>,
244    /// The obsolete SST files.
245    obsolete_sst_files: Vec<u32>,
246    /// Byte offset within the mmap where the AMQF data region starts (i.e. the header length).
247    /// Entry AMQF offsets and used-keys offsets are relative to this position.
248    amqf_data_start: u32,
249    /// The offset of the start of the "used keys" AMQF data relative to the AMQF data region.
250    start_of_used_keys_amqf_data_offset: u32,
251    /// The offset of the end of the "used keys" AMQF data relative to the AMQF data region.
252    end_of_used_keys_amqf_data_offset: u32,
253    /// The memory mapped file.
254    /// The entire memory-mapped file. Must be the last field that matters for drop order —
255    /// `entries` contains `FilterRef`s that borrow from this mmap.
256    mmap: Mmap,
257}
258
259impl MetaFile {
260    /// Opens a meta file at the given path. Memory maps the entire file and eagerly deserializes
261    /// all AMQF filters as zero-copy [`qfilter::FilterRef`]s that borrow from the mmap.
262    pub fn open(db_path: &Path, sequence_number: u32) -> Result<Self> {
263        let filename = format!("{sequence_number:08}.meta");
264        let path = db_path.join(&filename);
265        Self::open_internal(db_path.to_path_buf(), sequence_number, &path)
266            .with_context(|| format!("Unable to open meta file {filename}"))
267    }
268
269    fn open_internal(db_path: PathBuf, sequence_number: u32, path: &Path) -> Result<Self> {
270        let file = File::open(path)?;
271        let mmap = unsafe { MmapOptions::new().map(file.file()) }.context("Failed to mmap")?;
272        #[cfg(unix)]
273        mmap.advise(memmap2::Advice::Random)
274            .context("Failed to advise mmap")?;
275        advise_mmap_for_persistence(&mmap)?;
276        // Parse the header from the mmap via ReadBytesExt on &[u8].
277        let mut reader: &[u8] = &mmap;
278        let magic = reader.read_u32::<BE>()?;
279        if magic != 0xFE4ADA4A {
280            bail!("Invalid magic number");
281        }
282        let family = reader.read_u32::<BE>()?;
283        let obsolete_count = reader.read_u32::<BE>()?;
284        let mut obsolete_sst_files = Vec::with_capacity(obsolete_count as usize);
285        for _ in 0..obsolete_count {
286            obsolete_sst_files.push(reader.read_u32::<BE>()?);
287        }
288
289        let count = reader.read_u32::<BE>()?;
290
291        // Compute where the AMQF data region starts so we can deserialize filters inline.
292        // Remaining header: count * ENTRY_HEADER_SIZE + used_keys_end_offset.
293        let header_so_far = (mmap.len() - reader.len()) as u32;
294        let amqf_data_start =
295            header_so_far + count * (size_of::<EntryHeader>() as u32) + size_of::<u32>() as u32;
296        let amqf_data = &mmap[amqf_data_start as usize..];
297
298        // Parse entries and eagerly deserialize AMQF filters as zero-copy FilterRefs.
299        let mut entries = Vec::with_capacity(count as usize);
300        let mut start_of_amqf_data_offset: u32 = 0;
301        for _ in 0..count {
302            let (header, rest): (Ref<&[u8], EntryHeader>, _) = Ref::from_prefix(reader)
303                .ok()
304                .context("Entry header out of bounds")?;
305            reader = rest;
306            let sst_data = StaticSortedFileMetaData {
307                sequence_number: header.sequence_number.get(),
308                block_count: header.block_count.get(),
309            };
310            let min_hash = header.min_hash.get();
311            let max_hash = header.max_hash.get();
312            let size = header.size.get();
313            let flags = MetaEntryFlags(header.flags.get());
314            let end_of_amqf_data_offset = header.amqf_end_offset.get();
315
316            let amqf_bytes = amqf_data
317                .get(start_of_amqf_data_offset as usize..end_of_amqf_data_offset as usize)
318                .expect("AMQF data out of bounds");
319            // Deserialize the filter borrowing from the mmap, then erase the lifetime.
320            let amqf: qfilter::FilterRef<'_> =
321                postcard::from_bytes(amqf_bytes).with_context(|| {
322                    format!(
323                        "Failed to deserialize AMQF from {:08}.meta for {:08}.sst",
324                        sequence_number, sst_data.sequence_number
325                    )
326                })?;
327            // Safety: the mmap is kept alive by MetaFile and is dropped after entries (field
328            // declaration order), so the borrow remains valid for the lifetime of the MetaEntry.
329            let amqf: qfilter::FilterRef<'static> = unsafe { std::mem::transmute(amqf) };
330
331            entries.push(MetaEntry {
332                sst_data,
333                family,
334                min_hash,
335                max_hash,
336                size,
337                flags,
338                amqf_data_offset: start_of_amqf_data_offset..end_of_amqf_data_offset,
339                amqf,
340                sst: OnceLock::new(),
341            });
342            start_of_amqf_data_offset = end_of_amqf_data_offset;
343        }
344
345        let start_of_used_keys_amqf_data_offset = start_of_amqf_data_offset;
346        let end_of_used_keys_amqf_data_offset = reader.read_u32::<BE>()?;
347
348        Ok(Self {
349            db_path,
350            sequence_number,
351            family,
352            entries,
353            obsolete_entries: Vec::new(),
354            obsolete_sst_files,
355            amqf_data_start,
356            start_of_used_keys_amqf_data_offset,
357            end_of_used_keys_amqf_data_offset,
358            mmap,
359        })
360    }
361
362    pub fn clear_cache(&mut self) {
363        for entry in self.entries.iter_mut() {
364            entry.sst.take();
365        }
366    }
367
368    pub fn prepare_sst_cache(&self) {
369        for entry in self.entries.iter() {
370            let _ = entry.sst(self);
371        }
372    }
373
374    pub fn sequence_number(&self) -> u32 {
375        self.sequence_number
376    }
377
378    pub fn family(&self) -> u32 {
379        self.family
380    }
381
382    /// The on-disk size of this meta file in bytes (the length of its memory map).
383    pub fn byte_size(&self) -> u64 {
384        self.mmap.len() as u64
385    }
386
387    pub fn entries(&self) -> &[MetaEntry] {
388        &self.entries
389    }
390
391    pub fn entry(&self, index: u32) -> &MetaEntry {
392        let index = index as usize;
393        &self.entries[index]
394    }
395
396    pub fn amqf_data(&self) -> &[u8] {
397        &self.mmap[self.amqf_data_start as usize..]
398    }
399
400    pub fn deserialize_used_key_hashes_amqf(&self) -> Result<Option<qfilter::FilterRef<'_>>> {
401        if self.start_of_used_keys_amqf_data_offset == self.end_of_used_keys_amqf_data_offset {
402            return Ok(None);
403        }
404        let amqf = &self.amqf_data()[self.start_of_used_keys_amqf_data_offset as usize
405            ..self.end_of_used_keys_amqf_data_offset as usize];
406        Ok(Some(postcard::from_bytes(amqf).with_context(|| {
407            format!(
408                "Failed to deserialize used key hashes AMQF from {:08}.meta",
409                self.sequence_number
410            )
411        })?))
412    }
413
414    pub fn retain_entries(&mut self, mut predicate: impl FnMut(u32) -> bool) -> bool {
415        let old_len = self.entries.len();
416        self.entries.retain(|entry| {
417            if predicate(entry.sst_data.sequence_number) {
418                true
419            } else {
420                self.obsolete_entries.push(entry.sst_data.sequence_number);
421                false
422            }
423        });
424        old_len != self.entries.len()
425    }
426
427    pub fn obsolete_entries(&self) -> &[u32] {
428        &self.obsolete_entries
429    }
430
431    pub fn has_active_entries(&self) -> bool {
432        !self.entries.is_empty()
433    }
434
435    pub fn obsolete_sst_files(&self) -> &[u32] {
436        &self.obsolete_sst_files
437    }
438
439    /// Looks up a key in this meta file.
440    ///
441    /// If `FIND_ALL` is false, returns after finding the first match.
442    /// If `FIND_ALL` is true, returns all entries with the same key from all SST files
443    /// (useful for keyspaces where keys are hashes and collisions are possible).
444    pub fn lookup<K: QueryKey, const FIND_ALL: bool>(
445        &self,
446        key_family: u32,
447        key_hash: u64,
448        key: &K,
449        key_block_cache: &BlockCache,
450        value_block_cache: &BlockCache,
451    ) -> Result<MetaLookupResult> {
452        if key_family != self.family {
453            return Ok(MetaLookupResult::FamilyMiss);
454        }
455        let mut miss_result = MetaLookupResult::RangeMiss;
456        let mut all_results: SmallVec<[LookupValue; 1]> = SmallVec::new();
457
458        for entry in self.entries.iter().rev() {
459            if key_hash < entry.min_hash || key_hash > entry.max_hash {
460                continue;
461            }
462            if !entry.amqf.contains_fingerprint(key_hash) {
463                miss_result = MetaLookupResult::QuickFilterMiss;
464                continue;
465            }
466
467            let result = entry.sst(self)?.lookup::<K, FIND_ALL>(
468                key_hash,
469                key,
470                key_block_cache,
471                value_block_cache,
472            )?;
473
474            match result {
475                SstLookupResult::NotFound => {
476                    // continue searching other sst files
477                }
478                SstLookupResult::Found(values) => {
479                    if !FIND_ALL {
480                        // Return immediately with the first result
481                        return Ok(MetaLookupResult::SstLookup(SstLookupResult::Found(values)));
482                    }
483                    // Check for tombstone — stops search across older SSTs within this meta file.
484                    // Since tombstones sort last within a key group, if the last value is Deleted,
485                    // we have a tombstone.
486                    let has_tombstone = values.last().is_some_and(|v| *v == LookupValue::Deleted);
487                    all_results.extend(values);
488                    if has_tombstone {
489                        return Ok(MetaLookupResult::SstLookup(SstLookupResult::Found(
490                            all_results,
491                        )));
492                    }
493                }
494            }
495        }
496
497        if FIND_ALL && !all_results.is_empty() {
498            return Ok(MetaLookupResult::SstLookup(SstLookupResult::Found(
499                all_results,
500            )));
501        }
502
503        Ok(miss_result)
504    }
505
506    pub fn batch_lookup<K: QueryKey>(
507        &self,
508        key_family: u32,
509        keys: &[K],
510        cells: &mut [(u64, usize, Option<LookupValue>)],
511        empty_cells: &mut usize,
512        key_block_cache: &BlockCache,
513        value_block_cache: &BlockCache,
514    ) -> Result<MetaBatchLookupResult> {
515        if key_family != self.family {
516            #[cfg(feature = "stats")]
517            return Ok(MetaBatchLookupResult {
518                family_miss: true,
519                ..Default::default()
520            });
521            #[cfg(not(feature = "stats"))]
522            return Ok(MetaBatchLookupResult {});
523        }
524        debug_assert!(
525            cells.is_sorted_by_key(|(hash, _, _)| *hash),
526            "Cells must be sorted by key hash"
527        );
528        #[allow(unused_mut, reason = "It's used when stats are enabled")]
529        let mut lookup_result = MetaBatchLookupResult::default();
530        for entry in self.entries.iter().rev() {
531            let start_index = cells
532                .binary_search_by(|(hash, _, _)| hash.cmp(&entry.min_hash).then(Ordering::Greater))
533                .err()
534                .unwrap();
535            if start_index >= cells.len() {
536                #[cfg(feature = "stats")]
537                {
538                    lookup_result.range_misses += 1;
539                }
540                continue;
541            }
542            let end_index = cells
543                .binary_search_by(|(hash, _, _)| hash.cmp(&entry.max_hash).then(Ordering::Less))
544                .err()
545                .unwrap()
546                .checked_sub(1);
547            let Some(end_index) = end_index else {
548                #[cfg(feature = "stats")]
549                {
550                    lookup_result.range_misses += 1;
551                }
552                continue;
553            };
554            if start_index > end_index {
555                #[cfg(feature = "stats")]
556                {
557                    lookup_result.range_misses += 1;
558                }
559                continue;
560            }
561            for (hash, index, result) in &mut cells[start_index..=end_index] {
562                debug_assert!(
563                    *hash >= entry.min_hash && *hash <= entry.max_hash,
564                    "Key hash out of range"
565                );
566                if result.is_some() {
567                    continue;
568                }
569                if !entry.amqf.contains_fingerprint(*hash) {
570                    #[cfg(feature = "stats")]
571                    {
572                        lookup_result.quick_filter_misses += 1;
573                    }
574                    continue;
575                }
576                let sst_result = entry.sst(self)?.lookup::<_, false>(
577                    *hash,
578                    &keys[*index],
579                    key_block_cache,
580                    value_block_cache,
581                )?;
582                if let SstLookupResult::Found(mut values) = sst_result {
583                    // find_all=false guarantees exactly one result
584                    debug_assert!(values.len() == 1);
585                    let Some(value) = values.pop() else {
586                        unreachable!()
587                    };
588                    *result = Some(value);
589                    *empty_cells -= 1;
590                    #[cfg(feature = "stats")]
591                    {
592                        lookup_result.hits += 1;
593                    }
594                    if *empty_cells == 0 {
595                        return Ok(lookup_result);
596                    }
597                } else {
598                    #[cfg(feature = "stats")]
599                    {
600                        lookup_result.sst_misses += 1;
601                    }
602                }
603            }
604        }
605        Ok(lookup_result)
606    }
607}