Repository: vercel/next.js
Pull request: fix(turbo-persistence): Fix `verify_sst_content` feature, improve unit test performance

## Why?
I'm trying to run https://github.com/boshen/cargo-shear over our workspace to remove unused dependencies. It runs `cargo expand` with `--all-features`, which was breaking on this crate.

## Changes Summary

- Add `opt-level = 1` in dev for this crate.
- Reduce the iterations used in tests. If we really want tons of iterations, we should use benchmarks or write a fuzzer utility. Unit tests should be small and fast.
- Increase the test for blobs to use a 65KiB blob size and assert that's larger than the blob threshold.
- Cleanup: Use the more-correct `u32` type for the `family` field in a few more places instead of `usize` and make the numeric representations of the `KeySpace` enum a bit more explicit.

## Prior to these changes
- Building with `--features turbo-persistence/verify_sst_content` would fail with compilation errors
- Unit tests would take multiple minutes to run
- The blob test didn't appear to actually test blobs (10KiB < 64KiB)

## After these changes
- The `verify_sst_content` feature works and passes unit tests
- Unit tests finish in about a minute with `cargo test` and `cargo nextest`, this is still slow, but much better than it was.
- Unit tests actually test blobs as intended.

![Screenshot 2025-04-21 at 2.12.32 PM.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/HAZVitxRNnZz8QMiPn4a/ff7aa104-c193-4514-9fde-1a6fcc11acf8.png)

## Potential future work
It would be useful to be able to adjust the medium and large (blob) size thresholds to allow us to lower those thresholds in tests.

Closes PACK-4402

