flush key spaces when done with them

Author: sokra

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

@@ -1,6 +1,7 @@
 #![feature(once_cell_try)]
 #![feature(new_zeroed_alloc)]
 #![feature(get_mut_unchecked)]
+#![feature(sync_unsafe_cell)]
 
 mod arc_slice;
 mod collector;

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

@@ -1,5 +1,5 @@
 use std::{
-    cell::UnsafeCell,
+    cell::SyncUnsafeCell,
     fs::File,
     io::Write,
     mem::{replace, take},
@@ -68,7 +68,7 @@ pub struct WriteBatch<K: StoreKey + Send, const FAMILIES: usize> {
     /// The current sequence number counter. Increased for every new SST file or blob file.
     current_sequence_number: AtomicU32,
     /// The thread local state.
-    thread_locals: ThreadLocal<UnsafeCell<ThreadLocalState<K, FAMILIES>>>,
+    thread_locals: ThreadLocal<SyncUnsafeCell<ThreadLocalState<K, FAMILIES>>>,
     /// Collectors in use. The thread local collectors flush into these when they are full.
     collectors: [Mutex<GlobalCollectorState<K>>; FAMILIES],
     /// The list of new SST files that have been created.
@@ -109,7 +109,7 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
     #[allow(clippy::mut_from_ref)]
     fn thread_local_state(&self) -> &mut ThreadLocalState<K, FAMILIES> {
         let cell = self.thread_locals.get_or(|| {
-            UnsafeCell::new(ThreadLocalState {
+            SyncUnsafeCell::new(ThreadLocalState {
                 collectors: [const { None }; FAMILIES],
                 new_blob_files: Vec::new(),
             })
@@ -219,6 +219,45 @@ impl<K: StoreKey + Send + Sync, const FAMILIES: usize> WriteBatch<K, FAMILIES> {
         Ok(())
     }
 
+    /// Flushes a family of the write batch, reducing the amount of buffered memory used.
+    /// Does not commit any data persistently.
+    ///
+    /// Safety: Caller must ensure that no concurrent put or delete operation is happening on the
+    /// flushed family.
+    pub unsafe fn flush(&self, family: u32) -> Result<()> {
+        let mut collectors = Vec::new();
+        for cell in self.thread_locals.iter() {
+            let state = unsafe { &mut *cell.get() };
+            if let Some(collector) = state.collectors[usize_from_u32(family)].take() {
+                if !collector.is_empty() {
+                    collectors.push(collector);
+                }
+            }
+        }
+
+        let shared_error = Mutex::new(Ok(()));
+        scope(|scope| {
+            for mut collector in collectors {
+                let this = &self;
+                let shared_error = &shared_error;
+                let span = Span::current();
+                scope.spawn(move |_| {
+                    let _span = span.entered();
+                    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);
+                });
+            }
+        });
+
+        shared_error.into_inner()?;
+
+        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> {

turbopack/crates/turbo-tasks-backend/src/database/noop_kv.rs

@@ -90,6 +90,10 @@ impl SerialWriteBatch<'_> for NoopWriteBatch {
     fn delete(&mut self, _key_space: KeySpace, _key: WriteBuffer<'_>) -> Result<()> {
         Ok(())
     }
+
+    fn flush(&mut self, _key_space: KeySpace) -> Result<()> {
+        Ok(())
+    }
 }
 
 impl ConcurrentWriteBatch<'_> for NoopWriteBatch {
@@ -105,4 +109,8 @@ impl ConcurrentWriteBatch<'_> for NoopWriteBatch {
     fn delete(&self, _key_space: KeySpace, _key: WriteBuffer<'_>) -> Result<()> {
         Ok(())
     }
+
+    unsafe fn flush(&self, _key_space: KeySpace) -> Result<()> {
+        Ok(())
+    }
 }

turbopack/crates/turbo-tasks-backend/src/database/turbo.rs

@@ -160,6 +160,10 @@ impl<'a> ConcurrentWriteBatch<'a> for TurboWriteBatch<'a> {
     fn delete(&self, key_space: KeySpace, key: WriteBuffer<'_>) -> Result<()> {
         self.batch.delete(key_space as u32, key.into_static())
     }
+
+    unsafe fn flush(&self, key_space: KeySpace) -> Result<()> {
+        self.batch.flush(key_space as u32)
+    }
 }
 
 impl KeyBase for WriteBuffer<'_> {

turbopack/crates/turbo-tasks-backend/src/database/write_batch.rs

@@ -66,11 +66,18 @@ pub trait SerialWriteBatch<'a>: BaseWriteBatch<'a> {
         value: WriteBuffer<'_>,
     ) -> Result<()>;
     fn delete(&mut self, key_space: KeySpace, key: WriteBuffer<'_>) -> Result<()>;
+    fn flush(&mut self, key_space: KeySpace) -> Result<()>;
 }
 
 pub trait ConcurrentWriteBatch<'a>: BaseWriteBatch<'a> + Sync + Send {
     fn put(&self, key_space: KeySpace, key: WriteBuffer<'_>, value: WriteBuffer<'_>) -> Result<()>;
     fn delete(&self, key_space: KeySpace, key: WriteBuffer<'_>) -> Result<()>;
+    /// Flushes a key space of the write batch, reducing the amount of buffered memory used.
+    /// Does not commit any data persistently.
+    ///
+    /// Safety: Caller must ensure that no concurrent put or delete operation is happening on the
+    /// flushed key space.
+    unsafe fn flush(&self, key_space: KeySpace) -> Result<()>;
 }
 
 pub enum WriteBatch<'a, S, C>
@@ -164,6 +171,16 @@ where
             WriteBatch::Concurrent(c, _) => c.delete(key_space, key),
         }
     }
+
+    fn flush(&mut self, key_space: KeySpace) -> Result<()> {
+        match self {
+            WriteBatch::Serial(s) => s.flush(key_space),
+            WriteBatch::Concurrent(c, _) => {
+                // Safety: the &mut self ensures that no concurrent operation is happening
+                unsafe { c.flush(key_space) }
+            }
+        }
+    }
 }
 
 pub enum WriteBatchRef<'r, 'a, S, C>
@@ -241,6 +258,16 @@ where
             WriteBatchRef::Concurrent(c, _) => c.delete(key_space, key),
         }
     }
+
+    fn flush(&mut self, key_space: KeySpace) -> Result<()> {
+        match self {
+            WriteBatchRef::Serial(s) => s.flush(key_space),
+            WriteBatchRef::Concurrent(c, _) => {
+                // Safety: the &mut self ensures that no concurrent operation is happening
+                unsafe { c.flush(key_space) }
+            }
+        }
+    }
 }
 
 pub struct UnimplementedWriteBatch;
