Turbopack: Decompress medium values lazily to reduce memory during compaction

Author: sokra

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

@@ -915,8 +915,8 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
                                     });
                                 }
 
-                                fn create_sst_file(
-                                    entries: &[LookupEntry],
+                                fn create_sst_file<'l>(
+                                    entries: &[LookupEntry<'l>],
                                     total_key_size: usize,
                                     total_value_size: usize,
                                     path: &Path,
@@ -959,7 +959,7 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
 
                                 let mut total_key_size = 0;
                                 let mut total_value_size = 0;
-                                let mut current: Option<LookupEntry> = None;
+                                let mut current: Option<LookupEntry<'_>> = None;
                                 let mut entries = Vec::new();
                                 let mut last_entries = Vec::new();
                                 let mut last_entries_total_sizes = (0, 0);
@@ -970,7 +970,8 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
                                     if let Some(current) = current.take() {
                                         if current.key != entry.key {
                                             let key_size = current.key.len();
-                                            let value_size = current.value.size_in_sst();
+                                            let value_size =
+                                                current.value.uncompressed_size_in_sst();
                                             total_key_size += key_size;
                                             total_value_size += value_size;
 
@@ -1017,7 +1018,7 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
                                 }
                                 if let Some(entry) = current {
                                     total_key_size += entry.key.len();
-                                    total_value_size += entry.value.size_in_sst();
+                                    total_value_size += entry.value.uncompressed_size_in_sst();
                                     entries.push(entry);
                                 }
 

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

@@ -14,28 +14,43 @@ pub enum LookupValue {
     Blob { sequence_number: u32 },
 }
 
-impl LookupValue {
+/// A value from a SST file lookup.
+pub enum LazyLookupValue<'l> {
+    /// A LookupValue
+    Eager(LookupValue),
+    /// A medium sized value that is still compressed.
+    Medium {
+        uncompressed_size: u32,
+        block: &'l [u8],
+        dictionary: &'l [u8],
+    },
+}
+
+impl LazyLookupValue<'_> {
     /// Returns the size of the value in the SST file.
-    pub fn size_in_sst(&self) -> usize {
+    pub fn uncompressed_size_in_sst(&self) -> usize {
         match self {
-            LookupValue::Slice { value } => value.len(),
-            LookupValue::Deleted => 0,
-            LookupValue::Blob { .. } => 0,
+            LazyLookupValue::Eager(LookupValue::Slice { value }) => value.len(),
+            LazyLookupValue::Eager(LookupValue::Deleted) => 0,
+            LazyLookupValue::Eager(LookupValue::Blob { .. }) => 0,
+            LazyLookupValue::Medium {
+                uncompressed_size, ..
+            } => *uncompressed_size as usize,
         }
     }
 }
 
 /// An entry from a SST file lookup.
-pub struct LookupEntry {
+pub struct LookupEntry<'l> {
     /// The hash of the key.
     pub hash: u64,
     /// The key.
     pub key: ArcSlice<u8>,
     /// The value.
-    pub value: LookupValue,
+    pub value: LazyLookupValue<'l>,
 }
 
-impl Entry for LookupEntry {
+impl Entry for LookupEntry<'_> {
     fn key_hash(&self) -> u64 {
         self.hash
     }
@@ -50,17 +65,26 @@ impl Entry for LookupEntry {
 
     fn value(&self) -> EntryValue<'_> {
         match &self.value {
-            LookupValue::Deleted => EntryValue::Deleted,
-            LookupValue::Slice { value } => {
+            LazyLookupValue::Eager(LookupValue::Deleted) => EntryValue::Deleted,
+            LazyLookupValue::Eager(LookupValue::Slice { value }) => {
                 if value.len() > MAX_SMALL_VALUE_SIZE {
                     EntryValue::Medium { value }
                 } else {
                     EntryValue::Small { value }
                 }
             }
-            LookupValue::Blob { sequence_number } => EntryValue::Large {
+            LazyLookupValue::Eager(LookupValue::Blob { sequence_number }) => EntryValue::Large {
                 blob: *sequence_number,
             },
+            LazyLookupValue::Medium {
+                uncompressed_size,
+                block,
+                dictionary,
+            } => EntryValue::MediumCompressed {
+                uncompressed_size: *uncompressed_size,
+                block,
+                dictionary,
+            },
         }
     }
 }

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

@@ -6,27 +6,27 @@ use crate::lookup_entry::LookupEntry;
 
 /// An active iterator that is being merged. It has peeked the next element and can be compared
 /// according to that element. The `order` is used when multiple iterators have the same key.
-struct ActiveIterator<T: Iterator<Item = Result<LookupEntry>>> {
+struct ActiveIterator<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> {
     iter: T,
     order: usize,
-    entry: LookupEntry,
+    entry: LookupEntry<'l>,
 }
 
-impl<T: Iterator<Item = Result<LookupEntry>>> PartialEq for ActiveIterator<T> {
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> PartialEq for ActiveIterator<'l, T> {
     fn eq(&self, other: &Self) -> bool {
         self.entry.hash == other.entry.hash && *self.entry.key == *other.entry.key
     }
 }
 
-impl<T: Iterator<Item = Result<LookupEntry>>> Eq for ActiveIterator<T> {}
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> Eq for ActiveIterator<'l, T> {}
 
-impl<T: Iterator<Item = Result<LookupEntry>>> PartialOrd for ActiveIterator<T> {
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> PartialOrd for ActiveIterator<'l, T> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl<T: Iterator<Item = Result<LookupEntry>>> Ord for ActiveIterator<T> {
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> Ord for ActiveIterator<'l, T> {
     fn cmp(&self, other: &Self) -> Ordering {
         self.entry
             .hash
@@ -39,11 +39,11 @@ impl<T: Iterator<Item = Result<LookupEntry>>> Ord for ActiveIterator<T> {
 
 /// An iterator that merges multiple sorted iterators into a single sorted iterator. Internal it
 /// uses an heap of iterators to iterate them in order.
-pub struct MergeIter<T: Iterator<Item = Result<LookupEntry>>> {
-    heap: BinaryHeap<ActiveIterator<T>>,
+pub struct MergeIter<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> {
+    heap: BinaryHeap<ActiveIterator<'l, T>>,
 }
 
-impl<T: Iterator<Item = Result<LookupEntry>>> MergeIter<T> {
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> MergeIter<'l, T> {
     pub fn new(iters: impl Iterator<Item = T>) -> Result<Self> {
         let mut heap = BinaryHeap::new();
         for (order, mut iter) in iters.enumerate() {
@@ -56,8 +56,8 @@ impl<T: Iterator<Item = Result<LookupEntry>>> MergeIter<T> {
     }
 }
 
-impl<T: Iterator<Item = Result<LookupEntry>>> Iterator for MergeIter<T> {
-    type Item = Result<LookupEntry>;
+impl<'l, T: Iterator<Item = Result<LookupEntry<'l>>>> Iterator for MergeIter<'l, T> {
+    type Item = Result<LookupEntry<'l>>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let ActiveIterator {

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

@@ -18,7 +18,7 @@ use rustc_hash::FxHasher;
 use crate::{
     QueryKey,
     arc_slice::ArcSlice,
-    lookup_entry::{LookupEntry, LookupValue},
+    lookup_entry::{LazyLookupValue, LookupEntry, LookupValue},
 };
 
 /// The block header for an index block.
@@ -339,6 +339,14 @@ impl StaticSortedFile {
 
     /// Reads a block from the file.
     fn read_block(&self, block_index: u16, compression_dictionary: &[u8]) -> Result<ArcSlice<u8>> {
+        let (uncompressed_length, block) = self.get_compressed_block(block_index)?;
+
+        let buffer = decompress_into_arc(uncompressed_length, block, compression_dictionary)?;
+        Ok(ArcSlice::from(buffer))
+    }
+
+    /// Gets the slice of the compressed block from the memory mapped file.
+    fn get_compressed_block(&self, block_index: u16) -> Result<(u32, &[u8])> {
         #[cfg(feature = "strict_checks")]
         if block_index >= self.meta.block_count {
             bail!(
@@ -386,17 +394,9 @@ impl StaticSortedFile {
                 self.meta.blocks_start()
             );
         }
-        let uncompressed_length =
-            (&self.mmap[block_start..block_start + 4]).read_u32::<BE>()? as usize;
-        let block = self.mmap[block_start + 4..block_end].to_vec();
-
-        let buffer = Arc::new_zeroed_slice(uncompressed_length);
-        // Safety: MaybeUninit<u8> can be safely transmuted to u8.
-        let mut buffer = unsafe { transmute::<Arc<[MaybeUninit<u8>]>, Arc<[u8]>>(buffer) };
-        // Safety: We know that the buffer is not shared yet.
-        let decompressed = unsafe { Arc::get_mut_unchecked(&mut buffer) };
-        decompress_with_dict(&block, decompressed, compression_dictionary)?;
-        Ok(ArcSlice::from(buffer))
+        let uncompressed_length = (&self.mmap[block_start..block_start + 4]).read_u32::<BE>()?;
+        let block = &self.mmap[block_start + 4..block_end];
+        Ok((uncompressed_length, block))
     }
 }
 
@@ -423,15 +423,15 @@ struct CurrentIndexBlock {
     index: usize,
 }
 
-impl Iterator for StaticSortedFileIter<'_> {
-    type Item = Result<LookupEntry>;
+impl<'l> Iterator for StaticSortedFileIter<'l> {
+    type Item = Result<LookupEntry<'l>>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.next_internal().transpose()
     }
 }
 
-impl StaticSortedFileIter<'_> {
+impl<'l> StaticSortedFileIter<'l> {
     /// Enters a block at the given index.
     fn enter_block(&mut self, block_index: u16) -> Result<()> {
         let block_arc = self.this.get_key_block(block_index, self.key_block_cache)?;
@@ -468,7 +468,7 @@ impl StaticSortedFileIter<'_> {
     }
 
     /// Gets the next entry in the file and moves the cursor.
-    fn next_internal(&mut self) -> Result<Option<LookupEntry>> {
+    fn next_internal(&mut self) -> Result<Option<LookupEntry<'l>>> {
         loop {
             if let Some(CurrentKeyBlock {
                 offsets,
@@ -479,9 +479,22 @@ impl StaticSortedFileIter<'_> {
             {
                 let GetKeyEntryResult { hash, key, ty, val } =
                     get_key_entry(&offsets, &entries, entry_count, index)?;
-                let value = self
-                    .this
-                    .handle_key_match(ty, val, self.value_block_cache)?;
+                let value = if ty == KEY_BLOCK_ENTRY_TYPE_MEDIUM {
+                    let mut val = val;
+                    let block = val.read_u16::<BE>()?;
+                    let (uncompressed_size, block) = self.this.get_compressed_block(block)?;
+                    LazyLookupValue::Medium {
+                        uncompressed_size,
+                        block,
+                        dictionary: &self.this.mmap
+                            [self.this.meta.value_compression_dictionary_range()],
+                    }
+                } else {
+                    let value = self
+                        .this
+                        .handle_key_match(ty, val, self.value_block_cache)?;
+                    LazyLookupValue::Eager(value)
+                };
                 let entry = LookupEntry {
                     hash,
                     // Safety: The key is a valid slice of the entries.
@@ -573,3 +586,27 @@ fn get_key_entry<'l>(
         }
     })
 }
+
+pub fn decompress_into_arc(
+    uncompressed_length: u32,
+    block: &[u8],
+    compression_dictionary: &[u8],
+) -> Result<Arc<[u8]>> {
+    // We directly allocate the buffer in an Arc to avoid copying it into an Arc and avoiding
+    // double indirection. This is a dynamically sized arc.
+    let buffer = Arc::new_zeroed_slice(uncompressed_length as usize);
+    // Safety: MaybeUninit<u8> can be safely transmuted to u8.
+    // Safety: decompress_with_dict will only write to `buffer` and not read from it.
+    // Safety: We check that decompress_with_dict will write all bytes.
+    let mut buffer = unsafe { transmute::<Arc<[MaybeUninit<u8>]>, Arc<[u8]>>(buffer) };
+    // Safety: We know that the buffer is not shared yet.
+    let decompressed = unsafe { Arc::get_mut_unchecked(&mut buffer) };
+    // Safety: decompress_with_dict will only write to `decompressed` and not read from it.
+    let bytes_writes = decompress_with_dict(&block, decompressed, compression_dictionary)?;
+    assert_eq!(
+        bytes_writes, uncompressed_length as usize,
+        "Decompressed length does not match expected length"
+    );
+    // Safety: The buffer is now fully initialized and can be used.
+    Ok(buffer)
+}

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

@@ -12,7 +12,7 @@ use lzzzz::lz4::{ACC_LEVEL_DEFAULT, max_compressed_size};
 
 use crate::static_sorted_file::{
     BLOCK_TYPE_INDEX, BLOCK_TYPE_KEY, KEY_BLOCK_ENTRY_TYPE_BLOB, KEY_BLOCK_ENTRY_TYPE_DELETED,
-    KEY_BLOCK_ENTRY_TYPE_MEDIUM, KEY_BLOCK_ENTRY_TYPE_SMALL,
+    KEY_BLOCK_ENTRY_TYPE_MEDIUM, KEY_BLOCK_ENTRY_TYPE_SMALL, decompress_into_arc,
 };
 
 /// The maximum number of entries that should go into a single key block
@@ -68,6 +68,12 @@ pub enum EntryValue<'l> {
     Small { value: &'l [u8] },
     /// Medium-sized value. They are stored in their own value block.
     Medium { value: &'l [u8] },
+    /// Medium-sized value. They are stored in their own value block. Precompressed.
+    MediumCompressed {
+        uncompressed_size: u32,
+        block: &'l [u8],
+        dictionary: &'l [u8],
+    },
     /// Large-sized value. They are stored in a blob file.
     Large { blob: u32 },
     /// Tombstone. The value was removed.
@@ -386,7 +392,18 @@ fn write_value_blocks(
                 value_locations.push((block_index, 0));
                 writer.write_value_block(value, value_compression_dictionary)?;
             }
-            _ => {
+            EntryValue::MediumCompressed {
+                uncompressed_size,
+                block,
+                dictionary,
+            } => {
+                let block_index = writer.next_block_index();
+                value_locations.push((block_index, 0));
+                // Recompress block with a different dictionary
+                let decompressed = decompress_into_arc(uncompressed_size, block, dictionary)?;
+                writer.write_value_block(&decompressed, value_compression_dictionary)?;
+            }
+            EntryValue::Deleted | EntryValue::Large { .. } => {
                 value_locations.push((0, 0));
             }
         }
@@ -438,7 +455,7 @@ fn write_key_blocks_and_compute_amqf(
                     value.len().try_into().unwrap(),
                 );
             }
-            EntryValue::Medium { .. } => {
+            EntryValue::Medium { .. } | EntryValue::MediumCompressed { .. } => {
                 block.put_medium(entry, value_location.0);
             }
             EntryValue::Large { blob } => {