capture used key hashes and store them in the meta file

Author: sokra

Cargo.lock

@@ -424,6 +424,7 @@ napi = { version = "2", default-features = false, features = [
   "napi5",
   "compat-mode",
 ] }
+nohash-hasher = "0.2.0"
 notify = "8.1.0"
 once_cell = "1.17.1"
 owo-colors = "4.2.2"

Cargo.toml

@@ -16,7 +16,7 @@ fn main() -> Result<()> {
         bail!("The provided path does not exist: {}", path.display());
     }
 
-    let db: TurboPersistence<SerialScheduler> = TurboPersistence::open_read_only(path)?;
+    let db: TurboPersistence<SerialScheduler, 0> = TurboPersistence::open_read_only(path)?;
     let meta_info = db
         .meta_info()
         .context("Failed to retrieve meta information")?;

turbopack/crates/turbo-persistence-tools/src/main.rs

@@ -13,12 +13,14 @@ print_stats = ["stats"]
 
 [dependencies]
 anyhow = { workspace = true }
+dashmap = { workspace = true}
 either = { workspace = true }
 pot = "3.0.0"
 byteorder = { workspace = true }
 jiff = "0.2.10"
 lzzzz = "1.1.0"
 memmap2 = "0.9.5"
+nohash-hasher = { workspace = true }
 parking_lot = { workspace = true }
 qfilter = { version = "0.2.4", features = ["serde"] }
 quick_cache = { workspace = true }

turbopack/crates/turbo-persistence/Cargo.toml

@@ -50,8 +50,10 @@ A meta file can contain metadata about multiple SST files. The metadata is store
     - 8 bytes max hash
     - 8 bytes SST file size
     - 4 bytes end of AMQF offset relative to start of all AMQF data
+  - 4 bytes end of AMQF offset relative to start of all AMQF data of the "used key hashes" AMQF
 - foreach described SST file
   - serialized AMQF
+- serialized "used key hashes" AMQF
 
 ### SST file
 

turbopack/crates/turbo-persistence/README.md

