Add more tests, fix double-counting of ResolveNative

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

@@ -717,9 +717,20 @@ impl Task {
         };
         aggregation_context.apply_queued_updates();
         self.clear_dependencies(dependencies, backend, turbo_tasks);
+
         if let TaskType::Persistent { ty } = &self.ty {
-            backend.task_statistics().increment_execution_count(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 })
     }
 
@@ -1448,7 +1459,17 @@ impl Task {
                 drop(state);
 
                 if let TaskType::Persistent { ty } = &self.ty {
-                    backend.task_statistics().increment_finished_read_count(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))

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

@@ -1,15 +1,15 @@
-use std::sync::{Arc, OnceLock};
+use std::sync::OnceLock;
 
 use dashmap::DashMap;
 use serde::{ser::SerializeMap, Serialize, Serializer};
 use turbo_tasks::{backend::PersistentTaskType, registry, FunctionId, TraitTypeId};
 
 #[derive(Default, Serialize)]
 struct TaskStatistics {
-    /// How many times the function was executed (roughly a cache miss)
+    /// How many times the function began execution (roughly a cache miss)
     execution_count: u32,
-    /// How many times the function was read (roughly a cache hit or a cache
-    /// miss)
+    /// How many times the function was read (either from cache, or after an
+    /// execution)
     finished_read_count: u32,
 }
 
@@ -50,23 +50,23 @@ impl AllTasksStatistics {
         self.inner.get().is_some()
     }
 
-    pub(crate) fn increment_execution_count(&self, task_type: &Arc<PersistentTaskType>) {
+    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_finished_read_count(&self, task_type: &Arc<PersistentTaskType>) {
+    pub(crate) fn increment_finished_read_count(&self, task_type: &PersistentTaskType) {
         self.with_task_type_statistics(task_type, |stats| stats.finished_read_count += 1)
     }
 
     fn with_task_type_statistics(
         &self,
-        task_type: &Arc<PersistentTaskType>,
+        task_type: &PersistentTaskType,
         func: impl Fn(&mut TaskStatistics),
     ) {
         if let Some(all_stats) = self.inner.get() {
             func(
                 all_stats
-                    .entry((&**task_type).into())
+                    .entry(task_type.into())
                     .or_insert(TaskStatistics::default())
                     .value_mut(),
             )

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

@@ -2,6 +2,7 @@
 
 use std::{future::Future, sync::Arc};
 
+use anyhow::Result;
 use once_cell::sync::Lazy;
 use regex::Regex;
 use serde_json::json;
@@ -12,24 +13,149 @@ use turbo_tasks_testing::register;
 register!();
 
 #[tokio::test]
-async fn statistics_counts() {
+async fn test_simple_task() {
     run_with_tt(|tt| async move {
-        for i in 0..100 {
-            let _ = double(i).await;
+        for i in 0..10 {
+            double(i).await.unwrap();
+            // use cached results
+            double(i).await.unwrap();
+        }
+        for i in 0..10 {
+            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,
+                },
+            })
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+async fn test_vc_receiving_task() {
+    run_with_tt(|tt| async move {
+        for i in 0..10 {
+            let dvc = double(i);
+            double_vc(dvc).await.unwrap();
+            // use cached results
+            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,
+                },
+                "turbo-tasks-memory::::double_vc": {
+                    "execution_count": 10,
+                    "finished_read_count": 20,
+                },
+            })
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+async fn test_trait_methods() {
+    run_with_tt(|tt| async move {
+        for i in 0..10 {
+            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::::double": {
-                    "execution_count": 100,
-                    "finished_read_count": 100,
-                }
+                "turbo-tasks-memory::::wrap": {
+                    "execution_count": 10,
+                    "finished_read_count": 20,
+                },
+                "turbo-tasks-memory::::WrappedU64::Doublable::double": {
+                    "execution_count": 10,
+                    "finished_read_count": 10,
+                },
+                "turbo-tasks-memory::::WrappedU64::Doublable::double_vc": {
+                    "execution_count": 10,
+                    "finished_read_count": 10,
+                },
             })
         );
     })
     .await;
 }
 
+// creates Vcs, but doesn't ever execute them
+#[tokio::test]
+async fn test_no_execution() {
+    run_with_tt(|tt| async move {
+        // don't await this!
+        let _ = wrap_vc(double_vc(double(123))).double().double_vc();
+        assert_eq!(
+            remove_hashes(serde_json::to_value(tt.backend().task_statistics()).unwrap()),
+            json!({})
+        );
+    })
+    .await;
+}
+
+// Internally, this function uses `PersistentTaskType::Native`.
+#[turbo_tasks::function]
+fn double(val: u64) -> Vc<u64> {
+    Vc::cell(val * 2)
+}
+
+// Internally, this function uses `PersistentTaskType::ResolveNative`.
+#[turbo_tasks::function]
+async fn double_vc(val: Vc<u64>) -> Result<Vc<u64>> {
+    let val = *val.await?;
+    Ok(Vc::cell(val * 2))
+}
+
+#[turbo_tasks::value]
+struct WrappedU64(u64);
+
+#[turbo_tasks::function]
+fn wrap(val: u64) -> Vc<WrappedU64> {
+    WrappedU64(val).cell()
+}
+
+#[turbo_tasks::function]
+async fn wrap_vc(val: Vc<u64>) -> Result<Vc<WrappedU64>> {
+    Ok(WrappedU64(*val.await?).cell())
+}
+
+#[turbo_tasks::value_trait]
+pub trait Doublable {
+    fn double(&self) -> Vc<Self>;
+    fn double_vc(self: Vc<Self>) -> Vc<Self>;
+}
+
+#[turbo_tasks::value_impl]
+impl Doublable for WrappedU64 {
+    #[turbo_tasks::function]
+    fn double(&self) -> Vc<Self> {
+        WrappedU64(self.0 * 2).cell()
+    }
+
+    #[turbo_tasks::function]
+    async fn double_vc(self: Vc<Self>) -> Result<Vc<Self>> {
+        let val = self.await?.0;
+        Ok(WrappedU64(val * 2).cell())
+    }
+}
+
+#[turbo_tasks::function]
+async fn fail(val: u64) -> Result<Vc<()>> {
+    anyhow::bail!("failed using {val}");
+}
+
 async fn run_with_tt<Fut>(func: impl FnOnce(Arc<TurboTasks<MemoryBackend>>) -> Fut)
 where
     Fut: Future<Output = ()> + Send + 'static,
@@ -46,11 +172,6 @@ where
     .unwrap();
 }
 
-#[turbo_tasks::function]
-async fn double(val: u64) -> Vc<u64> {
-    Vc::cell(val * 2)
-}
-
 // Global task identifiers can contain a hash of the crate and dependencies.
 // Remove that so that we can compare against a stable value in tests.
 fn remove_hashes(mut json: serde_json::Value) -> serde_json::Value {