Repository: vercel/next.js
Pull request: Turbopack: use small thread local collector that flushes to global collector

### What?

Storing many large thread local collectors can lead to high memory usage on systems with many cores. It also unnecessarily delays writing out files to disk.

This changes this to use very small thread local collector that flushes to a single global collector (per key family). The small thread local collectors still avoid contention. This also creates more compact files.

Closes PACK-4480

Changed files:
filename: turbopack/crates/turbo-persistence/src/collector.rs
patch: @@ -9,19 +9,19 @@ use crate::{
 
 /// A collector accumulates entries that should be eventually written to a file. It keeps track of
 /// count and size of the entries to decide when it's "full". Accessing the entries sorts them.
-pub struct Collector<K: StoreKey> {
+pub struct Collector<K: StoreKey, const SIZE_SHIFT: usize = 0> {
     total_key_size: usize,
     total_value_size: usize,
     entries: Vec<CollectorEntry<K>>,
 }
 
-impl<K: StoreKey> Collector<K> {
+impl<K: StoreKey, const SIZE_SHIFT: usize> Collector<K, SIZE_SHIFT> {
     /// Creates a new collector. Note that this allocates the full capacity for the entries.
     pub fn new() -> Self {
         Self {
             total_key_size: 0,
             total_value_size: 0,
-            entries: Vec::with_capacity(MAX_ENTRIES_PER_INITIAL_FILE),
+            entries: Vec::with_capacity(MAX_ENTRIES_PER_INITIAL_FILE >> SIZE_SHIFT),
         }
     }
 
@@ -32,8 +32,9 @@ impl<K: StoreKey> Collector<K> {
 
     /// Returns true if the collector is full.
     pub fn is_full(&self) -> bool {
-        self.entries.len() >= MAX_ENTRIES_PER_INITIAL_FILE
-            || self.total_key_size + self.total_value_size > DATA_THRESHOLD_PER_INITIAL_FILE
+        self.entries.len() >= MAX_ENTRIES_PER_INITIAL_FILE >> SIZE_SHIFT
+            || self.total_key_size + self.total_value_size
+                > DATA_THRESHOLD_PER_INITIAL_FILE >> SIZE_SHIFT
     }
 
     /// Adds a normal key-value pair to the collector.
@@ -110,9 +111,4 @@ impl<K: StoreKey> Collector<K> {
         self.total_value_size = 0;
         self.entries.drain(..)
     }
-
-    /// Returns the number of entries in the collector.
-    pub fn len(&self) -> usize {
-        self.entries.len()
-    }
 }

filename: turbopack/crates/turbo-persistence/src/constants.rs
patch: @@ -17,6 +17,10 @@ pub const DATA_THRESHOLD_PER_INITIAL_FILE: usize = 256 * 1024 * 1024;
 /// Finish file when total amount of data exceeds this
 pub const DATA_THRESHOLD_PER_COMPACTED_FILE: usize = 256 * 1024 * 1024;
 
+/// Reduction factor (as bit shift) for the size of the thread-local buffer as shift of
+/// MAX_ENTRIES_PER_INITIAL_FILE and DATA_THRESHOLD_PER_INITIAL_FILE.
+pub const THREAD_LOCAL_SIZE_SHIFT: usize = 7;
+
 /// Maximum RAM bytes for AQMF cache
 pub const AQMF_CACHE_SIZE: u64 = 300 * 1024 * 1024;
 pub const AQMF_AVG_SIZE: usize = 37399;

filename: turbopack/crates/turbo-persistence/src/write_batch.rs
patch: @@ -2,7 +2,7 @@ use std::{
     cell::UnsafeCell,
     fs::File,
     io::Write,
-    mem::{replace, swap},
+    mem::{replace, take},
     path::PathBuf,
     sync::atomic::{AtomicU32, Ordering},
 };
@@ -13,13 +13,18 @@ use lzzzz::lz4::{self, ACC_LEVEL_DEFAULT};
 use parking_lot::Mutex;
 use rayon::{
     iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator},
-    scope, Scope,
+    scope,
 };
+use smallvec::SmallVec;
 use thread_local::ThreadLocal;
 
 use crate::{
-    collector::Collector, collector_entry::CollectorEntry, constants::MAX_MEDIUM_VALUE_SIZE,
-    key::StoreKey, static_sorted_file_builder::StaticSortedFileBuilder, ValueBuffer,
+    collector::Collector,
+    collector_entry::CollectorEntry,
+    constants::{MAX_MEDIUM_VALUE_SIZE, THREAD_LOCAL_SIZE_SHIFT},
+    key::StoreKey,
+    static_sorted_file_builder::StaticSortedFileBuilder,
+    ValueBuffer,
 };
 
 /// The thread local state of a `WriteBatch`. `FAMILIES` should fit within a `u32`.
@@ -29,10 +34,7 @@ use crate::{
 // `min_generic_const_args` feature.
 struct ThreadLocalState<K: StoreKey + Send, const FAMILIES: usize> {
     /// The collectors for each family.
-    collectors: [Option<Collector<K>>; FAMILIES],
-    /// The list of new SST files that have been created.
-    /// Tuple of (sequence number, file).
-    new_sst_files: Vec<(u32, File)>,
+    collectors: [Option<Collector<K, THREAD_LOCAL_SIZE_SHIFT>>; FAMILIES],
     /// The list of new blob files that have been created.
     /// Tuple of (sequence number, file).
     new_blob_files: Vec<(u32, File)>,
@@ -55,8 +57,15 @@ pub struct WriteBatch<K: StoreKey + Send, const FAMILIES: usize> {
     current_sequence_number: AtomicU32,
     /// The thread local state.
     thread_locals: ThreadLocal<UnsafeCell<ThreadLocalState<K, FAMILIES>>>,
-    /// Collectors are are current unused, but have memory preallocated.
+    /// Collectors in use. The thread local collectors flush into these when they are full.
+    collectors: [Mutex<Collector<K>>; FAMILIES],
+    /// The list of new SST files that have been created.
+    /// Tuple of (sequence number, file).
+    new_sst_files: Mutex<Vec<(u32, File)>>,
+    /// Collectors that are currently unused, but have memory preallocated.
     idle_collectors: Mutex<Vec<Collector<K>>>,
+    /// Collectors that are currently unused, but have memory preallocated.
+    idle_thread_local_collectors: Mutex<Vec<Collector<K, THREAD_LOCAL_SIZE_SHIFT>>>,
 }
 
 impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
@@ -69,7 +78,10 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
             path,
             current_sequence_number: AtomicU32::new(current),
             thread_locals: ThreadLocal::new(),
+            collectors: [(); FAMILIES].map(|_| Mutex::new(Collector::new())),
+            new_sst_files: Mutex::new(Vec::new()),
             idle_collectors: Mutex::new(Vec::new()),
+            idle_thread_local_collectors: Mutex::new(Vec::new()),
         }
     }
 
@@ -86,7 +98,6 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
         let cell = self.thread_locals.get_or(|| {
             UnsafeCell::new(ThreadLocalState {
                 collectors: [const { None }; FAMILIES],
-                new_sst_files: Vec::new(),
                 new_blob_files: Vec::new(),
             })
         });
@@ -95,31 +106,59 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     }
 
     /// Returns the collector for a family for the current thread.
-    fn collector_mut<'l>(
+    fn thread_local_collector_mut<'l>(
         &self,
         state: &'l mut ThreadLocalState<K, FAMILIES>,
         family: u32,
-    ) -> Result<&'l mut Collector<K>> {
-        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
+    ) -> Result<&'l mut Collector<K, THREAD_LOCAL_SIZE_SHIFT>> {
+        debug_assert!(usize_from_u32(family) < FAMILIES);
+        let collector = state.collectors[usize_from_u32(family)].get_or_insert_with(|| {
+            self.idle_thread_local_collectors
                 .lock()
                 .pop()
                 .unwrap_or_else(|| Collector::new())
         });
         if collector.is_full() {
-            let sst = self.create_sst_file(family, collector.sorted())?;
-            collector.clear();
-            state.new_sst_files.push(sst);
+            self.flush_thread_local_collector(family, collector)?;
         }
         Ok(collector)
     }
 
