Turbopack: Lazy decompress medium value compressed blocks (#82257)

### What?

During compaction we can lazily recompress medium-sized values to avoid the large temporary memory usage.

Author: sokra

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

@@ -0,0 +1,56 @@
+use std::{mem::MaybeUninit, sync::Arc};
+
+use anyhow::{Context, Result};
+use lzzzz::lz4::{ACC_LEVEL_DEFAULT, decompress, decompress_with_dict};
+
+#[tracing::instrument(level = "trace", skip_all)]
+pub fn decompress_into_arc(
+    uncompressed_length: u32,
+    block: &[u8],
+    compression_dictionary: Option<&[u8]>,
+    _long_term: bool,
+) -> 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<[MaybeUninit<u8>]> = Arc::new_zeroed_slice(uncompressed_length as usize);
+    // Assume that the buffer is initialized.
+    let buffer = Arc::into_raw(buffer);
+    // Safety: Assuming that the buffer is initialized is safe because we just created it as
+    // zeroed slice and u8 doesn't require initialization.
+    let mut buffer = unsafe { Arc::from_raw(buffer as *mut [u8]) };
+    // Safety: We know that the buffer is not shared yet.
+    let decompressed = unsafe { Arc::get_mut_unchecked(&mut buffer) };
+    let bytes_writes = if let Some(dict) = compression_dictionary {
+        // Safety: decompress_with_dict will only write to `decompressed` and not read from it.
+        decompress_with_dict(block, decompressed, dict)?
+    } else {
+        // Safety: decompress will only write to `decompressed` and not read from it.
+        decompress(block, decompressed)?
+    };
+    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)
+}
+
+#[tracing::instrument(level = "trace", skip_all)]
+pub fn compress_into_buffer(
+    block: &[u8],
+    dict: Option<&[u8]>,
+    _long_term: bool,
+    buffer: &mut Vec<u8>,
+) -> Result<()> {
+    let mut compressor = if let Some(dict) = dict {
+        lzzzz::lz4::Compressor::with_dict(dict)
+    } else {
+        lzzzz::lz4::Compressor::new()
+    }
+    .context("LZ4 compressor creation failed")?;
+    let acc_factor = ACC_LEVEL_DEFAULT;
+    compressor
+        .next_to_vec(block, buffer, acc_factor)
+        .context("Compression failed")?;
+    Ok(())
+}

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

@@ -3,19 +3,15 @@ use std::{
     collections::HashSet,
     fs::{self, File, OpenOptions, ReadDir},
     io::{BufWriter, Write},
-    mem::{MaybeUninit, swap, transmute},
+    mem::swap,
     ops::RangeInclusive,
     path::{Path, PathBuf},
-    sync::{
-        Arc,
-        atomic::{AtomicBool, AtomicU32, Ordering},
-    },
+    sync::atomic::{AtomicBool, AtomicU32, Ordering},
 };
 
 use anyhow::{Context, Result, bail};
 use byteorder::{BE, ReadBytesExt, WriteBytesExt};
 use jiff::Timestamp;
-use lzzzz::lz4::decompress;
 use memmap2::Mmap;
 use parking_lot::{Mutex, RwLock};
 
@@ -24,6 +20,7 @@ use crate::{
     QueryKey,
     arc_slice::ArcSlice,
     compaction::selector::{Compactable, compute_metrics, get_merge_segments},
+    compression::decompress_into_arc,
     constants::{
         AMQF_AVG_SIZE, AMQF_CACHE_SIZE, DATA_THRESHOLD_PER_COMPACTED_FILE, KEY_BLOCK_AVG_SIZE,
         KEY_BLOCK_CACHE_SIZE, MAX_ENTRIES_PER_COMPACTED_FILE, VALUE_BLOCK_AVG_SIZE,
@@ -392,14 +389,9 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
         #[cfg(target_os = "linux")]
         mmap.advise(memmap2::Advice::Unmergeable)?;
         let mut compressed = &mmap[..];
-        let uncompressed_length = compressed.read_u32::<BE>()? as usize;
-
-        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(compressed, decompressed)?;
+        let uncompressed_length = compressed.read_u32::<BE>()?;
+
+        let buffer = decompress_into_arc(uncompressed_length, compressed, None, true)?;
         Ok(ArcSlice::from(buffer))
     }
 
@@ -918,9 +910,9 @@ impl<S: ParallelScheduler> TurboPersistence<S> {
                                 });
                             }
 
-                            fn create_sst_file<S: ParallelScheduler>(
+                            fn create_sst_file<'l, S: ParallelScheduler>(
                                 parallel_scheduler: &S,
-                                entries: &[LookupEntry],
+                                entries: &[LookupEntry<'l>],
                                 total_key_size: usize,
                                 total_value_size: usize,
                                 path: &Path,
@@ -964,7 +956,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);
@@ -975,7 +967,7 @@ 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;
 
@@ -1023,7 +1015,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/lib.rs

@@ -8,6 +8,7 @@ mod arc_slice;
 mod collector;
 mod collector_entry;
 mod compaction;
+mod compression;
 mod constants;
 mod db;
 mod key;

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 {