From 38ec68d9900ca6db7209876e3bf6d87259e49395 Mon Sep 17 00:00:00 2001 From: Benjamin Woodruff Date: Mon, 5 Aug 2024 17:34:32 -0400 Subject: [PATCH] [turbopack] Implement resolution for local Vcs (#68472) *This is a migrated PR. This was in the turbo repository before the next.js merge.* **Migrated From:** https://github.com/vercel/turbo/pull/8826 ### Description Builds on the minimal implementation of local Vcs in #68469 At some point, we need to return local Vcs or pass them to another function as an argument. However, local Vcs are only valid within the context of the current task. The solution is that we need to convert local `Vc`s to non-local `Vc`s, and the easiest way to do that is to `.resolve()` it! This PR creates the new non-local cell on the current task, re-using the SharedReference the was used for the local cell, along with it's type id information. `RawVc` lacks information about the compile-time type `T`. Even `Vc` may know the real type of `T` if the `Vc` has been downcast to `Vc>`. To solve for this, we perform a lookup of the `VcCellMode::raw_cell` method using the type id (I added `raw_cell` in #68467). ### Testing Instructions ``` cargo nextest r -p turbo-tasks -p turbo-tasks-memory ``` --- .../turbo-tasks-memory/tests/local_cell.rs | 55 ++++++++++++++++++- turbopack/crates/turbo-tasks/src/backend.rs | 45 ++++++++++++++- turbopack/crates/turbo-tasks/src/manager.rs | 21 +++++-- .../crates/turbo-tasks/src/persisted_graph.rs | 4 +- turbopack/crates/turbo-tasks/src/raw_vc.rs | 31 ++++------- turbopack/crates/turbo-tasks/src/read_ref.rs | 2 +- .../turbo-tasks/src/task/shared_reference.rs | 21 +++++-- .../crates/turbo-tasks/src/value_type.rs | 24 +++++++- turbopack/crates/turbo-tasks/src/vc/mod.rs | 7 +-- 9 files changed, 169 insertions(+), 41 deletions(-) diff --git a/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs b/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs index e2f20b5501fc6..9274c789c6404 100644 --- a/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs +++ b/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs @@ -12,7 +12,7 @@ struct Wrapper(u32); struct TransparentWrapper(u32); #[tokio::test] -async fn store_and_read() { +async fn test_store_and_read() { run(®ISTRATION, async { let a: Vc = Vc::local_cell(42); assert_eq!(*a.await.unwrap(), 42); @@ -27,7 +27,7 @@ async fn store_and_read() { } #[tokio::test] -async fn store_and_read_generic() { +async fn test_store_and_read_generic() { run(®ISTRATION, async { // `Vc>>` is stored as `Vc>>` and requires special // transmute handling @@ -43,3 +43,54 @@ async fn store_and_read_generic() { }) .await } + +#[turbo_tasks::function] +async fn returns_resolved_local_vc() -> Vc { + Vc::::local_cell(42).resolve().await.unwrap() +} + +#[tokio::test] +async fn test_return_resolved() { + run(®ISTRATION, async { + assert_eq!(*returns_resolved_local_vc().await.unwrap(), 42); + }) + .await +} + +#[turbo_tasks::value(eq = "manual")] +#[derive(Default)] +struct Untracked { + #[turbo_tasks(debug_ignore, trace_ignore)] + #[serde(skip)] + cell: Vc, +} + +impl PartialEq for Untracked { + fn eq(&self, other: &Self) -> bool { + std::ptr::eq(self as *const _, other as *const _) + } +} + +impl Eq for Untracked {} + +#[turbo_tasks::function] +async fn get_untracked_local_cell() -> Vc { + Untracked { + cell: Vc::local_cell(42), + } + .cell() +} + +#[tokio::test] +#[should_panic(expected = "Local Vcs must only be accessed within their own task")] +async fn test_panics_on_local_cell_escape() { + run(®ISTRATION, async { + get_untracked_local_cell() + .await + .unwrap() + .cell + .await + .unwrap(); + }) + .await +} diff --git a/turbopack/crates/turbo-tasks/src/backend.rs b/turbopack/crates/turbo-tasks/src/backend.rs index 257aa71cca5bb..aeb61f180f728 100644 --- a/turbopack/crates/turbo-tasks/src/backend.rs +++ b/turbopack/crates/turbo-tasks/src/backend.rs @@ -20,6 +20,7 @@ use crate::{ manager::TurboTasksBackendApi, raw_vc::CellId, registry, + task::shared_reference::TypedSharedReference, trait_helpers::{get_trait_method, has_trait, traits}, triomphe_utils::unchecked_sidecast_triomphe_arc, FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdProvider, TaskIdSet, TraitRef, @@ -364,7 +365,7 @@ impl TypedCellContent { .1 .0 .ok_or_else(|| anyhow!("Cell is empty"))? - .typed(self.0); + .into_typed(self.0); Ok( // Safety: It is a TypedSharedReference TraitRef::new(shared_reference), @@ -376,12 +377,54 @@ impl TypedCellContent { } } +impl From for TypedCellContent { + fn from(value: TypedSharedReference) -> Self { + TypedCellContent(value.0, CellContent(Some(value.1))) + } +} + +impl TryFrom for TypedSharedReference { + type Error = TypedCellContent; + + fn try_from(content: TypedCellContent) -> Result { + if let TypedCellContent(type_id, CellContent(Some(shared_reference))) = content { + Ok(TypedSharedReference(type_id, shared_reference)) + } else { + Err(content) + } + } +} + impl CellContent { pub fn into_typed(self, type_id: ValueTypeId) -> TypedCellContent { TypedCellContent(type_id, self) } } +impl From for CellContent { + fn from(value: SharedReference) -> Self { + CellContent(Some(value)) + } +} + +impl From> for CellContent { + fn from(value: Option) -> Self { + CellContent(value) + } +} + +impl TryFrom for SharedReference { + type Error = CellContent; + + fn try_from(content: CellContent) -> Result { + if let CellContent(Some(shared_reference)) = content { + Ok(shared_reference) + } else { + Err(content) + } + } +} + pub type TaskCollectiblesMap = AutoMap, 1>; pub trait Backend: Sync + Send { diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs index dfbd6ce3ca65d..3feea84080492 100644 --- a/turbopack/crates/turbo-tasks/src/manager.rs +++ b/turbopack/crates/turbo-tasks/src/manager.rs @@ -35,6 +35,7 @@ use crate::{ magic_any::MagicAny, raw_vc::{CellId, RawVc}, registry, + task::shared_reference::TypedSharedReference, trace::TraceRawVcs, trait_helpers::get_trait_method, util::StaticOrArc, @@ -274,7 +275,7 @@ struct CurrentTaskState { /// Cells for locally allocated Vcs (`RawVc::LocalCell`). This is freed /// (along with `CurrentTaskState`) when the task finishes executing. - local_cells: Vec, + local_cells: Vec, } impl CurrentTaskState { @@ -1542,7 +1543,12 @@ pub(crate) async fn read_task_cell( } } -#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +/// A reference to a task's cell with methods that allow updating the contents +/// of the cell. +/// +/// Mutations should not outside of the task that that owns this cell. Doing so +/// is a logic error, and may lead to incorrect caching behavior. +#[derive(Clone, Copy, Serialize, Deserialize)] pub struct CurrentCellRef { current_task: TaskId, index: CellId, @@ -1719,7 +1725,7 @@ pub fn find_cell_by_type(ty: ValueTypeId) -> CurrentCellRef { }) } -pub(crate) fn create_local_cell(value: TypedCellContent) -> (ExecutionId, LocalCellId) { +pub(crate) fn create_local_cell(value: TypedSharedReference) -> (ExecutionId, LocalCellId) { CURRENT_TASK_STATE.with(|cell| { let CurrentTaskState { execution_id, @@ -1742,11 +1748,18 @@ pub(crate) fn create_local_cell(value: TypedCellContent) -> (ExecutionId, LocalC }) } +/// Returns the contents of the given local cell. Panics if a local cell is +/// attempted to be accessed outside of its task. +/// +/// Returns [`TypedSharedReference`] instead of [`TypedCellContent`] because +/// local cells are always filled. The returned value can be cheaply converted +/// with `.into()`. +/// /// Panics if the ExecutionId does not match the expected value. pub(crate) fn read_local_cell( execution_id: ExecutionId, local_cell_id: LocalCellId, -) -> TypedCellContent { +) -> TypedSharedReference { CURRENT_TASK_STATE.with(|cell| { let CurrentTaskState { execution_id: expected_execution_id, diff --git a/turbopack/crates/turbo-tasks/src/persisted_graph.rs b/turbopack/crates/turbo-tasks/src/persisted_graph.rs index e317dc903714b..14609f8b6071d 100644 --- a/turbopack/crates/turbo-tasks/src/persisted_graph.rs +++ b/turbopack/crates/turbo-tasks/src/persisted_graph.rs @@ -41,7 +41,7 @@ struct SerializableTaskCell(Option>); impl From for TaskCell { fn from(val: SerializableTaskCell) -> Self { match val.0 { - Some(d) => TaskCell::Content(CellContent(d.map(|d| d.untyped().1))), + Some(d) => TaskCell::Content(d.map(TypedSharedReference::into_untyped).into()), None => TaskCell::NeedComputation, } } @@ -56,7 +56,7 @@ impl Serialize for TaskCells { for (cell_id, cell) in &self.0 { let task_cell = SerializableTaskCell(match cell { TaskCell::Content(CellContent(opt)) => { - Some(opt.as_ref().map(|d| d.typed(cell_id.type_id))) + Some(opt.clone().map(|d| d.into_typed(cell_id.type_id))) } TaskCell::NeedComputation => None, }); diff --git a/turbopack/crates/turbo-tasks/src/raw_vc.rs b/turbopack/crates/turbo-tasks/src/raw_vc.rs index bf8185c2fd917..9dc748bb282f4 100644 --- a/turbopack/crates/turbo-tasks/src/raw_vc.rs +++ b/turbopack/crates/turbo-tasks/src/raw_vc.rs @@ -164,26 +164,15 @@ impl RawVc { /// See [`crate::Vc::resolve`]. pub(crate) async fn resolve(self) -> Result { - let tt = turbo_tasks(); - let mut current = self; - let mut notified = false; - loop { - match current { - RawVc::TaskOutput(task) => { - if !notified { - tt.notify_scheduled_tasks(); - notified = true; - } - current = read_task_output(&*tt, task, false).await?; - } - RawVc::TaskCell(_, _) => return Ok(current), - RawVc::LocalCell(_, _) => todo!(), - } - } + self.resolve_inner(/* strongly_consistent */ false).await } /// See [`crate::Vc::resolve_strongly_consistent`]. pub(crate) async fn resolve_strongly_consistent(self) -> Result { + self.resolve_inner(/* strongly_consistent */ true).await + } + + pub(crate) async fn resolve_inner(self, strongly_consistent: bool) -> Result { let tt = turbo_tasks(); let mut current = self; let mut notified = false; @@ -194,10 +183,14 @@ impl RawVc { tt.notify_scheduled_tasks(); notified = true; } - current = read_task_output(&*tt, task, true).await?; + current = read_task_output(&*tt, task, strongly_consistent).await?; } RawVc::TaskCell(_, _) => return Ok(current), - RawVc::LocalCell(_, _) => todo!(), + RawVc::LocalCell(execution_id, local_cell_id) => { + let shared_reference = read_local_cell(execution_id, local_cell_id); + let value_type = get_value_type(shared_reference.0); + return Ok((value_type.raw_cell)(shared_reference)); + } } } } @@ -355,7 +348,7 @@ impl Future for ReadRawVcFuture { } } RawVc::LocalCell(execution_id, local_cell_id) => { - return Poll::Ready(Ok(read_local_cell(execution_id, local_cell_id))); + return Poll::Ready(Ok(read_local_cell(execution_id, local_cell_id).into())); } }; // SAFETY: listener is from previous pinned this diff --git a/turbopack/crates/turbo-tasks/src/read_ref.rs b/turbopack/crates/turbo-tasks/src/read_ref.rs index b5525a29b4dc0..97f55bc2d2d74 100644 --- a/turbopack/crates/turbo-tasks/src/read_ref.rs +++ b/turbopack/crates/turbo-tasks/src/read_ref.rs @@ -251,7 +251,7 @@ where }; Vc { node: >::raw_cell( - SharedReference::new(value).typed(type_id), + SharedReference::new(value).into_typed(type_id), ), _t: PhantomData, } diff --git a/turbopack/crates/turbo-tasks/src/task/shared_reference.rs b/turbopack/crates/turbo-tasks/src/task/shared_reference.rs index 0194bfdfc597d..bb1e94dda1d03 100644 --- a/turbopack/crates/turbo-tasks/src/task/shared_reference.rs +++ b/turbopack/crates/turbo-tasks/src/task/shared_reference.rs @@ -2,6 +2,7 @@ use std::{ any::Any, fmt::{Debug, Display}, hash::Hash, + ops::Deref, }; use anyhow::Result; @@ -36,14 +37,26 @@ impl SharedReference { } } - pub(crate) fn typed(&self, type_id: ValueTypeId) -> TypedSharedReference { - TypedSharedReference(type_id, self.clone()) + pub fn downcast_ref(&self) -> Option<&T> { + self.0.downcast_ref() + } + + pub fn into_typed(self, type_id: ValueTypeId) -> TypedSharedReference { + TypedSharedReference(type_id, self) } } impl TypedSharedReference { - pub(crate) fn untyped(&self) -> (ValueTypeId, SharedReference) { - (self.0, self.1.clone()) + pub fn into_untyped(self) -> SharedReference { + self.1 + } +} + +impl Deref for TypedSharedReference { + type Target = SharedReference; + + fn deref(&self) -> &Self::Target { + &self.1 } } diff --git a/turbopack/crates/turbo-tasks/src/value_type.rs b/turbopack/crates/turbo-tasks/src/value_type.rs index d9d6e0246d8c2..067df22ca9b20 100644 --- a/turbopack/crates/turbo-tasks/src/value_type.rs +++ b/turbopack/crates/turbo-tasks/src/value_type.rs @@ -16,10 +16,14 @@ use crate::{ id::{FunctionId, TraitTypeId}, magic_any::{AnyDeserializeSeed, MagicAny, MagicAnyDeserializeSeed, MagicAnySerializeSeed}, registry::{register_trait_type, register_value_type}, + task::shared_reference::TypedSharedReference, + vc::VcCellMode, + RawVc, VcValueType, }; type MagicSerializationFn = fn(&dyn MagicAny) -> &dyn erased_serde::Serialize; type AnySerializationFn = fn(&(dyn Any + Sync + Send)) -> &dyn erased_serde::Serialize; +type RawCellFactoryFn = fn(TypedSharedReference) -> RawVc; // TODO this type need some refactoring when multiple languages are added to // turbo-task In this case a trait_method might be of a different function type. @@ -41,6 +45,17 @@ pub struct ValueType { /// Functors for serialization magic_serialization: Option<(MagicSerializationFn, MagicAnyDeserializeSeed)>, any_serialization: Option<(AnySerializationFn, AnyDeserializeSeed)>, + + /// An implementation of + /// [`VcCellMode::raw_cell`][crate::vc::cell_mode::VcCellMode::raw_cell]. + /// + /// Allows dynamically constructing a cell using the type id. Used inside of + /// [`RawVc`] where we have a type id, but not the concrete type `T` of + /// `Vc`. + /// + /// Because we allow resolving `Vc`, it's otherwise not possible + /// for `RawVc` to know what the appropriate `VcCellMode` is. + pub(crate) raw_cell: RawCellFactoryFn, } impl Hash for ValueType { @@ -88,19 +103,20 @@ pub fn any_as_serialize( impl ValueType { /// This is internally used by `#[turbo_tasks::value]` - pub fn new() -> Self { + pub fn new() -> Self { Self { name: std::any::type_name::().to_string(), traits: AutoSet::new(), trait_methods: AutoMap::new(), magic_serialization: None, any_serialization: None, + raw_cell: >::raw_cell, } } /// This is internally used by `#[turbo_tasks::value]` pub fn new_with_magic_serialization< - T: Debug + Eq + Hash + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static, + T: VcValueType + Debug + Eq + Hash + Serialize + for<'de> Deserialize<'de>, >() -> Self { Self { name: std::any::type_name::().to_string(), @@ -111,12 +127,13 @@ impl ValueType { MagicAnyDeserializeSeed::new::(), )), any_serialization: Some((any_as_serialize::, AnyDeserializeSeed::new::())), + raw_cell: >::raw_cell, } } /// This is internally used by `#[turbo_tasks::value]` pub fn new_with_any_serialization< - T: Any + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static, + T: VcValueType + Any + Serialize + for<'de> Deserialize<'de>, >() -> Self { Self { name: std::any::type_name::().to_string(), @@ -124,6 +141,7 @@ impl ValueType { trait_methods: AutoMap::new(), magic_serialization: None, any_serialization: Some((any_as_serialize::, AnyDeserializeSeed::new::())), + raw_cell: >::raw_cell, } } diff --git a/turbopack/crates/turbo-tasks/src/vc/mod.rs b/turbopack/crates/turbo-tasks/src/vc/mod.rs index 5c10f18769d49..539e787ddb0a1 100644 --- a/turbopack/crates/turbo-tasks/src/vc/mod.rs +++ b/turbopack/crates/turbo-tasks/src/vc/mod.rs @@ -25,7 +25,6 @@ pub use self::{ traits::{Dynamic, TypedForInput, Upcast, VcValueTrait, VcValueType}, }; use crate::{ - backend::CellContent, debug::{ValueDebug, ValueDebugFormat, ValueDebugFormatString}, manager::create_local_cell, registry, @@ -282,10 +281,8 @@ where // cells aren't stored across executions, so there can be no concept of // "updating" the cell across multiple executions. let (execution_id, local_cell_id) = create_local_cell( - CellContent(Some(SharedReference::new(triomphe::Arc::new( - T::Read::target_to_repr(inner), - )))) - .into_typed(T::get_value_type_id()), + SharedReference::new(triomphe::Arc::new(T::Read::target_to_repr(inner))) + .into_typed(T::get_value_type_id()), ); Vc { node: RawVc::LocalCell(execution_id, local_cell_id),