+    fn flush_thread_local_collector(
+        &self,
+        family: u32,
+        collector: &mut Collector<K, THREAD_LOCAL_SIZE_SHIFT>,
+    ) -> Result<()> {
+        let mut full_collectors = SmallVec::<[_; 2]>::new();
+        {
+            let mut global_collector = self.collectors[usize_from_u32(family)].lock();
+            for entry in collector.drain() {
+                global_collector.add_entry(entry);
+                if global_collector.is_full() {
+                    full_collectors.push(replace(
+                        &mut *global_collector,
+                        self.idle_collectors
+                            .lock()
+                            .pop()
+                            .unwrap_or_else(|| Collector::new()),
+                    ));
+                }
+            }
+        }
+        for mut global_collector in full_collectors {
+            // When the global collector is full, we create a new SST file.
+            let sst = self.create_sst_file(family, global_collector.sorted())?;
+            global_collector.clear();
+            self.new_sst_files.lock().push(sst);
+            self.idle_collectors.lock().push(global_collector);
+        }
+        Ok(())
+    }
+
     /// Puts a key-value pair into the write batch.
     pub fn put(&self, family: u32, key: K, value: ValueBuffer<'_>) -> Result<()> {
         let state = self.thread_local_state();
-        let collector = self.collector_mut(state, family)?;
+        let collector = self.thread_local_collector_mut(state, family)?;
         if value.len() <= MAX_MEDIUM_VALUE_SIZE {
             collector.put(key, value);
         } else {
@@ -133,106 +172,77 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     /// Puts a delete operation into the write batch.
     pub fn delete(&self, family: u32, key: K) -> Result<()> {
         let state = self.thread_local_state();
-        let collector = self.collector_mut(state, family)?;
+        let collector = self.thread_local_collector_mut(state, family)?;
         collector.delete(key);
         Ok(())
     }
 
     /// Finishes the write batch by returning the new sequence number and the new SST files. This
     /// writes all outstanding thread local data to disk.
     pub(crate) fn finish(&mut self) -> Result<FinishResult> {
-        let mut new_sst_files = Vec::new();
         let mut new_blob_files = Vec::new();
-        let mut all_collectors = [(); FAMILIES].map(|_| Vec::new());
-        for cell in self.thread_locals.iter_mut() {
-            let state = cell.get_mut();
-            new_sst_files.append(&mut state.new_sst_files);
-            new_blob_files.append(&mut state.new_blob_files);
-            for (family, global_collector) in all_collectors.iter_mut().enumerate() {
-                if let Some(collector) = state.collectors[family].take() {
-                    if !collector.is_empty() {
-                        global_collector.push(Some(collector));
-                    }
-                }
-            }
-        }
-        let shared_new_sst_files = Mutex::new(&mut new_sst_files);
         let shared_error = Mutex::new(Ok(()));
+
+        // First, we flush all thread local collectors to the global collectors.
         scope(|scope| {
-            fn handle_done_collector<'scope, K: StoreKey + Send + Sync, const FAMILIES: usize>(
-                this: &'scope WriteBatch<K, FAMILIES>,
-                scope: &Scope<'scope>,
-                family: u32,
-                mut collector: Collector<K>,
-                shared_new_sst_files: &'scope Mutex<&mut Vec<(u32, File)>>,
-                shared_error: &'scope Mutex<Result<()>>,
-            ) {
-                scope.spawn(
-                    move |_| match this.create_sst_file(family, collector.sorted()) {
-                        Ok(sst) => {
-                            collector.clear();
-                            this.idle_collectors.lock().push(collector);
-                            shared_new_sst_files.lock().push(sst);
+            let mut collectors = [const { Vec::new() }; FAMILIES];
+            for cell in self.thread_locals.iter_mut() {
+                let state = cell.get_mut();
+                new_blob_files.append(&mut state.new_blob_files);
+                for (family, thread_local_collector) in state.collectors.iter_mut().enumerate() {
+                    if let Some(collector) = thread_local_collector.take() {
+                        if !collector.is_empty() {
+                            collectors[family].push(collector);
                         }
-                        Err(err) => {
+                    }
+                }
+            }
+            for (family, thread_local_collectors) in collectors.into_iter().enumerate() {
+                for mut collector in thread_local_collectors {
+                    let this = &self;
+                    let shared_error = &shared_error;
+                    scope.spawn(move |_| {
+                        if let Err(err) =
+                            this.flush_thread_local_collector(family as u32, &mut collector)
+                        {
                             *shared_error.lock() = Err(err);
                         }
-                    },
-                );
+                        this.idle_thread_local_collectors.lock().push(collector);
+                    });
+                }
             }
-
-            all_collectors
-                .into_par_iter()
-                .enumerate()
-                .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) {
-                            (Some(mut a), Some(mut b)) => {
-                                if a.len() < b.len() {
-                                    swap(&mut a, &mut b);
-                                }
-                                for entry in b.drain() {
-                                    if a.is_full() {
-                                        let full_collector = replace(
-                                            &mut a,
-                                            self.idle_collectors
-                                                .lock()
-                                                .pop()
-                                                .unwrap_or_else(|| Collector::new()),
-                                        );
-                                        handle_done_collector(
-                                            self,
-                                            scope,
-                                            family,
-                                            full_collector,
-                                            &shared_new_sst_files,
-                                            &shared_error,
-                                        );
-                                    }
-                                    a.add_entry(entry);
-                                }
-                                self.idle_collectors.lock().push(b);
-                                Some(a)
-                            }
-                            (Some(a), None) => Some(a),
-                            (None, Some(b)) => Some(b),
-                            (None, None) => None,
-                        },
-                    );
-                    if let Some(collector) = final_collector {
-                        handle_done_collector(
-                            self,
-                            scope,
-                            family,
-                            collector,
-                            &shared_new_sst_files,
-                            &shared_error,
-                        );
-                    }
-                });
         });
+
+        // Now we reduce the global collectors in parallel
+        let mut new_sst_files = take(self.new_sst_files.get_mut());
+        let shared_new_sst_files = Mutex::new(&mut new_sst_files);
+
+        let collectors = replace(
+            &mut self.collectors,
+            [(); FAMILIES].map(|_| {
+                Mutex::new(
+                    self.idle_collectors
+                        .lock()
+                        .pop()
+                        .unwrap_or_else(|| Collector::new()),
+                )
+            }),
+        );
+        collectors
+            .into_par_iter()
+            .enumerate()
+            .try_for_each(|(family, collector)| {
+                let family = family as u32;
+                let mut collector = collector.into_inner();
+                if !collector.is_empty() {
+                    let sst = self.create_sst_file(family, collector.sorted())?;
+                    collector.clear();
+                    self.idle_collectors.lock().push(collector);
+                    shared_new_sst_files.lock().push(sst);
+                }
+                anyhow::Ok(())
+            })?;
+
         shared_error.into_inner()?;
         let seq = self.current_sequence_number.load(Ordering::SeqCst);
         new_sst_files.sort_by_key(|(seq, _)| *seq);