refactor(turbo-tasks): Port task statistics to the new backend (#75303)

This was available on the old memory backend, this ports it to the new backend, which is important now that it's the default. Plus, the code's cleaner now that the trait resolution hack I had was cleaned up in #75214!
vercel · Feb 3, 2025 · 7a0eb3c · 7a0eb3c
1 parent 8db22f7
commit 7a0eb3c
Show file tree

Hide file tree

Showing 15 changed files with 351 additions and 320 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/napi/src/next_api/project.rs b/crates/napi/src/next_api/project.rs
@@ -1,4 +1,4 @@
-use std::{path::PathBuf, sync::Arc, thread, time::Duration};
+use std::{io::Write, path::PathBuf, sync::Arc, thread, time::Duration};
 
 use anyhow::{anyhow, bail, Context, Result};
 use napi::{
@@ -386,27 +386,21 @@ pub async fn project_new(
         memory_limit,
         dependency_tracking,
     )?;
-    if !persistent_caching {
-        use std::io::Write;
-        let stats_path = std::env::var_os("NEXT_TURBOPACK_TASK_STATISTICS");
-        if let Some(stats_path) = stats_path {
-            let Some(backend) = turbo_tasks.memory_backend() else {
-                return Err(anyhow!("task statistics require a memory backend").into());
-            };
-            let task_stats = backend.task_statistics().enable().clone();
-            exit.on_exit(async move {
-                tokio::task::spawn_blocking(move || {
-                    let mut file = std::fs::File::create(&stats_path)
-                        .with_context(|| format!("failed to create or open {stats_path:?}"))?;
-                    serde_json::to_writer(&file, &task_stats)
-                        .context("failed to serialize or write task statistics")?;
-                    file.flush().context("failed to flush file")
-                })
-                .await
-                .unwrap()
-                .unwrap();
-            });
-        }
+    let stats_path = std::env::var_os("NEXT_TURBOPACK_TASK_STATISTICS");
+    if let Some(stats_path) = stats_path {
+        let task_stats = turbo_tasks.task_statistics().enable().clone();
+        exit.on_exit(async move {
+            tokio::task::spawn_blocking(move || {
+                let mut file = std::fs::File::create(&stats_path)
+                    .with_context(|| format!("failed to create or open {stats_path:?}"))?;
+                serde_json::to_writer(&file, &task_stats)
+                    .context("failed to serialize or write task statistics")?;
+                file.flush().context("failed to flush file")
+            })
+            .await
+            .unwrap()
+            .unwrap();
+        });
     }
     let options: ProjectOptions = options.into();
     let container = turbo_tasks

diff --git a/crates/napi/src/next_api/utils.rs b/crates/napi/src/next_api/utils.rs
@@ -10,7 +10,8 @@ use napi::{
 };
 use serde::Serialize;
 use turbo_tasks::{
-    trace::TraceRawVcs, OperationVc, ReadRef, TaskId, TryJoinIterExt, TurboTasks, UpdateInfo, Vc,
+    task_statistics::TaskStatisticsApi, trace::TraceRawVcs, OperationVc, ReadRef, TaskId,
+    TryJoinIterExt, TurboTasks, TurboTasksApi, UpdateInfo, Vc,
 };
 use turbo_tasks_backend::{
     default_backing_storage, noop_backing_storage, DefaultBackingStorage, NoopBackingStorage,
@@ -108,17 +109,17 @@ impl NextTurboTasks {
         }
     }
 
-    pub fn memory_backend(&self) -> Option<&turbo_tasks_memory::MemoryBackend> {
+    pub async fn stop_and_wait(&self) {
         match self {
-            NextTurboTasks::Memory(_) => None,
-            NextTurboTasks::PersistentCaching(_) => None,
+            NextTurboTasks::Memory(turbo_tasks) => turbo_tasks.stop_and_wait().await,
+            NextTurboTasks::PersistentCaching(turbo_tasks) => turbo_tasks.stop_and_wait().await,
         }
     }
 
-    pub async fn stop_and_wait(&self) {
+    pub fn task_statistics(&self) -> &TaskStatisticsApi {
         match self {
-            NextTurboTasks::Memory(turbo_tasks) => turbo_tasks.stop_and_wait().await,
-            NextTurboTasks::PersistentCaching(turbo_tasks) => turbo_tasks.stop_and_wait().await,
+            NextTurboTasks::Memory(turbo_tasks) => turbo_tasks.task_statistics(),
+            NextTurboTasks::PersistentCaching(turbo_tasks) => turbo_tasks.task_statistics(),
         }
     }
 }

diff --git a/turbopack/crates/turbo-tasks-backend/Cargo.toml b/turbopack/crates/turbo-tasks-backend/Cargo.toml
@@ -54,6 +54,8 @@ turbo-tasks-testing = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true, features = ["async_tokio"] }
+regex = { workspace = true }
+serde_json = { workspace = true }
 
 [build-dependencies]
 turbo-tasks-build = { workspace = true }

diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
@@ -31,6 +31,7 @@ use turbo_tasks::{
     },
     event::{Event, EventListener},
     registry,
+    task_statistics::TaskStatisticsApi,
     util::IdFactoryWithReuse,
     CellId, FunctionId, RawVc, ReadConsistency, SessionId, TaskId, TraitTypeId,
     TurboTasksBackendApi, ValueTypeId, TRANSIENT_TASK_BIT,
@@ -180,6 +181,8 @@ struct TurboTasksBackendInner<B: BackingStorage> {
     idle_start_event: Event,
     idle_end_event: Event,
 
+    task_statistics: TaskStatisticsApi,
+
     backing_storage: B,
 }
 
@@ -224,6 +227,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             stopping_event: Event::new(|| "TurboTasksBackend::stopping_event".to_string()),
             idle_start_event: Event::new(|| "TurboTasksBackend::idle_start_event".to_string()),
             idle_end_event: Event::new(|| "TurboTasksBackend::idle_end_event".to_string()),
+            task_statistics: TaskStatisticsApi::default(),
             backing_storage,
         }
     }
@@ -341,6 +345,16 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
     fn should_track_children(&self) -> bool {
         self.options.children_tracking
     }
+
+    fn track_cache_hit(&self, task_type: &CachedTaskType) {
+        self.task_statistics
+            .map(|stats| stats.increment_cache_hit(task_type.fn_type));
+    }
+
+    fn track_cache_miss(&self, task_type: &CachedTaskType) {
+        self.task_statistics
+            .map(|stats| stats.increment_cache_miss(task_type.fn_type));
+    }
 }
 
 pub(crate) struct OperationGuard<'a, B: BackingStorage> {
