Repository: vercel/next.js
Pull request: [turbopack] add task type infromation to the print_cache_item_size feature

Enhance our debugging output to include information about TaskType values

example output:
```
Task type cache stats: 337.82MiB (364.22MiB)
  49.98MiB (53.43MiB) <ModuleAssetContext as AssetContext>::resolve_asset  x 417238 avg 125B (134B)
  29.38MiB (31.46MiB) resolve  x 297288 avg 103B (110B)
  17.97MiB (22.58MiB) AssetIdent::new_inner  x 135310 avg 139B (174B)
  15.16MiB (21.86MiB) <ExternalCjsModulesResolvePlugin as AfterResolvePlugin>::after_resolve  x 90748 avg 175B (252B)
  12.36MiB (13.32MiB) resolve_internal  x 138970 avg 93B (100B)
  8.44MiB (7.86MiB) ResolveOrigin::resolve_options  x 301586 avg 29B (27B)
  8.43MiB (9.47MiB) AfterResolvePluginCondition::matches  x 87740 avg 100B (113B)
  6.70MiB (6.23MiB) <EsmAssetReference as ModuleReference>::resolve_reference  x 472105 avg 14B (13B)
  6.48MiB (7.81MiB) read_matches  x 81046 avg 83B (101B)
  6.40MiB (5.95MiB) <EsmAssetReference as ChunkableModuleReference>::chunking_type  x 451598 avg 14B (13B)
  6.39MiB (5.94MiB) <EsmAssetReference as ChunkableModuleReference>::binding_usage  x 450844 avg 14B (13B)
  6.03MiB (5.68MiB) <ModuleAssetContext as AssetContext>::process_resolve_result  x 184982 avg 34B (32B)
  5.99MiB (5.57MiB) EsmAssetReference::get_referenced_asset  x 422883 avg 14B (13B)
```

also add summary stats to the 'task cache stats'

```
Task cache stats: 3.23GiB (6.49GiB)
  302.72MiB (769.45MiB) <Code as GenerateSourceMap>::generate_source_map = 313.04KiB (308.92KiB) meta 8524 x 37B (37B), 302.42MiB (769.15MiB) data 8524 x 36.33KiB (92.40KiB)
  212.75MiB (592.88MiB) module_factory_with_code_generation_issue = 1.59MiB (1.51MiB) meta 45906 x 36B (34B), 211.16MiB (591.37MiB) data 45906 x 4.71KiB (13.19KiB)
  207.89MiB (526.37MiB) EcmascriptBuildNodeChunkContent::code = 1.08MiB (1.12MiB) meta 5163 x 219B (228B), 206.81MiB (525.24MiB) data 5163 x 41.02KiB (104.17KiB)
  205.21MiB (577.16MiB) EcmascriptChunkItemContent::new = 1.56MiB (1.63MiB) meta 24577 x 66B (69B), 203.65MiB (575.53MiB) data 24577 x 8.48KiB (23.98KiB)
  179.00MiB (441.68MiB) EcmascriptBrowserChunkContent::code = 905.40KiB (967.46KiB) meta 2485 x 373B (398B), 178.11MiB (440.74MiB) data 2485 x 73.40KiB (181.62KiB)
```

Using this information i can see that the largest task type is ~365b so we can safely decrease the size of our scratch buffer

