feat(turbo-tasks): Trace transient tasks

Author: bgw

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

@@ -5,6 +5,7 @@ mod storage;
 
 use std::{
     borrow::Cow,
+    fmt::{self, Write},
     future::Future,
     hash::BuildHasherDefault,
     mem::take,
@@ -18,6 +19,7 @@ use std::{
 
 use anyhow::{bail, Result};
 use auto_hash_map::{AutoMap, AutoSet};
+use indexmap::IndexSet;
 use parking_lot::{Condvar, Mutex};
 use rustc_hash::{FxHashMap, FxHashSet, FxHasher};
 use smallvec::smallvec;
@@ -28,8 +30,9 @@ use turbo_tasks::{
         TransientTaskType, TypedCellContent,
     },
     event::{Event, EventListener},
-    registry,
+    registry::{self, get_value_type_global_name},
     task_statistics::TaskStatisticsApi,
+    trace::TraceRawVcs,
     util::IdFactoryWithReuse,
     CellId, FunctionId, FxDashMap, RawVc, ReadCellOptions, ReadConsistency, SessionId, TaskId,
     TraitTypeId, TurboTasksBackendApi, ValueTypeId, TRANSIENT_TASK_BIT,
@@ -425,7 +428,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> Result<Result<RawVc, EventListener>> {
         if let Some(reader) = reader {
-            self.assert_not_persistent_calling_transient(reader, task_id);
+            self.assert_not_persistent_calling_transient(reader, task_id, /* cell_id */ None);
         }
 
         let mut ctx = self.execute_context(turbo_tasks);
@@ -625,7 +628,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> Result<Result<TypedCellContent, EventListener>> {
         if let Some(reader) = reader {
-            self.assert_not_persistent_calling_transient(reader, task_id);
+            self.assert_not_persistent_calling_transient(reader, task_id, Some(cell));
         }
 
         fn add_cell_dependency<B: BackingStorage>(
@@ -985,9 +988,10 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> TaskId {
         if !parent_task.is_transient() {
-            panic_persistent_calling_transient(
+            self.panic_persistent_calling_transient(
                 self.lookup_task_type(parent_task).as_deref(),
                 Some(&task_type),
+                /* cell_id */ None,
             );
         }
         if let Some(task_id) = self.task_cache.lookup_forward(&task_type) {
@@ -1013,6 +1017,84 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         task_id
     }
 
+    /// Generate an object that implements [`fmt::Display`] explaining why the given
+    /// [`CachedTaskType`] is transient.
+    fn debug_trace_transient_task(
+        &self,
+        task_type: &CachedTaskType,
+        cell_id: Option<CellId>,
+    ) -> DebugTraceTransientTask {
+        // it shouldn't be possible to have cycles in tasks, but we could have an exponential blowup
+        // from tracing the same task many times, so use a visited_set
+        fn inner_id(
+            backend: &TurboTasksBackendInner<impl BackingStorage>,
+            task_id: TaskId,
+            cell_type_id: Option<ValueTypeId>,
+            visited_set: &mut FxHashSet<TaskId>,
+        ) -> DebugTraceTransientTask {
+            if let Some(task_type) = backend.lookup_task_type(task_id) {
+                if visited_set.contains(&task_id) {
+                    let task_name = task_type.get_name();
+                    DebugTraceTransientTask::Collapsed {
+                        task_name,
+                        cell_type_id,
+                    }
+                } else {
+                    inner_cached(backend, &task_type, cell_type_id, visited_set)
+                }
+            } else {
+                DebugTraceTransientTask::Uncached { cell_type_id }
+            }
+        }
+        fn inner_cached(
+            backend: &TurboTasksBackendInner<impl BackingStorage>,
+            task_type: &CachedTaskType,
+            cell_type_id: Option<ValueTypeId>,
+            visited_set: &mut FxHashSet<TaskId>,
+        ) -> DebugTraceTransientTask {
+            let task_name = task_type.get_name();
+
+            let cause_self = task_type.this.and_then(|cause_self_raw_vc| {
+                let task_id = cause_self_raw_vc.get_task_id();
+                if task_id.is_transient() {
+                    Some(Box::new(inner_id(
+                        backend,
+                        cause_self_raw_vc.get_task_id(),
+                        cause_self_raw_vc.try_get_type_id(),
+                        visited_set,
+                    )))
+                } else {
+                    None
+                }
+            });
+            let cause_args = task_type
+                .arg
+                .get_raw_vcs()
+                .into_iter()
+                .map(|raw_vc| (raw_vc.get_task_id(), raw_vc.try_get_type_id()))
+                .filter(|(task_id, _)| task_id.is_transient())
+                .collect::<IndexSet<_>>() // dedupe
+                .into_iter()
+                .map(|(task_id, cell_type_id)| {
+                    inner_id(backend, task_id, cell_type_id, visited_set)
+                })
+                .collect();
+
+            DebugTraceTransientTask::Cached {
+                task_name,
+                cell_type_id,
+                cause_self,
+                cause_args,
+            }
+        }
+        inner_cached(
+            self,
+            task_type,
+            cell_id.map(|c| c.type_id),
+            &mut FxHashSet::default(),
+        )
+    }
+
     fn invalidate_task(
         &self,
         task_id: TaskId,
@@ -1987,14 +2069,42 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         }
     }
 
-    fn assert_not_persistent_calling_transient(&self, parent_id: TaskId, child_id: TaskId) {
+    fn assert_not_persistent_calling_transient(
+        &self,
+        parent_id: TaskId,
+        child_id: TaskId,
+        cell_id: Option<CellId>,
+    ) {
         if !parent_id.is_transient() && child_id.is_transient() {
-            panic_persistent_calling_transient(
+            self.panic_persistent_calling_transient(
                 self.lookup_task_type(parent_id).as_deref(),
                 self.lookup_task_type(child_id).as_deref(),
+                cell_id,
             );
         }
     }
+
+    fn panic_persistent_calling_transient(
+        &self,
+        parent: Option<&CachedTaskType>,
+        child: Option<&CachedTaskType>,
+        cell_id: Option<CellId>,
+    ) {
+        let transient_reason = if let Some(child) = child {
+            format!(
+                " The callee is transient because it depends on:\n{}",
+                self.debug_trace_transient_task(child, cell_id),
+            )
+        } else {
+            String::new()
+        };
+        panic!(
+            "Persistent task {} is not allowed to call or read transient tasks {}.{}",
+            parent.map_or("unknown", |t| t.get_name()),
+            child.map_or("unknown", |t| t.get_name()),
+            transient_reason,
+        );
+    }
 }
 
 impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
@@ -2272,15 +2382,90 @@ impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
     }
 }
 
-fn panic_persistent_calling_transient(
-    parent: Option<&CachedTaskType>,
-    child: Option<&CachedTaskType>,
-) {
-    panic!(
-        "Persistent task {} is not allowed to call or read transient tasks {}",
-        parent.map_or("unknown", |t| t.get_name()),
-        child.map_or("unknown", |t| t.get_name()),
-    );
+enum DebugTraceTransientTask {
+    Cached {
+        task_name: &'static str,
+        cell_type_id: Option<ValueTypeId>,
+        cause_self: Option<Box<DebugTraceTransientTask>>,
+        cause_args: Vec<DebugTraceTransientTask>,
+    },
+    /// This representation is used when this task is a duplicate of one previously shown
+    Collapsed {
+        task_name: &'static str,
+        cell_type_id: Option<ValueTypeId>,
+    },
+    Uncached {
+        cell_type_id: Option<ValueTypeId>,
+    },
+}
+
+impl DebugTraceTransientTask {
+    fn fmt_indented(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
+        let indent = "    ".repeat(level);
+        f.write_str(&indent)?;
+
+        fn fmt_cell_type_id(
+            f: &mut fmt::Formatter<'_>,
+            cell_type_id: Option<ValueTypeId>,
+        ) -> fmt::Result {
+            if let Some(ty) = cell_type_id {
+                write!(f, " (read cell of type {})", get_value_type_global_name(ty))
+            } else {
+                Ok(())
+            }
+        }
+
+        // write the name and type
+        match self {
+            Self::Cached {
+                task_name,
+                cell_type_id,
+                ..
+            }
+            | Self::Collapsed {
+                task_name,
+                cell_type_id,
+                ..
+            } => {
+                f.write_str(task_name)?;
+                fmt_cell_type_id(f, *cell_type_id)?;
+                if matches!(self, Self::Collapsed { .. }) {
+                    f.write_str(" (collapsed)")?;
+                }
+            }
+            Self::Uncached { cell_type_id } => {
+                f.write_str("unknown transient task")?;
+                fmt_cell_type_id(f, *cell_type_id)?;
+            }
+        }
+        f.write_char('\n')?;
+
+        // write any extra "cause" information we might have
+        if let Self::Cached {
+            cause_self,
+            cause_args,
+            ..
+        } = self
+        {
+            if let Some(c) = cause_self {
+                writeln!(f, "{indent}  self:")?;
+                c.fmt_indented(f, level + 1)?;
+            }
+            if !cause_args.is_empty() {
+                writeln!(f, "{indent}  args:")?;
+                for c in cause_args {
+                    c.fmt_indented(f, level + 1)?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+impl fmt::Display for DebugTraceTransientTask {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_indented(f, 0)
+    }
 }
 
 // from https://github.com/tokio-rs/tokio/blob/29cd6ec1ec6f90a7ee1ad641c03e0e00badbcb0e/tokio/src/time/instant.rs#L57-L63

turbopack/crates/turbo-tasks-backend/tests/trace_transient.rs

@@ -0,0 +1,74 @@
+#![feature(arbitrary_self_types)]
+#![feature(arbitrary_self_types_pointers)]
+
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use turbo_tasks::{trace::TraceRawVcs, NonLocalValue, ResolvedVc, TaskInput, Vc};
+use turbo_tasks_testing::{register, run_without_cache_check, Registration};
+
+static REGISTRATION: Registration = register!();
+
+const EXPECTED_TRACE: &str = "\
+Adder::add_method (read cell of type turbo-tasks@TODO::::primitives::u64)
+  self:
+    Adder::new (read cell of type turbo-tasks-backend@TODO::::Adder)
+      args:
+        unknown transient task (read cell of type turbo-tasks@TODO::::primitives::unit)
+  args:
+    unknown transient task (read cell of type turbo-tasks@TODO::::primitives::u16)
+    unknown transient task (read cell of type turbo-tasks@TODO::::primitives::u32)";
+
+#[tokio::test]
+async fn test_trace_transient() {
+    let result = run_without_cache_check(&REGISTRATION, async {
+        read_incorrect_task_input_operation(IncorrectTaskInput(
+            Adder::new(Vc::cell(()))
+                .add_method(Vc::cell(2), Vc::cell(3))
+                .to_resolved()
+                .await?,
+        ))
+        .read_strongly_consistent()
+        .await?;
+        anyhow::Ok(())
+    })
+    .await;
+    assert!(result
+        .unwrap_err()
+        .to_string()
+        .contains(&EXPECTED_TRACE.escape_debug().to_string()));
+}
+
+#[turbo_tasks::value]
+struct Adder;
+
+#[turbo_tasks::value_impl]
+impl Adder {
+    #[turbo_tasks::function]
+    fn new(arg: ResolvedVc<()>) -> Vc<Adder> {
+        let _ = arg; // Make sure unused argument filtering doesn't remove the arg
+        Adder.cell()
+    }
+
+    #[turbo_tasks::function]
+    async fn add_method(&self, arg1: ResolvedVc<u16>, arg2: ResolvedVc<u32>) -> Result<Vc<u64>> {
+        Ok(Vc::cell(u64::from(*arg1.await?) + u64::from(*arg2.await?)))
+    }
+}
+
+#[turbo_tasks::function(operation)]
+async fn read_incorrect_task_input_operation(value: IncorrectTaskInput) -> Result<Vc<u64>> {
+    Ok(Vc::cell(*value.0.await?))
+}
+
+/// Has an intentionally incorrect `TaskInput` implementation, representing some code that the debug
+/// tracing might be particularly useful with.
+#[derive(
+    Copy, Clone, Debug, PartialEq, Eq, Hash, TraceRawVcs, Serialize, Deserialize, NonLocalValue,
+)]
+struct IncorrectTaskInput(ResolvedVc<u64>);
+
+impl TaskInput for IncorrectTaskInput {
+    fn is_transient(&self) -> bool {
+        false
+    }
+}

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

@@ -68,6 +68,8 @@ pub struct CachedTaskType {
 }
 
 impl CachedTaskType {
+    /// Get the name of the function from the registry. Equivalent to the
+    /// [`fmt::Display::to_string`] implementation, but does not allocate a `String`.
     pub fn get_name(&self) -> &'static str {
         &registry::get_function(self.fn_type).name
     }

turbopack/crates/turbo-tasks/src/raw_vc.rs

@@ -216,6 +216,13 @@ impl RawVc {
         }
     }
 
+    pub fn try_get_type_id(&self) -> Option<ValueTypeId> {
+        match self {
+            RawVc::TaskCell(_, CellId { type_id, .. }) => Some(*type_id),
+            RawVc::TaskOutput(_) | RawVc::LocalOutput(_, _) => None,
+        }
+    }
+
     /// For a cell that's already resolved, synchronously check if it implements a trait using the
     /// type information in `RawVc::TaskCell` (we don't actualy need to read the cell!).
     pub(crate) fn resolved_has_trait(&self, trait_id: TraitTypeId) -> bool {