Move memo dropping off to another thread

Veykril · Veykril · commit 0a4b3fd004ff · 2025-02-06T13:55:19.000+01:00
diff --git a/components/salsa-macro-rules/src/setup_tracked_fn.rs b/components/salsa-macro-rules/src/setup_tracked_fn.rs
@@ -201,6 +201,7 @@ macro_rules! setup_tracked_fn {
                     &self,
                     aux: &dyn $zalsa::JarAux,
                     first_index: $zalsa::IngredientIndex,
+                    memo_drop_sender: $zalsa::MemoDropSender,
                 ) -> Vec<Box<dyn $zalsa::Ingredient>> {
                     let struct_index = $zalsa::macro_if! {
                         if $needs_interner {
@@ -217,6 +218,7 @@ macro_rules! setup_tracked_fn {
                         struct_index,
                         first_index,
                         aux,
+                        memo_drop_sender
                     );
                     fn_ingredient.set_capacity($lru);
                     $zalsa::macro_if! {
diff --git a/src/accumulator.rs b/src/accumulator.rs
@@ -12,7 +12,7 @@ use accumulated::AnyAccumulated;
 use crate::{
     cycle::CycleRecoveryStrategy,
     ingredient::{fmt_index, Ingredient, Jar, MaybeChangedAfter},
-    plumbing::JarAux,
+    plumbing::{JarAux, MemoDropSender},
     table::Table,
     zalsa::IngredientIndex,
     zalsa_local::QueryOrigin,
@@ -50,6 +50,7 @@ impl<A: Accumulator> Jar for JarImpl<A> {
         &self,
         _aux: &dyn JarAux,
         first_index: IngredientIndex,
+        _: MemoDropSender,
     ) -> Vec<Box<dyn Ingredient>> {
         vec![Box::new(<IngredientImpl<A>>::new(first_index))]
     }
diff --git a/src/exclusive.rs b/src/exclusive.rs
@@ -0,0 +1,15 @@
+/// `std::sync::Exclusive`
+pub struct Exclusive<T: ?Sized> {
+    inner: T,
+}
+
+unsafe impl<T> Sync for Exclusive<T> {}
+
+impl<T> Exclusive<T> {
+    pub fn new(inner: T) -> Self {
+        Self { inner }
+    }
+    pub fn get_mut(&mut self) -> &mut T {
+        &mut self.inner
+    }
+}
diff --git a/src/function.rs b/src/function.rs
@@ -1,4 +1,4 @@
-use std::{any::Any, fmt, mem::ManuallyDrop, sync::Arc};
+use std::{any::Any, fmt, marker::PhantomData, sync::Arc};
 
 use crate::{
     accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues},
@@ -7,19 +7,16 @@ use crate::{
     key::DatabaseKeyIndex,
     plumbing::JarAux,
     salsa_struct::SalsaStructInDb,
-    table::Table,
+    table::{memo::MemoDropSender, Table},
     zalsa::{IngredientIndex, MemoIngredientIndex, Zalsa},
     zalsa_local::QueryOrigin,
     Cycle, Database, Id, Revision,
 };
 
-use self::delete::DeletedEntries;
-
 use super::ingredient::Ingredient;
 
 mod accumulated;
 mod backdate;
-mod delete;
 mod diff_outputs;
 mod execute;
 mod fetch;
@@ -113,7 +110,8 @@ pub struct IngredientImpl<C: Configuration> {
     /// current revision: you would be right, but we are being defensive, because
     /// we don't know that we can trust the database to give us the same runtime
     /// everytime and so forth.
-    deleted_entries: DeletedEntries<C>,
+    delete: MemoDropSender,
+    config: PhantomData<fn(C) -> C>,
 }
 
 /// True if `old_value == new_value`. Invoked by the generated
@@ -127,12 +125,18 @@ impl<C> IngredientImpl<C>
 where
     C: Configuration,
 {
-    pub fn new(struct_index: IngredientIndex, index: IngredientIndex, aux: &dyn JarAux) -> Self {
+    pub fn new(
+        struct_index: IngredientIndex,
+        index: IngredientIndex,
+        aux: &dyn JarAux,
+        delete: MemoDropSender,
+    ) -> Self {
         Self {
             index,
             memo_ingredient_index: aux.next_memo_ingredient_index(struct_index, index),
             lru: Default::default(),
-            deleted_entries: Default::default(),
+            delete,
+            config: PhantomData,
         }
     }
 
@@ -171,14 +175,7 @@ where
         // Unsafety conditions: memo must be in the map (it's not yet, but it will be by the time this
         // value is returned) and anything removed from map is added to deleted entries (ensured elsewhere).
         let db_memo = unsafe { self.extend_memo_lifetime(&memo) };
-        // Safety: We delay the drop of `old_value` until a new revision starts which ensures no
-        // references will exist for the memo contents.
-        if let Some(old_value) = unsafe { self.insert_memo_into_table_for(zalsa, id, memo) } {
-            // In case there is a reference to the old memo out there, we have to store it
-            // in the deleted entries. This will get cleared when a new revision starts.
-            self.deleted_entries
-                .push(ManuallyDrop::into_inner(old_value));
-        }
+        self.insert_memo_into_table_for(zalsa, id, memo);
         db_memo
     }
 }
@@ -236,7 +233,6 @@ where
     fn reset_for_new_revision(&mut self, table: &mut Table) {
         self.lru
             .for_each_evicted(|evict| self.evict_value_from_memo_for(table.memos_mut(evict)));
-        std::mem::take(&mut self.deleted_entries);
     }
 
     fn fmt_index(&self, index: Option<crate::Id>, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
diff --git a/src/function/delete.rs b/src/function/delete.rs
diff --git a/src/function/memo.rs b/src/function/memo.rs
@@ -35,25 +35,22 @@ impl<C: Configuration> IngredientImpl<C> {
     }
 
     /// Inserts the memo for the given key; (atomically) overwrites and returns any previously existing memo
-    ///
-    /// # Safety
-    ///
-    /// The caller needs to make sure to not drop the returned value until no more references into
-    /// the database exist as there may be outstanding borrows into the `Arc` contents.
-    pub(super) unsafe fn insert_memo_into_table_for<'db>(
+    pub(super) fn insert_memo_into_table_for<'db>(
         &'db self,
         zalsa: &'db Zalsa,
         id: Id,
         memo: ArcMemo<'db, C>,
-    ) -> Option<ManuallyDrop<ArcMemo<'db, C>>> {
+    ) {
         let static_memo = unsafe { self.to_static(memo) };
-        let old_static_memo = unsafe {
+        // SAFETY: We delay the deletion of the old memo until the next revision starts.
+        let old_memo = unsafe {
             zalsa
                 .memo_table_for(id)
                 .insert(self.memo_ingredient_index, static_memo)
-        }?;
-        let old_static_memo = ManuallyDrop::into_inner(old_static_memo);
-        Some(ManuallyDrop::new(unsafe { self.to_self(old_static_memo) }))
+        };
+        if let Some(old_memo) = old_memo {
+            self.delete.delay(ManuallyDrop::into_inner(old_memo));
+        }
     }
 
     /// Loads the current memo for `key_index`. This does not hold any sort of
@@ -111,12 +108,11 @@ impl<C: Configuration> IngredientImpl<C> {
                 }
             }
         };
-        // SAFETY: We queue the old value for deletion, delaying its drop until the next revision bump.
-        let old = unsafe { table.map_memo(self.memo_ingredient_index, map) };
-        if let Some(old) = old {
-            // In case there is a reference to the old memo out there, we have to store it
-            // in the deleted entries. This will get cleared when a new revision starts.
-            self.deleted_entries.push(ManuallyDrop::into_inner(old));
+        // SAFETY: We queue the old value for deletion, delaying its drop until the next revision
+        // bump.
+        let old_memo = unsafe { table.map_memo(self.memo_ingredient_index, map) };
+        if let Some(old_memo) = old_memo {
+            self.delete.delay(ManuallyDrop::into_inner(old_memo));
         }
     }
 }
diff --git a/src/ingredient.rs b/src/ingredient.rs
@@ -6,6 +6,7 @@ use std::{
 use crate::{
     accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues},
     cycle::CycleRecoveryStrategy,
+    plumbing::MemoDropSender,
     table::Table,
     zalsa::{IngredientIndex, MemoIngredientIndex},
     zalsa_local::QueryOrigin,
@@ -23,6 +24,7 @@ pub trait Jar: Any {
         &self,
         aux: &dyn JarAux,
         first_index: IngredientIndex,
+        memo_drop_sender: MemoDropSender,
     ) -> Vec<Box<dyn Ingredient>>;
 
     /// If this jar's first ingredient is a salsa struct, return its `TypeId`
diff --git a/src/input.rs b/src/input.rs
@@ -17,7 +17,7 @@ use crate::{
     ingredient::{fmt_index, Ingredient, MaybeChangedAfter},
     input::singleton::{Singleton, SingletonChoice},
     key::{DatabaseKeyIndex, InputDependencyIndex},
-    plumbing::{Jar, JarAux, Stamp},
+    plumbing::{Jar, JarAux, MemoDropSender, Stamp},
     table::{memo::MemoTable, sync::SyncTable, Slot, Table},
     zalsa::{IngredientIndex, Zalsa},
     zalsa_local::QueryOrigin,
@@ -58,6 +58,7 @@ impl<C: Configuration> Jar for JarImpl<C> {
         &self,
         _aux: &dyn JarAux,
         struct_index: crate::zalsa::IngredientIndex,
+        _: MemoDropSender,
     ) -> Vec<Box<dyn Ingredient>> {
         let struct_ingredient: IngredientImpl<C> = IngredientImpl::new(struct_index);
 
diff --git a/src/interned.rs b/src/interned.rs
@@ -4,7 +4,7 @@ use crate::accumulator::accumulated_map::InputAccumulatedValues;
 use crate::durability::Durability;
 use crate::ingredient::{fmt_index, MaybeChangedAfter};
 use crate::key::InputDependencyIndex;
-use crate::plumbing::{Jar, JarAux};
+use crate::plumbing::{Jar, JarAux, MemoDropSender};
 use crate::table::memo::MemoTable;
 use crate::table::sync::SyncTable;
 use crate::table::{Slot, Table};
@@ -103,6 +103,7 @@ impl<C: Configuration> Jar for JarImpl<C> {
         &self,
         _aux: &dyn JarAux,
         first_index: IngredientIndex,
+        _: MemoDropSender,
     ) -> Vec<Box<dyn Ingredient>> {
         vec![Box::new(IngredientImpl::<C>::new(first_index)) as _]
     }
diff --git a/src/lib.rs b/src/lib.rs
@@ -8,6 +8,7 @@ mod database;
 mod database_impl;
 mod durability;
 mod event;
+mod exclusive;
 mod function;
 mod hash;
 mod id;
@@ -90,6 +91,7 @@ pub mod plumbing {
     pub use crate::salsa_struct::SalsaStructInDb;
     pub use crate::storage::HasStorage;
     pub use crate::storage::Storage;
+    pub use crate::table::memo::MemoDropSender;
     pub use crate::tracked_struct::TrackedStructInDb;
     pub use crate::update::always_update;
     pub use crate::update::helper::Dispatch as UpdateDispatch;
diff --git a/src/runtime.rs b/src/runtime.rs
@@ -5,14 +5,24 @@ use std::{
         atomic::{AtomicBool, Ordering},
         Arc,
     },
-    thread::ThreadId,
+    thread::{JoinHandle, ThreadId},
 };
 
+use crossbeam::sync::{Parker, Unparker};
 use parking_lot::Mutex;
 
 use crate::{
-    active_query::ActiveQuery, cycle::CycleRecoveryStrategy, durability::Durability,
-    key::DatabaseKeyIndex, revision::AtomicRevision, table::Table, zalsa_local::ZalsaLocal,
+    active_query::ActiveQuery,
+    cycle::CycleRecoveryStrategy,
+    durability::Durability,
+    exclusive::Exclusive,
+    key::DatabaseKeyIndex,
+    revision::AtomicRevision,
+    table::{
+        memo::{memo_drop_channel, spawn_memo_drop_thread, MemoDropReceiver, MemoDropSender},
+        Table,
+    },
+    zalsa_local::ZalsaLocal,
     Cancelled, Cycle, Database, Event, EventKind, Revision,
 };
 
@@ -43,6 +53,14 @@ pub struct Runtime {
 
     /// Data for instances
     table: Table,
+    // DROP ORDER: `memo_drop_sender` and `unparker` need to drop before `memo_drop_thread`
+    // as `memo_drop_thread` exits once `memo_drop_sender` drops and its last parker barrier unparks
+    memo_drop_unparker: UnparkOnDrop,
+    memo_drop_sender: MemoDropSender,
+    memo_drop_receiver: Exclusive<Option<MemoDropReceiver>>,
+    memo_drop_thread: JoinOnDrop,
+    // Do this once we can do via TAIT instead of storing the receiver temporarily in a field.
+    // memo_drop_thread: LazyLock<Thread, impl FnOnce() -> Thread>,
 }
 
 #[derive(Clone, Debug)]
@@ -80,11 +98,20 @@ impl<V> StampedValue<V> {
 
 impl Default for Runtime {
     fn default() -> Self {
+        let (memo_drop_sender, memo_drop_receiver) = memo_drop_channel();
+
+        let parker = Parker::new();
+        let unparker = UnparkOnDrop(parker.unparker().clone());
+        memo_drop_sender.park(parker);
         Runtime {
             revisions: [const { AtomicRevision::start() }; Durability::LEN],
             revision_canceled: Default::default(),
             dependency_graph: Default::default(),
             table: Default::default(),
+            memo_drop_unparker: unparker,
+            memo_drop_sender,
+            memo_drop_receiver: Exclusive::new(Some(memo_drop_receiver)),
+            memo_drop_thread: Default::default(),
         }
     }
 }
@@ -151,9 +178,20 @@ impl Runtime {
         let r_new = r_old.next();
         self.revisions[0].store(r_new);
         self.revision_canceled.store(false, Ordering::Release);
+        self.memo_drop_thread.0.get_or_insert_with(|| {
+            spawn_memo_drop_thread(self.memo_drop_receiver.get_mut().take().unwrap())
+        });
         r_new
     }
 
+    /// Releases the previous barrier and acquires a new one, effectively kicking off a destruction
+    /// cycle for all collected memos up to this point.
+    pub(crate) fn memo_drop_barrier(&mut self) {
+        let parker = Parker::new();
+        mem::replace(&mut self.memo_drop_unparker.0, parker.unparker().clone()).unpark();
+        self.memo_drop_sender.park(parker);
+    }
+
     /// Block until `other_id` completes executing `database_key`;
     /// panic or unwind in the case of a cycle.
     ///
@@ -353,4 +391,27 @@ impl Runtime {
             .lock()
             .unblock_runtimes_blocked_on(database_key, wait_result);
     }
+
+    pub(crate) fn memo_drop_sender(&self) -> MemoDropSender {
+        self.memo_drop_sender.clone()
+    }
+}
+
+#[derive(Default)]
+struct JoinOnDrop(Option<JoinHandle<()>>);
+
+impl Drop for JoinOnDrop {
+    fn drop(&mut self) {
+        if let Some(drop_thread) = self.0.take() {
+            drop_thread.join().unwrap();
+        }
+    }
+}
+
+struct UnparkOnDrop(Unparker);
+
+impl Drop for UnparkOnDrop {
+    fn drop(&mut self) {
+        self.0.unpark();
+    }
 }
diff --git a/src/table/memo.rs b/src/table/memo.rs
diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs
diff --git a/src/zalsa.rs b/src/zalsa.rs
diff --git a/tests/lru.rs b/tests/lru.rs