@@ -11,8 +11,10 @@ use std::{
 
 use anyhow::{Context, Result, bail};
 use byteorder::{BE, ReadBytesExt, WriteBytesExt};
+use dashmap::DashSet;
 use jiff::Timestamp;
 use memmap2::Mmap;
+use nohash_hasher::BuildNoHashHasher;
 use parking_lot::{Mutex, RwLock};
 
 pub use crate::compaction::selector::CompactConfig;
@@ -104,15 +106,15 @@ struct TrackedStats {
 
 /// TurboPersistence is a persistent key-value store. It is limited to a single writer at a time
 /// using a single write batch. It allows for concurrent reads.
-pub struct TurboPersistence<S: ParallelScheduler> {
+pub struct TurboPersistence<S: ParallelScheduler, const FAMILIES: usize> {
     parallel_scheduler: S,
     /// The path to the directory where the database is stored
     path: PathBuf,
     /// If true, the database is opened in read-only mode. In this mode, no writes are allowed and
     /// no modification on the database is performed.
     read_only: bool,
     /// The inner state of the database. Writing will update that.
-    inner: RwLock<Inner>,
+    inner: RwLock<Inner<FAMILIES>>,
     /// A flag to indicate if a write operation is currently active. Prevents multiple concurrent
     /// write operations.
     active_write_operation: AtomicBool,
@@ -128,11 +130,15 @@ pub struct TurboPersistence<S: ParallelScheduler> {
 }
 
 /// The inner state of the database.
-struct Inner {
+struct Inner<const FAMILIES: usize> {
     /// The list of meta files in the database. This is used to derive the SST files.
     meta_files: Vec<MetaFile>,
     /// The current sequence number for the database.
     current_sequence_number: u32,
+    /// The in progress set of hashes of keys that have been accessed.
+    /// It will be flushed onto disk (into a meta file) on next commit.
+    /// It's a dashset to allow modification while only tracking a read lock on Inner.
+    accessed_key_hashes: [DashSet<u64, BuildNoHashHasher<u64>>; FAMILIES],
 }
 
 pub struct CommitOptions {
@@ -145,7 +151,7 @@ pub struct CommitOptions {
     keys_written: u64,
 }
 
-impl<S: ParallelScheduler + Default> TurboPersistence<S> {
+impl<S: ParallelScheduler + Default, const FAMILIES: usize> TurboPersistence<S, FAMILIES> {
     /// Open a TurboPersistence database at the given path.
     /// This will read the directory and might performance cleanup when the database was not closed
     /// properly. Cleanup only requires to read a few bytes from a few files and to delete
@@ -161,7 +167,7 @@ impl<S: ParallelScheduler + Default> TurboPersistence<S> {
     }
 }
 
-impl<S: ParallelScheduler> TurboPersistence<S> {
+impl<S: ParallelScheduler, const FAMILIES: usize> TurboPersistence<S, FAMILIES> {
     fn new(path: PathBuf, read_only: bool, parallel_scheduler: S) -> Self {
         Self {
             parallel_scheduler,
@@ -170,6 +176,8 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
             inner: RwLock::new(Inner {
                 meta_files: Vec::new(),
                 current_sequence_number: 0,
+                accessed_key_hashes: [(); FAMILIES]
+                    .map(|_| DashSet::with_hasher(BuildNoHashHasher::default())),
             }),
             active_write_operation: AtomicBool::new(false),
             amqf_cache: AmqfCache::with(
@@ -405,7 +413,7 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
     /// time. The WriteBatch need to be committed with [`TurboPersistence::commit_write_batch`].
     /// Note that the WriteBatch might start writing data to disk while it's filled up with data.
     /// This data will only become visible after the WriteBatch is committed.
-    pub fn write_batch<K: StoreKey + Send + Sync + 'static, const FAMILIES: usize>(
+    pub fn write_batch<K: StoreKey + Send + Sync + 'static>(
         &self,
     ) -> Result<WriteBatch<K, S, FAMILIES>> {
         if self.read_only {
@@ -443,7 +451,7 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
 
     /// Commits a WriteBatch to the database. This will finish writing the data to disk and make it
     /// visible to readers.
-    pub fn commit_write_batch<K: StoreKey + Send + Sync + 'static, const FAMILIES: usize>(
+    pub fn commit_write_batch<K: StoreKey + Send + Sync + 'static>(
         &self,
         mut write_batch: WriteBatch<K, S, FAMILIES>,
     ) -> Result<()> {
@@ -456,7 +464,27 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
             new_sst_files,
             new_blob_files,
             keys_written,
-        } = write_batch.finish()?;
+        } = write_batch.finish(|family| {
+            let inner = self.inner.read();
+            let set = &inner.accessed_key_hashes[family as usize];
+            // len is only a snapshot at that time and it can change while we create the filter.
+            // So we give it 5% more space to make resizes less likely.
+            let initial_capacity = set.len() * 19 / 20;
+            let mut amqf =
+                qfilter::Filter::with_fingerprint_size(initial_capacity as u64, u64::BITS as u8)
+                    .unwrap();
+            // This drains items from the set. But due to concurrency it might not be empty
+            // afterwards, but that's fine. It will be part of the next commit.
+            set.retain(|hash| {
+                // Performance-wise it would usually be better to insert sorted fingerprints, but we
+                // assume that hashes are equally distributed, which makes it unnecessary.
+                // Good for cache locality is that we insert in the order of the dashset's buckets.
+                amqf.insert_fingerprint(false, *hash)
+                    .expect("Failed to insert fingerprint");
+                false
+            });
+            amqf
+        })?;
         self.commit(CommitOptions {
             new_meta_files,
             new_sst_files,
@@ -1178,24 +1206,27 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
                     self.stats.miss_amqf.fetch_add(1, Ordering::Relaxed);
                 }
                 MetaLookupResult::SstLookup(result) => match result {
-                    SstLookupResult::Found(result) => match result {
-                        LookupValue::Deleted => {
-                            #[cfg(feature = "stats")]
-                            self.stats.hits_deleted.fetch_add(1, Ordering::Relaxed);
-                            return Ok(None);
-                        }
-                        LookupValue::Slice { value } => {
-                            #[cfg(feature = "stats")]
-                            self.stats.hits_small.fetch_add(1, Ordering::Relaxed);
-                            return Ok(Some(value));
-                        }
-                        LookupValue::Blob { sequence_number } => {
-                            #[cfg(feature = "stats")]
-                            self.stats.hits_blob.fetch_add(1, Ordering::Relaxed);
-                            let blob = self.read_blob(sequence_number)?;
-                            return Ok(Some(blob));
+                    SstLookupResult::Found(result) => {
+                        inner.accessed_key_hashes[family].insert(hash);
+                        match result {
+                            LookupValue::Deleted => {
+                                #[cfg(feature = "stats")]
+                                self.stats.hits_deleted.fetch_add(1, Ordering::Relaxed);
+                                return Ok(None);
+                            }
+                            LookupValue::Slice { value } => {
+                                #[cfg(feature = "stats")]
+                                self.stats.hits_small.fetch_add(1, Ordering::Relaxed);
+                                return Ok(Some(value));
+                            }
+                            LookupValue::Blob { sequence_number } => {
+                                #[cfg(feature = "stats")]
+                                self.stats.hits_blob.fetch_add(1, Ordering::Relaxed);
+                                let blob = self.read_blob(sequence_number)?;
+                                return Ok(Some(blob));
+                            }
                         }
-                    },
+                    }
                     SstLookupResult::NotFound => {
                         #[cfg(feature = "stats")]
                         self.stats.miss_key.fetch_add(1, Ordering::Relaxed);

turbopack/crates/turbo-persistence/src/db.rs

@@ -183,6 +183,12 @@ pub struct MetaFile {
     obsolete_entries: Vec<u32>,
     /// The obsolete SST files.
     obsolete_sst_files: Vec<u32>,
+    /// The offset of the start of the "used keys" AMQF data in the meta file relative to the end
+    /// of the header.
+    start_of_used_keys_amqf_data_offset: u32,
+    /// The offset of the end of the "used keys" AMQF data in the the meta file relative to the end
+    /// of the header.
+    end_of_used_keys_amqf_data_offset: u32,
     /// The memory mapped file.
     mmap: Mmap,
 }
@@ -232,6 +238,9 @@ impl MetaFile {
             start_of_amqf_data_offset = entry.end_of_amqf_data_offset;
             entries.push(entry);
         }
+        let start_of_used_keys_amqf_data_offset = start_of_amqf_data_offset;
+        let end_of_used_keys_amqf_data_offset = file.read_u32::<BE>()?;
+
         let offset = file.stream_position()?;
         let file = file.into_inner();
         let mut options = MmapOptions::new();
@@ -246,6 +255,8 @@ impl MetaFile {
             entries,
             obsolete_entries: Vec::new(),
             obsolete_sst_files,
+            start_of_used_keys_amqf_data_offset,
+            end_of_used_keys_amqf_data_offset,
             mmap,
         };
         Ok(file)

turbopack/crates/turbo-persistence/src/meta_file.rs

@@ -6,6 +6,7 @@ use std::{
 
 use anyhow::{Context, Result};
 use byteorder::{BE, WriteBytesExt};
+use qfilter::Filter;
 
 use crate::static_sorted_file_builder::StaticSortedFileBuilderMeta;
 
@@ -15,6 +16,8 @@ pub struct MetaFileBuilder<'a> {
     entries: Vec<(u32, StaticSortedFileBuilderMeta<'a>)>,
     /// Obsolete SST files, represented by their sequence numbers
     obsolete_sst_files: Vec<u32>,
+    /// Optional AMQF for used key hashes
+    used_key_hashes_amqf: Option<Filter>,
 }
 
 impl<'a> MetaFileBuilder<'a> {
@@ -23,6 +26,7 @@ impl<'a> MetaFileBuilder<'a> {
             family,
             entries: Vec::new(),
             obsolete_sst_files: Vec::new(),
+            used_key_hashes_amqf: None,
         }
     }
 
@@ -34,6 +38,10 @@ impl<'a> MetaFileBuilder<'a> {
         self.obsolete_sst_files.push(sequence_number);
     }
 
+    pub fn set_used_key_hashes_amqf(&mut self, amqf: Filter) {
+        self.used_key_hashes_amqf = Some(amqf);
+    }
+
     #[tracing::instrument(level = "trace", skip_all)]
     pub fn write(self, db_path: &Path, seq: u32) -> Result<File> {
         let file = db_path.join(format!("{seq:08}.meta"));
@@ -65,10 +73,22 @@ impl<'a> MetaFileBuilder<'a> {
             amqf_offset += sst.amqf.len();
             file.write_u32::<BE>(amqf_offset as u32)?;
         }
+        let serialized_used_key_hashes = self
+            .used_key_hashes_amqf
+            .as_ref()
+            .map(|f| pot::to_vec(f).expect("AMQF serialization failed"));
+        amqf_offset += serialized_used_key_hashes
+            .as_ref()
+            .map(|bytes| bytes.len())
+            .unwrap_or(0);
+        file.write_u32::<BE>(amqf_offset as u32)?;
 
         for (_, sst) in &self.entries {
             file.write_all(&sst.amqf)?;
         }
+        if let Some(bytes) = &serialized_used_key_hashes {
+            file.write_all(bytes)?;
+        }
         Ok(file.into_inner()?)
     }
 }