Changed files:
filename: Cargo.toml
patch: @@ -22,7 +22,7 @@ exclude = [
 too_many_arguments = "allow"
 
 # This crate is particularly sensitive to compiler optimizations
-[profile.dev.package.turbo-tasks-memory]
+[profile.dev.package.turbo-persistence]
 opt-level = 1
 
 # Set the options for dependencies (not crates in the workspace), this mostly impacts cold builds

filename: turbopack/crates/turbo-persistence/src/tests.rs
patch: @@ -3,7 +3,7 @@ use std::time::Instant;
 use anyhow::Result;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
-use crate::{db::TurboPersistence, write_batch::WriteBatch};
+use crate::{constants::MAX_MEDIUM_VALUE_SIZE, db::TurboPersistence, write_batch::WriteBatch};
 
 #[test]
 fn full_cycle() -> Result<()> {
@@ -53,7 +53,7 @@ fn full_cycle() -> Result<()> {
         "Families",
         |batch| {
             for i in 0..16u8 {
-                batch.put(i as usize, vec![i], vec![i].into())?;
+                batch.put(u32::from(i), vec![i], vec![i].into())?;
             }
             Ok(())
         },
@@ -89,41 +89,46 @@ fn full_cycle() -> Result<()> {
         },
     );
 
+    const BLOB_SIZE: usize = 65 * 1024 * 1024;
+    #[expect(clippy::assertions_on_constants)]
+    {
+        assert!(BLOB_SIZE > MAX_MEDIUM_VALUE_SIZE);
+    }
     test_case(
         &mut test_cases,
         "Large keys and values (blob files)",
         |batch| {
-            for i in 0..20u8 {
-                batch.put(
-                    0,
-                    vec![i; 10 * 1024 * 1024],
-                    vec![i; 10 * 1024 * 1024].into(),
-                )?;
+            for i in 0..2u8 {
+                batch.put(0, vec![i; BLOB_SIZE], vec![i; BLOB_SIZE].into())?;
             }
             Ok(())
         },
         |db| {
-            for i in 0..20u8 {
-                let Some(value) = db.get(0, &vec![i; 10 * 1024 * 1024])? else {
+            for i in 0..2u8 {
+                let key_and_value = vec![i; BLOB_SIZE];
+                let Some(value) = db.get(0, &key_and_value)? else {
                     panic!("Value not found");
                 };
-                assert_eq!(&*value, &vec![i; 10 * 1024 * 1024]);
+                assert_eq!(&*value, &key_and_value);
             }
             Ok(())
         },
     );
 
+    fn different_sizes_range() -> impl Iterator<Item = u8> {
+        (10..20).map(|value| value * 10)
+    }
     test_case(
         &mut test_cases,
         "Different sizes keys and values",
         |batch| {
-            for i in 100..200u8 {
+            for i in different_sizes_range() {
                 batch.put(0, vec![i; i as usize], vec![i; i as usize].into())?;
             }
             Ok(())
         },
         |db| {
-            for i in 100..200u8 {
+            for i in different_sizes_range() {
                 let Some(value) = db.get(0, &vec![i; i as usize])? else {
                     panic!("Value not found");
                 };
@@ -351,18 +356,20 @@ fn persist_changes() -> Result<()> {
     let tempdir = tempfile::tempdir()?;
     let path = tempdir.path();
 
+    const READ_COUNT: u32 = 2_000; // we'll read every 10th value, so writes are 10x this value
     fn put(b: &WriteBatch<(u8, [u8; 4]), 1>, key: u8, value: u8) -> Result<()> {
-        for i in 0..2000000u32 {
+        for i in 0..(READ_COUNT * 10) {
             b.put(0, (key, i.to_be_bytes()), vec![value].into())?;
         }
         Ok(())
     }
     fn check(db: &TurboPersistence, key: u8, value: u8) -> Result<()> {
-        for i in 0..200000u32 {
+        for i in 0..READ_COUNT {
+            // read every 10th item
             let i = i * 10;
             assert_eq!(
                 db.get(0, &(key, i.to_be_bytes()))?.as_deref(),
-                Some(&[value][..])
+                Some(&[value][..]),
             );
         }
         Ok(())

filename: turbopack/crates/turbo-persistence/src/write_batch.rs
patch: @@ -22,7 +22,11 @@ use crate::{
     key::StoreKey, static_sorted_file_builder::StaticSortedFileBuilder, ValueBuffer,
 };
 
-/// The thread local state of a `WriteBatch`.
+/// The thread local state of a `WriteBatch`. `FAMILIES` should fit within a `u32`.
+//
+// NOTE: This type *must* use `usize`, even though the real type used in storage is `u32` because
+// there's no way to cast a `u32` to `usize` when declaring an array without the nightly
+// `min_generic_const_args` feature.
 struct ThreadLocalState<K: StoreKey + Send, const FAMILIES: usize> {
     /// The collectors for each family.
     collectors: [Option<Collector<K>>; FAMILIES],
@@ -54,7 +58,9 @@ pub struct WriteBatch<K: StoreKey + Send, const FAMILIES: usize> {
 impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     /// Creates a new write batch for a database.
     pub(crate) fn new(path: PathBuf, current: u32) -> Self {
-        assert!(FAMILIES <= u32::MAX as usize);
+        const {
+            assert!(FAMILIES <= usize_from_u32(u32::MAX));
+        };
         Self {
             path,
             current_sequence_number: AtomicU32::new(current),
@@ -88,10 +94,11 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     fn collector_mut<'l>(
         &self,
         state: &'l mut ThreadLocalState<K, FAMILIES>,
-        family: usize,
+        family: u32,
     ) -> Result<&'l mut Collector<K>> {
-        debug_assert!(family < FAMILIES);
-        let collector = state.collectors[family].get_or_insert_with(|| {
+        let family_idx = usize_from_u32(family);
+        debug_assert!(family_idx < FAMILIES);
+        let collector = state.collectors[family_idx].get_or_insert_with(|| {
             self.idle_collectors
                 .lock()
                 .pop()
@@ -106,7 +113,7 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     }
 
     /// Puts a key-value pair into the write batch.
-    pub fn put(&self, family: usize, key: K, value: ValueBuffer<'_>) -> Result<()> {
+    pub fn put(&self, family: u32, key: K, value: ValueBuffer<'_>) -> Result<()> {
         let state = self.thread_local_state();
         let collector = self.collector_mut(state, family)?;
         if value.len() <= MAX_MEDIUM_VALUE_SIZE {
@@ -120,7 +127,7 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     }
 
     /// Puts a delete operation into the write batch.
-    pub fn delete(&self, family: usize, key: K) -> Result<()> {
+    pub fn delete(&self, family: u32, key: K) -> Result<()> {
         let state = self.thread_local_state();
         let collector = self.collector_mut(state, family)?;
         collector.delete(key);
@@ -151,7 +158,7 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
             fn handle_done_collector<'scope, K: StoreKey + Send + Sync, const FAMILIES: usize>(
                 this: &'scope WriteBatch<K, FAMILIES>,
                 scope: &Scope<'scope>,
-                family: usize,
+                family: u32,
                 mut collector: Collector<K>,
                 shared_new_sst_files: &'scope Mutex<&mut Vec<(u32, File)>>,
                 shared_error: &'scope Mutex<Result<()>>,
@@ -173,7 +180,8 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
             all_collectors
                 .into_par_iter()
                 .enumerate()
-                .for_each(|(family, collectors)| {
+                .for_each(|(family_idx, collectors)| {
+                    let family = u32::try_from(family_idx).unwrap();
                     let final_collector = collectors.into_par_iter().reduce(
                         || None,
                         |a, b| match (a, b) {
@@ -250,14 +258,14 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     /// Creates a new SST file with the given collector data.
     fn create_sst_file(
         &self,
-        family: usize,
+        family: u32,
         collector_data: (&[CollectorEntry<K>], usize, usize),
     ) -> Result<(u32, File)> {
         let (entries, total_key_size, total_value_size) = collector_data;
         let seq = self.current_sequence_number.fetch_add(1, Ordering::SeqCst) + 1;
 
         let builder =
-            StaticSortedFileBuilder::new(family as u32, entries, total_key_size, total_value_size)?;
+            StaticSortedFileBuilder::new(family, entries, total_key_size, total_value_size)?;
 
         let path = self.path.join(format!("{:08}.sst", seq));
         let file = builder
@@ -272,6 +280,7 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
                 collector_entry::CollectorEntryValue,
                 key::hash_key,
                 static_sorted_file::{AqmfCache, BlockCache, LookupResult, StaticSortedFile},
+                static_sorted_file_builder::Entry,
             };
 
             file.sync_all()?;
@@ -297,24 +306,33 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
                 Default::default(),
                 Default::default(),
             );
+            let mut key_buf = Vec::new();
             for entry in entries {
-                let mut key = Vec::with_capacity(entry.key.len());
-                entry.key.write_to(&mut key);
+                entry.write_key_to(&mut key_buf);
                 let result = sst
-                    .lookup(hash_key(&key), &key, &cache1, &cache2, &cache3)
+                    .lookup(
+                        family,
+                        hash_key(&key_buf),
+                        &key_buf,
+                        &cache1,
+                        &cache2,
+                        &cache3,
+                    )
                     .expect("key found");
+                key_buf.clear();
                 match result {
                     LookupResult::Deleted => {}
-                    LookupResult::Small { value: val } => {
-                        if let EntryValue::Small { value } | EntryValue::Medium { value } =
-                            entry.value
-                        {
-                            assert_eq!(&*val, &*value);
-                        } else {
-                            panic!("Unexpected value");
-                        }
+                    LookupResult::Slice {
+                        value: lookup_value,
+                    } => {
+                        let expected_value_slice = match &entry.value {
+                            CollectorEntryValue::Small { value } => &**value,
+                            CollectorEntryValue::Medium { value } => &**value,
+                            _ => panic!("Unexpected value"),
+                        };
+                        assert_eq!(*lookup_value, *expected_value_slice);
                     }
-                    LookupResult::Blob { sequence_number } => {}
+                    LookupResult::Blob { sequence_number: _ } => {}
                     LookupResult::QuickFilterMiss => panic!("aqmf must include"),
                     LookupResult::RangeMiss => panic!("Index must cover"),
                     LookupResult::KeyMiss => panic!("All keys must exist"),
@@ -325,3 +343,13 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
         Ok((seq, file))
     }
 }
+
+#[inline(always)]
+const fn usize_from_u32(value: u32) -> usize {
+    // This should always be true, as we assume at least a 32-bit width architecture for Turbopack.
+    // Since this is a const expression, we expect it to be compiled away.
+    const {
+        assert!(u32::BITS < usize::BITS);
+    };
+    value as usize
+}

filename: turbopack/crates/turbo-tasks-backend/src/database/key_value_database.rs
patch: @@ -6,11 +6,11 @@ use crate::database::write_batch::{
 
 #[derive(Debug, Clone, Copy)]
 pub enum KeySpace {
-    Infra,
-    TaskMeta,
-    TaskData,
-    ForwardTaskCache,
-    ReverseTaskCache,
+    Infra = 0,
+    TaskMeta = 1,
+    TaskData = 2,
+    ForwardTaskCache = 3,
+    ReverseTaskCache = 4,
 }
 
 pub trait KeyValueDatabase {

filename: turbopack/crates/turbo-tasks-backend/src/database/turbo.rs
patch: @@ -138,11 +138,11 @@ impl<'a> BaseWriteBatch<'a> for TurboWriteBatch<'a> {
 impl<'a> ConcurrentWriteBatch<'a> for TurboWriteBatch<'a> {
     fn put(&self, key_space: KeySpace, key: WriteBuffer<'_>, value: WriteBuffer<'_>) -> Result<()> {
         self.batch
-            .put(key_space as usize, key.into_static(), value.into())
+            .put(key_space as u32, key.into_static(), value.into())
     }
 
     fn delete(&self, key_space: KeySpace, key: WriteBuffer<'_>) -> Result<()> {
-        self.batch.delete(key_space as usize, key.into_static())
+        self.batch.delete(key_space as u32, key.into_static())
     }
 }