Turbopack: parallel drop data before shutdown

Author: sokra

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

@@ -64,7 +64,8 @@ use crate::{
         InProgressCellState, InProgressState, InProgressStateInner, OutputValue, RootType,
     },
     utils::{
-        bi_map::BiMap, chunked_vec::ChunkedVec, ptr_eq_arc::PtrEqArc, sharded::Sharded, swap_retain,
+        bi_map::BiMap, chunked_vec::ChunkedVec, dash_map_drop_contents::drop_contents,
+        ptr_eq_arc::PtrEqArc, sharded::Sharded, swap_retain,
     },
 };
 
@@ -1216,6 +1217,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             self.is_idle.store(false, Ordering::Release);
             self.verify_aggregation_graph(turbo_tasks, false);
         }
+        self.task_cache.drop_contents();
+        drop_contents(&self.transient_tasks);
+        self.storage.drop_contents();
         if let Err(err) = self.backing_storage.shutdown() {
             println!("Shutting down failed: {err}");
         }

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

@@ -16,7 +16,10 @@ use crate::{
         CachedDataItemValue, CachedDataItemValueRef, CachedDataItemValueRefMut, OutputValue,
     },
     data_storage::{AutoMapStorage, OptionStorage},
-    utils::dash_map_multi::{RefMut, get_multiple_mut},
+    utils::{
+        dash_map_drop_contents::drop_contents,
+        dash_map_multi::{RefMut, get_multiple_mut},
+    },
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -806,6 +809,11 @@ impl Storage {
             },
         )
     }
+
+    pub fn drop_contents(&self) {
+        drop_contents(&self.map);
+        drop_contents(&self.modified);
+    }
 }
 
 pub struct StorageWriteGuard<'a> {

turbopack/crates/turbo-tasks-backend/src/utils/bi_map.rs

@@ -3,6 +3,8 @@ use std::{borrow::Borrow, hash::Hash};
 use dashmap::mapref::entry::Entry;
 use turbo_tasks::FxDashMap;
 
+use crate::utils::dash_map_drop_contents::drop_contents;
+
 /// A bidirectional [`FxDashMap`] that allows lookup by key or value.
 ///
 /// As keys and values are stored twice, they should be small types, such as
@@ -53,3 +55,14 @@ where
         }
     }
 }
+
+impl<K, V> BiMap<K, V>
+where
+    K: Eq + Hash + Send + Sync,
+    V: Eq + Hash + Send + Sync,
+{
+    pub fn drop_contents(&self) {
+        drop_contents(&self.forward);
+        drop_contents(&self.reverse);
+    }
+}

turbopack/crates/turbo-tasks-backend/src/utils/dash_map_drop_contents.rs

@@ -0,0 +1,17 @@
+use std::{
+    hash::{BuildHasher, Hash},
+    mem::take,
+};
+
+use dashmap::DashMap;
+use rayon::prelude::*;
+
+pub fn drop_contents<K: Hash + Eq + Send + Sync, V: Send + Sync, H: BuildHasher + Clone>(
+    map: &DashMap<K, V, H>,
+) {
+    let shards = map.shards();
+    shards.par_iter().for_each(|shard| {
+        let table = take(&mut *shard.write());
+        drop(table);
+    });
+}

turbopack/crates/turbo-tasks-backend/src/utils/mod.rs

@@ -1,5 +1,6 @@
 pub mod bi_map;
 pub mod chunked_vec;
+pub mod dash_map_drop_contents;
 pub mod dash_map_multi;
 pub mod ptr_eq_arc;
 pub mod sharded;