Move logic into get_or_create_persistent_task, add more tests (some f…

…ailing)

crates/turbo-tasks-memory/src/memory_backend.rs

@@ -535,8 +535,29 @@ impl Backend for MemoryBackend {
             self.lookup_and_connect_task(parent_task, &self.task_cache, &task_type, turbo_tasks)
         {
             // fast pass without creating a new task
+            match &*task_type {
+                PersistentTaskType::ResolveNative(..) | PersistentTaskType::Native(..) => {
+                    self.task_statistics().increment_cache_hit(&task_type);
+                }
+                PersistentTaskType::ResolveTrait(..) => {
+                    // this lookup recursively calls
+                    // `get_or_create_persistent_task` with `ResolveNative`, so
+                    // don't log it on this call, and attribute the cache hit to
+                    // the real function
+                }
+            }
             task
         } else {
+            match &*task_type {
+                PersistentTaskType::Native(..) => {
+                    self.task_statistics().increment_cache_miss(&task_type);
+                }
+                PersistentTaskType::ResolveTrait(..) | PersistentTaskType::ResolveNative(..) => {
+                    // these types re-execute themselves as `Native` after
+                    // resolving their arguments, skip counting their
+                    // executions here to avoid double-counting
+                }
+            }
             // It's important to avoid overallocating memory as this will go into the task
             // cache and stay there forever. We can to be as small as possible.
             let (task_type_hash, mut task_type) = PreHashed::into_parts(task_type);

crates/turbo-tasks-memory/src/task.rs

@@ -718,19 +718,6 @@ impl Task {
         aggregation_context.apply_queued_updates();
         self.clear_dependencies(dependencies, backend, turbo_tasks);
 
-        if let TaskType::Persistent { ty } = &self.ty {
-            match &***ty {
-                PersistentTaskType::Native(..) => {
-                    backend.task_statistics().increment_execution_count(ty);
-                }
-                PersistentTaskType::ResolveTrait(..) | PersistentTaskType::ResolveNative(..) => {
-                    // these types re-execute themselves as `Native` after
-                    // resolving their arguments, skip counting their
-                    // executions to avoid double-counting
-                }
-            }
-        }
-
         Some(TaskExecutionSpec { future, span })
     }
 
@@ -1457,21 +1444,6 @@ impl Task {
             Done { .. } => {
                 let result = func(&mut state.output)?;
                 drop(state);
-
-                if let TaskType::Persistent { ty } = &self.ty {
-                    match &***ty {
-                        PersistentTaskType::Native(..) => {
-                            backend.task_statistics().increment_finished_read_count(ty);
-                        }
-                        PersistentTaskType::ResolveTrait(..)
-                        | PersistentTaskType::ResolveNative(..) => {
-                            // these types are re-executed as `Native` after
-                            // resolving their arguments, skip counting their
-                            // executions to avoid double-counting
-                        }
-                    }
-                }
-
                 Ok(Ok(result))
             }
             Dirty {

crates/turbo-tasks-memory/src/task_statistics.rs

@@ -6,15 +6,12 @@ use turbo_tasks::{backend::PersistentTaskType, registry, FunctionId, TraitTypeId
 
 #[derive(Default, Serialize)]
 struct TaskStatistics {
-    /// How many times the function began execution (roughly a cache miss)
-    execution_count: u32,
-    /// How many times the function was read (either from cache, or after an
-    /// execution)
-    finished_read_count: u32,
+    cache_hit: u32,
+    cache_miss: u32,
 }
 
 /// A function or a trait method id.
-#[derive(Eq, Hash, PartialEq)]
+#[derive(Debug, Eq, Hash, PartialEq)]
 enum TaskTypeId {
     FunctionId(FunctionId),
     TraitTypeId(TraitTypeId),
@@ -50,12 +47,12 @@ impl AllTasksStatistics {
         self.inner.get().is_some()
     }
 
-    pub(crate) fn increment_execution_count(&self, task_type: &PersistentTaskType) {
-        self.with_task_type_statistics(task_type, |stats| stats.execution_count += 1)
+    pub(crate) fn increment_cache_hit(&self, task_type: &PersistentTaskType) {
+        self.with_task_type_statistics(task_type, |stats| stats.cache_hit += 1)
     }
 
-    pub(crate) fn increment_finished_read_count(&self, task_type: &PersistentTaskType) {
-        self.with_task_type_statistics(task_type, |stats| stats.finished_read_count += 1)
+    pub(crate) fn increment_cache_miss(&self, task_type: &PersistentTaskType) {
+        self.with_task_type_statistics(task_type, |stats| stats.cache_miss += 1)
     }
 
     fn with_task_type_statistics(

crates/turbo-tasks-memory/tests/task_statistics.rs

@@ -1,6 +1,9 @@
 #![feature(arbitrary_self_types)]
 
-use std::{future::Future, sync::Arc};
+use std::{
+    future::{Future, IntoFuture},
+    sync::Arc,
+};
 
 use anyhow::Result;
 use once_cell::sync::Lazy;
@@ -20,15 +23,35 @@ async fn test_simple_task() {
             // use cached results
             double(i).await.unwrap();
         }
-        for i in 0..10 {
+        for i in 0..5 {
             double(i).await.unwrap();
         }
         assert_eq!(
             remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
             json!({
                 "turbo-tasks-memory::::double": {
-                    "execution_count": 10,
-                    "finished_read_count": 30,
+                    "cache_miss": 10,
+                    "cache_hit": 15,
+                },
+            })
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+async fn test_await_same_vc_multiple_times() {
+    run_with_tt(|tt| async move {
+        let dvc = double(0);
+        // this is awaited multiple times, but only resolved once
+        tokio::try_join!(dvc.into_future(), dvc.into_future()).unwrap();
+        dvc.await.unwrap();
+        assert_eq!(
+            remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
+            json!({
+                "turbo-tasks-memory::::double": {
+                    "cache_miss": 1,
+                    "cache_hit": 0,
                 },
             })
         );
@@ -45,16 +68,20 @@ async fn test_vc_receiving_task() {
             // use cached results
             double_vc(dvc).await.unwrap();
         }
+        for i in 0..5 {
+            let dvc = double(i);
+            double_vc(dvc).await.unwrap();
+        }
         assert_eq!(
             remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
             json!({
                 "turbo-tasks-memory::::double": {
-                    "execution_count": 10,
-                    "finished_read_count": 10,
+                    "cache_miss": 10,
+                    "cache_hit": 5,
                 },
                 "turbo-tasks-memory::::double_vc": {
-                    "execution_count": 10,
-                    "finished_read_count": 20,
+                    "cache_miss": 10,
+                    "cache_hit": 15,
                 },
             })
         );
@@ -70,20 +97,67 @@ async fn test_trait_methods() {
             wvc.double().await.unwrap();
             wvc.double_vc().await.unwrap();
         }
+        // use cached results
+        for i in 0..5 {
+            let wvc = wrap(i);
+            wvc.double().await.unwrap();
+            wvc.double_vc().await.unwrap();
+        }
         assert_eq!(
             remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
             json!({
                 "turbo-tasks-memory::::wrap": {
-                    "execution_count": 10,
-                    "finished_read_count": 20,
+                    "cache_miss": 10,
+                    "cache_hit": 5,
                 },
                 "turbo-tasks-memory::::WrappedU64::Doublable::double": {
-                    "execution_count": 10,
-                    "finished_read_count": 10,
+                    "cache_miss": 10,
+                    "cache_hit": 5,
                 },
                 "turbo-tasks-memory::::WrappedU64::Doublable::double_vc": {
-                    "execution_count": 10,
-                    "finished_read_count": 10,
+                    "cache_miss": 10,
+                    "cache_hit": 0,
+                },
+            })
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+async fn test_dyn_trait_methods() {
+    run_with_tt(|tt| async move {
+        for i in 0..10 {
+            let wvc: Vc<Box<dyn Doublable>> = Vc::upcast(wrap(i));
+            wvc.double().resolve().await.unwrap();
+            wvc.double_vc().resolve().await.unwrap();
+        }
+        // use cached results
+        for i in 0..5 {
+            let wvc: Vc<Box<dyn Doublable>> = Vc::upcast(wrap(i));
+            wvc.double().resolve().await.unwrap();
+            wvc.double_vc().resolve().await.unwrap();
+        }
+        // use cached results without dynamic dispatch
+        for i in 0..2 {
+            let wvc = wrap(i);
+            wvc.double().await.unwrap();
+            wvc.double_vc().await.unwrap();
+        }
+        assert_eq!(
+            remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
+            json!({
+                "turbo-tasks-memory::::wrap": {
+                    "cache_miss": 10,
+                    "cache_hit": 7,
+                },
+                "turbo-tasks-memory::::WrappedU64::Doublable::double": {
+                    "cache_miss": 10,
+                    "cache_hit": 7,
+                },
+                "turbo-tasks-memory::::WrappedU64::Doublable::double_vc": {
+                    "cache_miss": 10,
+                    "cache_hit": 7,
                 },
             })
         );
@@ -99,7 +173,12 @@ async fn test_no_execution() {
         let _ = wrap_vc(double_vc(double(123))).double().double_vc();
         assert_eq!(
             remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
-            json!({})
+            json!({
+                "turbo-tasks-memory::::double": {
+                    "cache_miss": 1,
+                    "cache_hit": 0,
+                },
+            })
         );
     })
     .await;