@@ -831,10 +845,12 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> TaskId {
         if let Some(task_id) = self.task_cache.lookup_forward(&task_type) {
+            self.track_cache_hit(&task_type);
             self.connect_child(parent_task, task_id, turbo_tasks);
             return task_id;
         }
 
+        self.track_cache_miss(&task_type);
         let tx = self
             .should_restore()
             .then(|| self.backing_storage.start_read_transaction())
@@ -889,24 +905,23 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             );
         }
         if let Some(task_id) = self.task_cache.lookup_forward(&task_type) {
-            // Safety: `tx` is a valid transaction from `self.backend.backing_storage`.
+            self.track_cache_hit(&task_type);
             self.connect_child(parent_task, task_id, turbo_tasks);
             return task_id;
         }
 
+        self.track_cache_miss(&task_type);
         let task_type = Arc::new(task_type);
         let task_id = self.transient_task_id_factory.get();
         if let Err(existing_task_id) = self.task_cache.try_insert(task_type, task_id) {
             // Safety: We just created the id and failed to insert it.
             unsafe {
                 self.transient_task_id_factory.reuse(task_id);
             }
-            // Safety: `tx` is a valid transaction from `self.backend.backing_storage`.
             self.connect_child(parent_task, existing_task_id, turbo_tasks);
             return existing_task_id;
         }
 
-        // Safety: `tx` is a valid transaction from `self.backend.backing_storage`.
         self.connect_child(parent_task, task_id, turbo_tasks);
 
         task_id
@@ -2026,6 +2041,10 @@ impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
     fn dispose_root_task(&self, task_id: TaskId, turbo_tasks: &dyn TurboTasksBackendApi<Self>) {
         self.0.dispose_root_task(task_id, turbo_tasks);
     }
+
+    fn task_statistics(&self) -> &TaskStatisticsApi {
+        &self.0.task_statistics
+    }
 }
 
 // from https://github.com/tokio-rs/tokio/blob/29cd6ec1ec6f90a7ee1ad641c03e0e00badbcb0e/tokio/src/time/instant.rs#L57-L63

diff --git a/turbopack/crates/turbo-tasks-backend/tests/task_statistics.rs b/turbopack/crates/turbo-tasks-backend/tests/task_statistics.rs
@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/task_statistics.rs
diff --git a/turbopack/crates/turbo-tasks-memory/src/lib.rs b/turbopack/crates/turbo-tasks-memory/src/lib.rs
@@ -14,7 +14,5 @@ mod map_guard;
 mod memory_backend;
 mod output;
 mod task;
-mod task_statistics;
 
 pub use memory_backend::MemoryBackend;
-pub use task_statistics::{TaskStatistics, TaskStatisticsApi};
diff --git a/turbopack/crates/turbo-tasks-memory/src/memory_backend.rs b/turbopack/crates/turbo-tasks-memory/src/memory_backend.rs
@@ -23,6 +23,7 @@ use turbo_tasks::{
         TransientTaskType, TypedCellContent,
     },
     event::EventListener,
+    task_statistics::TaskStatisticsApi,
     util::{IdFactoryWithReuse, NoMoveVec},
     CellId, FunctionId, RawVc, ReadConsistency, TaskId, TaskIdSet, TraitTypeId,
     TurboTasksBackendApi, Unused, ValueTypeId, TRANSIENT_TASK_BIT,
@@ -36,7 +37,6 @@ use crate::{
     },
     output::Output,
     task::{ReadCellError, Task, TaskType},
-    task_statistics::TaskStatisticsApi,
 };
 
 fn prehash_task_type(task_type: CachedTaskType) -> PreHashed<CachedTaskType> {
@@ -333,16 +333,12 @@ impl MemoryBackend {
         }
     }
 
-    pub fn task_statistics(&self) -> &TaskStatisticsApi {
-        &self.task_statistics
-    }
-
-    fn track_cache_hit(&self, task_type: &PreHashed<CachedTaskType>) {
+    fn track_cache_hit(&self, task_type: &CachedTaskType) {
         self.task_statistics()
             .map(|stats| stats.increment_cache_hit(task_type.fn_type));
     }
 
-    fn track_cache_miss(&self, task_type: &PreHashed<CachedTaskType>) {
+    fn track_cache_miss(&self, task_type: &CachedTaskType) {
         self.task_statistics()
             .map(|stats| stats.increment_cache_miss(task_type.fn_type));
     }
@@ -777,6 +773,10 @@ impl Backend for MemoryBackend {
     fn dispose_root_task(&self, task: TaskId, turbo_tasks: &dyn TurboTasksBackendApi<Self>) {
         Task::unset_root(task, self, turbo_tasks);
     }
+
+    fn task_statistics(&self) -> &TaskStatisticsApi {
+        &self.task_statistics
+    }
 }
 
 pub(crate) enum Job {
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		../../turbo-tasks-testing/tests/task_statistics.rs