@@ -275,6 +302,9 @@ impl SerialWriteBatch<'_> for UnimplementedWriteBatch {
     fn delete(&mut self, _key_space: KeySpace, _key: WriteBuffer<'_>) -> Result<()> {
         todo!()
     }
+    fn flush(&mut self, _key_space: KeySpace) -> Result<()> {
+        todo!()
+    }
 }
 
 impl ConcurrentWriteBatch<'_> for UnimplementedWriteBatch {
@@ -289,4 +319,7 @@ impl ConcurrentWriteBatch<'_> for UnimplementedWriteBatch {
     fn delete(&self, _key_space: KeySpace, _key: WriteBuffer<'_>) -> Result<()> {
         todo!()
     }
+    unsafe fn flush(&self, _key_space: KeySpace) -> Result<()> {
+        todo!()
+    }
 }

turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs

@@ -184,6 +184,13 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorage
                 {
                     let _span = tracing::trace_span!("update task data").entered();
                     process_task_data(snapshots, Some(batch))?;
+                    [KeySpace::TaskMeta, KeySpace::TaskData]
+                        .into_par_iter()
+                        .try_for_each(|key_space| {
+                            // Safety: We already finished all processing of the task data and task
+                            // meta
+                            unsafe { batch.flush(key_space) }
+                        })?;
                 }
 
                 let mut next_task_id = get_next_free_task_id::<
@@ -500,6 +507,7 @@ where
             )
             .with_context(|| anyhow!("Unable to write operations"))?;
     }
+    batch.flush(KeySpace::Infra)?;
     Ok(())
 }