Changed files:
filename: turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
patch: @@ -1177,14 +1177,28 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     .into_iter()
                     .collect::<Vec<_>>();
                 if !task_cache_stats.is_empty() {
+                    use turbo_tasks::util::FormatBytes;
                     task_cache_stats.sort_unstable_by(|(key_a, stats_a), (key_b, stats_b)| {
                         (stats_b.data_compressed + stats_b.meta_compressed, key_b)
                             .cmp(&(stats_a.data_compressed + stats_a.meta_compressed, key_a))
                     });
-                    println!("Task cache stats:");
-                    for (task_desc, stats) in task_cache_stats {
-                        use turbo_tasks::util::FormatBytes;
+                    println!(
+                        "Task cache stats: {} ({})",
+                        FormatBytes(
+                            task_cache_stats
+                                .iter()
+                                .map(|(_, s)| s.data_compressed + s.meta_compressed)
+                                .sum::<usize>()
+                        ),
+                        FormatBytes(
+                            task_cache_stats
+                                .iter()
+                                .map(|(_, s)| s.data + s.meta)
+                                .sum::<usize>()
+                        )
+                    );
 
+                    for (task_desc, stats) in task_cache_stats {
                         println!(
                             "  {} ({}) {task_desc} = {} ({}) meta {} x {} ({}), {} ({}) data {} x \
                              {} ({})",

filename: turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs
patch: @@ -279,8 +279,13 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorageSealed
         let mut batch = self.inner.database.write_batch()?;
 
         // these buffers should be large, because they're temporary and re-used.
-        const INITIAL_ENCODE_BUFFER_CAPACITY: usize = 1024;
-
+        // From measuring a large application the largest TaskType was ~365b, so this should be big
+        // enough to trigger no resizes in the loop.
+        const INITIAL_ENCODE_BUFFER_CAPACITY: usize = 512;
+        #[cfg(feature = "print_cache_item_size")]
+        let all_stats: std::sync::Mutex<
+            std::collections::HashMap<&'static str, TaskTypeCacheStats>,
+        > = std::sync::Mutex::new(std::collections::HashMap::new());
         // Start organizing the updates in parallel
         match &mut batch {
             &mut WriteBatch::Concurrent(ref batch, _) => {
@@ -337,6 +342,13 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorageSealed
                                             "Unable to write task cache {task_type:?} => {task_id}"
                                         )
                                     })?;
+                                #[cfg(feature = "print_cache_item_size")]
+                                all_stats
+                                    .lock()
+                                    .unwrap()
+                                    .entry(task_type.get_name())
+                                    .or_default()
+                                    .add(&task_type_bytes);
                                 max_task_id = max_task_id.max(task_id);
                             }
 
@@ -411,6 +423,13 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorageSealed
                             .with_context(|| {
                                 format!("Unable to write task cache {task_type:?} => {task_id}")
                             })?;
+                        #[cfg(feature = "print_cache_item_size")]
+                        all_stats
+                            .lock()
+                            .unwrap()
+                            .entry(task_type.get_name())
+                            .or_default()
+                            .add(&task_type_bytes);
                         next_task_id = next_task_id.max(task_id + 1);
                     }
                 }
@@ -422,6 +441,8 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorageSealed
                 )?;
             }
         }
+        #[cfg(feature = "print_cache_item_size")]
+        print_task_type_cache_stats(all_stats.into_inner().unwrap());
 
         {
             let _span = tracing::trace_span!("commit").entered();
@@ -668,6 +689,67 @@ type SerializedTasks = Vec<
     )>,
 >;
 
+#[cfg(feature = "print_cache_item_size")]
+#[derive(Default)]
+struct TaskTypeCacheStats {
+    key_size: usize,
+    key_size_compressed: usize,
+    count: usize,
+}
+
+#[cfg(feature = "print_cache_item_size")]
+impl TaskTypeCacheStats {
+    fn compressed_size(data: &[u8]) -> Result<usize> {
+        Ok(lzzzz::lz4::Compressor::new()?.next_to_vec(
+            data,
+            &mut Vec::new(),
+            lzzzz::lz4::ACC_LEVEL_DEFAULT,
+        )?)
+    }
+    fn add(&mut self, key_bytes: &[u8]) {
+        self.key_size += key_bytes.len();
+        self.key_size_compressed += Self::compressed_size(key_bytes).unwrap_or(0);
+        self.count += 1;
+    }
+}
+
+#[cfg(feature = "print_cache_item_size")]
+fn print_task_type_cache_stats(stats: std::collections::HashMap<&'static str, TaskTypeCacheStats>) {
+    use turbo_tasks::util::FormatBytes;
+
+    let mut stats: Vec<_> = stats.into_iter().collect();
+    if stats.is_empty() {
+        return;
+    }
+    stats.sort_unstable_by(|(key_a, stats_a), (key_b, stats_b)| {
+        (stats_b.key_size_compressed, *key_b).cmp(&(stats_a.key_size_compressed, *key_a))
+    });
+    println!(
+        "Task type cache stats: {} ({})",
+        FormatBytes(
+            stats
+                .iter()
+                .map(|(_, s)| s.key_size_compressed)
+                .sum::<usize>()
+        ),
+        FormatBytes(stats.iter().map(|(_, s)| s.key_size).sum::<usize>())
+    );
+    for (fn_name, stats) in stats {
+        println!(
+            "  {} ({}) {fn_name}  x {} avg {} ({})",
+            FormatBytes(stats.key_size_compressed),
+            FormatBytes(stats.key_size),
+            stats.count,
+            FormatBytes(
+                stats
+                    .key_size_compressed
+                    .checked_div(stats.count)
+                    .unwrap_or(0)
+            ),
+            FormatBytes(stats.key_size.checked_div(stats.count).unwrap_or(0)),
+        );
+    }
+}
 fn process_task_data<'a, B: ConcurrentWriteBatch<'a> + Send + Sync, I>(
     tasks: Vec<I>,
     batch: Option<&B>,