[Turbopack] Memory improvements

turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs

@@ -22,9 +22,7 @@ use crate::{
     backend::{
         get_mut, get_mut_or_insert_with,
         operation::{invalidate::make_task_dirty, ExecuteContext, Operation, TaskGuard},
-        storage::{
-            count, get, get_many, iter_many, remove, update, update_count, update_ucount_and_get,
-        },
+        storage::{count, get, get_many, iter_many, remove, update, update_count},
         TaskDataCategory,
     },
     data::{
@@ -940,10 +938,7 @@ impl AggregationUpdateQueue {
 
                     // Add the same amount of upper edges
                     if update_count!(task, Upper { task: upper_id }, count) {
-                        if !upper_id.is_transient()
-                            && update_ucount_and_get!(task, PersistentUpperCount, 1)
-                                .is_power_of_two()
-                        {
+                        if count!(task, Upper).is_power_of_two() {
                             self.push_optimize_task(task_id);
                         }
                         // When this is a new inner node, update aggregated data and
@@ -1003,10 +998,6 @@ impl AggregationUpdateQueue {
                     #[cfg(feature = "trace_aggregation_update")]
                     let _span = trace_span!("make follower").entered();
 
-                    if !upper_id.is_transient() {
-                        update_ucount_and_get!(task, PersistentUpperCount, -1);
-                    }
-
                     let upper_ids: Vec<_> = get_uppers(&upper);
 
                     // Add the same amount of follower edges
@@ -1160,8 +1151,6 @@ impl AggregationUpdateQueue {
             keep_upper
         });
         if !upper_ids.is_empty() {
-            update_ucount_and_get!(follower, PersistentUpperCount, -persistent_uppers);
-
             let data = AggregatedDataUpdate::from_task(&mut follower).invert();
             let followers: Vec<_> = get_followers(&follower);
             drop(follower);
@@ -1258,10 +1247,6 @@ impl AggregationUpdateQueue {
                 Some(old - 1)
             });
             if remove_upper {
-                if !upper_id.is_transient() {
-                    update_ucount_and_get!(follower, PersistentUpperCount, -1);
-                }
-
                 let data = AggregatedDataUpdate::from_task(&mut follower).invert();
                 let followers: Vec<_> = get_followers(&follower);
                 drop(follower);
@@ -1415,8 +1400,7 @@ impl AggregationUpdateQueue {
             #[cfg(feature = "trace_aggregation_update")]
             let _span = trace_span!("new inner").entered();
             if !upper_ids.is_empty() {
-                let new_count =
-                    update_ucount_and_get!(follower, PersistentUpperCount, persistent_uppers);
+                let new_count = count!(follower, Upper);
                 if (new_count - persistent_uppers).next_power_of_two()
                     != new_count.next_power_of_two()
                 {
@@ -1546,10 +1530,7 @@ impl AggregationUpdateQueue {
             for &(follower_id, _) in followers_with_aggregation_number.iter() {
                 let mut follower = ctx.task(follower_id, TaskDataCategory::Meta);
                 if update_count!(follower, Upper { task: upper_id }, 1) {
-                    if !upper_id.is_transient()
-                        && update_ucount_and_get!(follower, PersistentUpperCount, 1)
-                            .is_power_of_two()
-                    {
+                    if !upper_id.is_transient() && count!(follower, Upper).is_power_of_two() {
                         self.push_optimize_task(follower_id);
                     }
 
@@ -1715,9 +1696,7 @@ impl AggregationUpdateQueue {
             drop(upper);
             let mut follower = ctx.task(new_follower_id, TaskDataCategory::Meta);
             if update_count!(follower, Upper { task: upper_id }, 1) {
-                if !upper_id.is_transient()
-                    && update_ucount_and_get!(follower, PersistentUpperCount, 1).is_power_of_two()
-                {
+                if !upper_id.is_transient() && count!(follower, Upper).is_power_of_two() {
                     self.push_optimize_task(new_follower_id);
                 }
                 // It's a new upper
@@ -1915,9 +1894,7 @@ impl AggregationUpdateQueue {
             }
             children_count
         };
-        let upper_count = get!(task, PersistentUpperCount)
-            .copied()
-            .unwrap_or_default() as usize;
+        let upper_count = count!(task, Upper);
         if upper_count <= 1
             || upper_count.saturating_sub(1) * follower_count
                 <= max(

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

@@ -147,6 +147,7 @@ impl InnerStorage {
         if let Some(i) = i {
             &mut self.map[i]
         } else {
+            self.map.reserve_exact(1);
             self.map.push(CachedDataItemStorage::new(ty));
             self.map.last_mut().unwrap()
         }
@@ -160,6 +161,18 @@ impl InnerStorage {
         self.map.iter_mut().position(|m| m.ty() == ty)
     }
 
+    fn get_or_create_map_index(&mut self, ty: CachedDataItemType) -> usize {
+        let i = self.map.iter().position(|m| m.ty() == ty);
+        if let Some(i) = i {
+            i
+        } else {
+            let i = self.map.len();
+            self.map.reserve_exact(1);
+            self.map.push(CachedDataItemStorage::new(ty));
+            i
+        }
+    }
+
     fn get_map(&self, ty: CachedDataItemType) -> Option<&CachedDataItemStorage> {
         self.map.iter().find(|m| m.ty() == ty)
     }
@@ -180,6 +193,7 @@ impl InnerStorage {
             let result = storage.remove(key);
             if result.is_some() && storage.is_empty() {
                 self.map.swap_remove(i);
+                self.map.shrink_to_fit();
             }
             result
         })
@@ -254,6 +268,7 @@ impl InnerStorage {
             .collect::<Vec<_>>();
         if self.map[i].is_empty() {
             self.map.swap_remove(i);
+            self.map.shrink_to_fit();
         }
         Either::Right(items.into_iter())
     }
@@ -263,9 +278,13 @@ impl InnerStorage {
         key: CachedDataItemKey,
         update: impl FnOnce(Option<CachedDataItemValue>) -> Option<CachedDataItemValue>,
     ) {
-        let map = self.get_or_create_map_mut(key.ty());
+        let i = self.get_or_create_map_index(key.ty());
+        let map = &mut self.map[i];
         if let Some(v) = update(map.remove(&key)) {
             map.insert(CachedDataItem::from_key_and_value(key, v));
+        } else if map.is_empty() {
+            self.map.swap_remove(i);
+            self.map.shrink_to_fit();
         }
     }
 
@@ -458,42 +477,6 @@ macro_rules! update {
     };
 }
 
-macro_rules! update_ucount_and_get {
-    ($task:ident, $key:ident $input:tt, -$update:expr) => {{
-        let update = $update;
-        let mut value = 0;
-        $crate::backend::storage::update!($task, $key $input, |old: Option<_>| {
-            if let Some(old) = old {
-                value = old - update;
-                (value != 0).then_some(value)
-            } else {
-                None
-            }
-        });
-        value
-    }};
-    ($task:ident, $key:ident $input:tt, $update:expr) => {{
-        let update = $update;
-        let mut value = 0;
-        $crate::backend::storage::update!($task, $key $input, |old: Option<_>| {
-            if let Some(old) = old {
-                value = old + update;
-                (value != 0).then_some(value)
-            } else {
-                value = update;
-                (update != 0).then_some(update)
-            }
-        });
-        value
-    }};
-    ($task:ident, $key:ident, -$update:expr) => {
-        $crate::backend::storage::update_ucount_and_get!($task, $key {}, -$update)
-    };
-    ($task:ident, $key:ident, $update:expr) => {
-        $crate::backend::storage::update_ucount_and_get!($task, $key {}, $update)
-    };
-}
-
 macro_rules! update_count {
     ($task:ident, $key:ident $input:tt, -$update:expr) => {{
         let update = $update;
@@ -562,4 +545,3 @@ pub(crate) use iter_many;
 pub(crate) use remove;
 pub(crate) use update;
 pub(crate) use update_count;
-pub(crate) use update_ucount_and_get;

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

@@ -418,10 +418,6 @@ pub enum CachedDataItem {
         task: TaskId,
         value: i32,
     },
-    PersistentUpperCount {
-        // Only counting persistent tasks
-        value: u32,
-    },
 
     // Aggregated Data
     AggregatedDirtyContainer {
@@ -505,7 +501,6 @@ impl CachedDataItem {
             CachedDataItem::AggregationNumber { .. } => true,
             CachedDataItem::Follower { task, .. } => !task.is_transient(),
             CachedDataItem::Upper { task, .. } => !task.is_transient(),
-            CachedDataItem::PersistentUpperCount { .. } => true,
             CachedDataItem::AggregatedDirtyContainer { task, .. } => !task.is_transient(),
             CachedDataItem::AggregatedCollectible { collectible, .. } => {
                 !collectible.cell.task.is_transient()
@@ -563,7 +558,6 @@ impl CachedDataItem {
             | Self::Dirty { .. }
             | Self::Follower { .. }
             | Self::Upper { .. }
-            | Self::PersistentUpperCount { .. }
             | Self::AggregatedDirtyContainer { .. }
             | Self::AggregatedCollectible { .. }
             | Self::AggregatedDirtyContainerCount { .. } => TaskDataCategory::Meta,
@@ -601,7 +595,6 @@ impl CachedDataItemKey {
             CachedDataItemKey::AggregationNumber { .. } => true,
             CachedDataItemKey::Follower { task, .. } => !task.is_transient(),
             CachedDataItemKey::Upper { task, .. } => !task.is_transient(),
-            CachedDataItemKey::PersistentUpperCount {} => true,
             CachedDataItemKey::AggregatedDirtyContainer { task, .. } => !task.is_transient(),
             CachedDataItemKey::AggregatedCollectible { collectible, .. } => {
                 !collectible.cell.task.is_transient()
@@ -647,7 +640,6 @@ impl CachedDataItemType {
             | Self::Dirty { .. }
             | Self::Follower { .. }
             | Self::Upper { .. }
-            | Self::PersistentUpperCount { .. }
             | Self::AggregatedDirtyContainer { .. }
             | Self::AggregatedCollectible { .. }
             | Self::AggregatedDirtyContainerCount { .. } => TaskDataCategory::Meta,

turbopack/crates/turbopack-core/src/resolve/mod.rs

@@ -1583,6 +1583,7 @@ pub async fn url_resolve(
     .await
 }
 
+#[tracing::instrument(level = "trace", skip_all)]
 async fn handle_before_resolve_plugins(
     lookup_path: Vc<FileSystemPath>,
     reference_type: Value<ReferenceType>,
@@ -1591,7 +1592,7 @@ async fn handle_before_resolve_plugins(
 ) -> Result<Option<Vc<ResolveResult>>> {
     for plugin in &options.await?.before_resolve_plugins {
         let condition = plugin.before_resolve_condition().resolve().await?;
-        if !condition.await?.matches(request).await? {
+        if !*condition.matches(request).await? {
             continue;
         }
 
@@ -1605,6 +1606,7 @@ async fn handle_before_resolve_plugins(
     Ok(None)
 }
 
+#[tracing::instrument(level = "trace", skip_all)]
 async fn handle_after_resolve_plugins(
     lookup_path: Vc<FileSystemPath>,
     reference_type: Value<ReferenceType>,

turbopack/crates/turbopack-core/src/resolve/plugin.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use anyhow::Result;
 use turbo_rcstr::RcStr;
 use turbo_tasks::{ResolvedVc, Value, Vc};
@@ -43,14 +45,14 @@ impl AfterResolvePluginCondition {
 #[turbo_tasks::value]
 pub enum BeforeResolvePluginCondition {
     Request(ResolvedVc<Glob>),
-    Modules(ResolvedVc<Vec<RcStr>>),
+    Modules(HashSet<RcStr>),
 }
 
 #[turbo_tasks::value_impl]
 impl BeforeResolvePluginCondition {
     #[turbo_tasks::function]
-    pub fn from_modules(modules: ResolvedVc<Vec<RcStr>>) -> Vc<Self> {
-        BeforeResolvePluginCondition::Modules(modules).cell()
+    pub async fn from_modules(modules: ResolvedVc<Vec<RcStr>>) -> Result<Vc<Self>> {
+        Ok(BeforeResolvePluginCondition::Modules(modules.await?.iter().cloned().collect()).cell())
     }
 
     #[turbo_tasks::function]
@@ -59,21 +61,23 @@ impl BeforeResolvePluginCondition {
     }
 }
 
+#[turbo_tasks::value_impl]
 impl BeforeResolvePluginCondition {
-    pub async fn matches(&self, request: Vc<Request>) -> Result<bool> {
-        Ok(match self {
+    #[turbo_tasks::function]
+    pub async fn matches(&self, request: Vc<Request>) -> Result<Vc<bool>> {
+        Ok(Vc::cell(match self {
             BeforeResolvePluginCondition::Request(glob) => match request.await?.request() {
                 Some(request) => glob.await?.execute(request.as_str()),
                 None => false,
             },
             BeforeResolvePluginCondition::Modules(modules) => {
                 if let Request::Module { module, .. } = &*request.await? {
-                    modules.await?.contains(module)
+                    modules.contains(module)
                 } else {
                     false
                 }
             }
-        })
+        }))
     }
 }