From 659d2412a3c67aa1511d7baa3da58fb781e9e6c5 Mon Sep 17 00:00:00 2001 From: Benjamin Woodruff Date: Mon, 5 Aug 2024 18:20:47 -0400 Subject: [PATCH] [turbopack] Implement remainder of local Vcs: casting and get_task_id (#68474) *This is a migrated PR. This was in the turbo repository before the next.js merge.* **Migrated From:** https://github.com/vercel/turbo/pull/8871 ## Description With these changes, local Vcs (introduced in #68469) should have a fully functional implementation. These remaining methods were pretty straightforward. After this, my focus will shift back towards the changes needed to tasks to allow us to make use of local Vcs. ## New `resolve_type_inner` helper method I merged the implementations of `resolve_trait` and `resolve_value`, since their implementations were 95% identical. I think this is an overall win, though the logic to prevent duplicate `ValueType` lookups (they're not expensive, but this code is also potentially very hot) is a bit messy. ## Testing Instructions ``` cargo nextest r -p turbo-tasks -p turbo-tasks-memory ``` --- .../turbo-tasks-memory/tests/local_cell.rs | 64 +++++++++++++- turbopack/crates/turbo-tasks/src/manager.rs | 29 +++++-- turbopack/crates/turbo-tasks/src/raw_vc.rs | 85 +++++++++++-------- 3 files changed, 133 insertions(+), 45 deletions(-) diff --git a/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs b/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs index 9274c789c6404..7ca3fa7103d94 100644 --- a/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs +++ b/turbopack/crates/turbo-tasks-memory/tests/local_cell.rs @@ -1,6 +1,6 @@ #![feature(arbitrary_self_types)] -use turbo_tasks::Vc; +use turbo_tasks::{debug::ValueDebug, test_helpers::current_task_for_testing, ValueDefault, Vc}; use turbo_tasks_testing::{register, run, Registration}; static REGISTRATION: Registration = register!(); @@ -57,6 +57,57 @@ async fn test_return_resolved() { .await } +#[turbo_tasks::value_trait] +trait UnimplementedTrait {} + +#[tokio::test] +async fn test_try_resolve_sidecast() { + run(®ISTRATION, async { + let trait_vc: Vc> = Vc::upcast(Vc::::local_cell(42)); + + // `u32` is both a `ValueDebug` and a `ValueDefault`, so this sidecast is valid + let sidecast_vc = Vc::try_resolve_sidecast::>(trait_vc) + .await + .unwrap(); + assert!(sidecast_vc.is_some()); + + // `u32` is not an `UnimplementedTrait` though, so this should return None + let wrongly_sidecast_vc = Vc::try_resolve_sidecast::>(trait_vc) + .await + .unwrap(); + assert!(wrongly_sidecast_vc.is_none()); + }) + .await +} + +#[tokio::test] +async fn test_try_resolve_downcast_type() { + run(®ISTRATION, async { + let trait_vc: Vc> = Vc::upcast(Vc::::local_cell(42)); + + let downcast_vc: Vc = Vc::try_resolve_downcast_type(trait_vc) + .await + .unwrap() + .unwrap(); + assert_eq!(*downcast_vc.await.unwrap(), 42); + + let wrongly_downcast_vc: Option> = + Vc::try_resolve_downcast_type(trait_vc).await.unwrap(); + assert!(wrongly_downcast_vc.is_none()); + }) + .await +} + +#[tokio::test] +async fn test_get_task_id() { + run(®ISTRATION, async { + // the task id as reported by the RawVc + let vc_task_id = Vc::into_raw(Vc::<()>::local_cell(())).get_task_id(); + assert_eq!(vc_task_id, current_task_for_testing()); + }) + .await +} + #[turbo_tasks::value(eq = "manual")] #[derive(Default)] struct Untracked { @@ -83,7 +134,7 @@ async fn get_untracked_local_cell() -> Vc { #[tokio::test] #[should_panic(expected = "Local Vcs must only be accessed within their own task")] -async fn test_panics_on_local_cell_escape() { +async fn test_panics_on_local_cell_escape_read() { run(®ISTRATION, async { get_untracked_local_cell() .await @@ -94,3 +145,12 @@ async fn test_panics_on_local_cell_escape() { }) .await } + +#[tokio::test] +#[should_panic(expected = "Local Vcs must only be accessed within their own task")] +async fn test_panics_on_local_cell_escape_get_task_id() { + run(®ISTRATION, async { + Vc::into_raw(get_untracked_local_cell().await.unwrap().cell).get_task_id(); + }) + .await +} diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs index 3feea84080492..80d2a18c3e30a 100644 --- a/turbopack/crates/turbo-tasks/src/manager.rs +++ b/turbopack/crates/turbo-tasks/src/manager.rs @@ -1755,7 +1755,8 @@ pub(crate) fn create_local_cell(value: TypedSharedReference) -> (ExecutionId, Lo /// local cells are always filled. The returned value can be cheaply converted /// with `.into()`. /// -/// Panics if the ExecutionId does not match the expected value. +/// Panics if the [`ExecutionId`] does not match the current task's +/// `execution_id`. pub(crate) fn read_local_cell( execution_id: ExecutionId, local_cell_id: LocalCellId, @@ -1766,12 +1767,28 @@ pub(crate) fn read_local_cell( local_cells, .. } = &*cell.borrow(); - assert_eq!( - execution_id, *expected_execution_id, - "This Vc is local. Local Vcs must only be accessed within their own task. Resolve the \ - Vc to convert it into a non-local version." - ); + assert_eq_local_cell(execution_id, *expected_execution_id); // local cell ids are one-indexed (they use NonZeroU32) local_cells[(*local_cell_id as usize) - 1].clone() }) } + +/// Panics if the [`ExecutionId`] does not match the current task's +/// `execution_id`. +pub(crate) fn assert_execution_id(execution_id: ExecutionId) { + CURRENT_TASK_STATE.with(|cell| { + let CurrentTaskState { + execution_id: expected_execution_id, + .. + } = &*cell.borrow(); + assert_eq_local_cell(execution_id, *expected_execution_id); + }) +} + +fn assert_eq_local_cell(actual: ExecutionId, expected: ExecutionId) { + assert_eq!( + actual, expected, + "This Vc is local. Local Vcs must only be accessed within their own task. Resolve the Vc \ + to convert it into a non-local version." + ); +} diff --git a/turbopack/crates/turbo-tasks/src/raw_vc.rs b/turbopack/crates/turbo-tasks/src/raw_vc.rs index 9dc748bb282f4..a0f4c3470b898 100644 --- a/turbopack/crates/turbo-tasks/src/raw_vc.rs +++ b/turbopack/crates/turbo-tasks/src/raw_vc.rs @@ -16,9 +16,12 @@ use crate::{ backend::{CellContent, TypedCellContent}, event::EventListener, id::{ExecutionId, LocalCellId}, - manager::{read_local_cell, read_task_cell, read_task_output, TurboTasksApi}, + manager::{ + assert_execution_id, current_task, read_local_cell, read_task_cell, read_task_output, + TurboTasksApi, + }, registry::{self, get_value_type}, - turbo_tasks, CollectiblesSource, TaskId, TraitTypeId, ValueTypeId, Vc, VcValueTrait, + turbo_tasks, CollectiblesSource, TaskId, TraitTypeId, ValueType, ValueTypeId, Vc, VcValueTrait, }; #[derive(Error, Debug)] @@ -100,38 +103,31 @@ impl RawVc { self, trait_type: TraitTypeId, ) -> Result, ResolveTypeError> { - let tt = turbo_tasks(); - tt.notify_scheduled_tasks(); - let mut current = self; - loop { - match current { - RawVc::TaskOutput(task) => { - current = read_task_output(&*tt, task, false) - .await - .map_err(|source| ResolveTypeError::TaskError { source })?; - } - RawVc::TaskCell(task, index) => { - let content = read_task_cell(&*tt, task, index) - .await - .map_err(|source| ResolveTypeError::ReadError { source })?; - if let TypedCellContent(value_type, CellContent(Some(_))) = content { - if get_value_type(value_type).has_trait(&trait_type) { - return Ok(Some(RawVc::TaskCell(task, index))); - } else { - return Ok(None); - } - } else { - return Err(ResolveTypeError::NoContent); - } - } - RawVc::LocalCell(_, _) => todo!(), - } - } + self.resolve_type_inner(|value_type_id| { + let value_type = get_value_type(value_type_id); + (value_type.has_trait(&trait_type), Some(value_type)) + }) + .await } pub(crate) async fn resolve_value( self, value_type: ValueTypeId, + ) -> Result, ResolveTypeError> { + self.resolve_type_inner(|cell_value_type| (cell_value_type == value_type, None)) + .await + } + + /// Helper for `resolve_trait` and `resolve_value`. + /// + /// After finding a cell, returns `Ok(Some(...))` when `conditional` returns + /// `true`, and `Ok(None)` when `conditional` returns `false`. + /// + /// As an optimization, `conditional` may return the `&'static ValueType` to + /// avoid a potential extra lookup later. + async fn resolve_type_inner( + self, + conditional: impl FnOnce(ValueTypeId) -> (bool, Option<&'static ValueType>), ) -> Result, ResolveTypeError> { let tt = turbo_tasks(); tt.notify_scheduled_tasks(); @@ -147,17 +143,29 @@ impl RawVc { let content = read_task_cell(&*tt, task, index) .await .map_err(|source| ResolveTypeError::ReadError { source })?; - if let TypedCellContent(cell_value_type, CellContent(Some(_))) = content { - if cell_value_type == value_type { - return Ok(Some(RawVc::TaskCell(task, index))); + if let TypedCellContent(value_type, CellContent(Some(_))) = content { + return Ok(if conditional(value_type).0 { + Some(RawVc::TaskCell(task, index)) } else { - return Ok(None); - } + None + }); } else { return Err(ResolveTypeError::NoContent); } } - RawVc::LocalCell(_, _) => todo!(), + RawVc::LocalCell(execution_id, local_cell_id) => { + let shared_reference = read_local_cell(execution_id, local_cell_id); + return Ok( + if let (true, value_type) = conditional(shared_reference.0) { + // re-use the `ValueType` lookup from `conditional`, if it exists + let value_type = + value_type.unwrap_or_else(|| get_value_type(shared_reference.0)); + Some((value_type.raw_cell)(shared_reference)) + } else { + None + }, + ); + } } } } @@ -172,7 +180,7 @@ impl RawVc { self.resolve_inner(/* strongly_consistent */ true).await } - pub(crate) async fn resolve_inner(self, strongly_consistent: bool) -> Result { + async fn resolve_inner(self, strongly_consistent: bool) -> Result { let tt = turbo_tasks(); let mut current = self; let mut notified = false; @@ -203,7 +211,10 @@ impl RawVc { pub fn get_task_id(&self) -> TaskId { match self { RawVc::TaskOutput(t) | RawVc::TaskCell(t, _) => *t, - RawVc::LocalCell(_, _) => todo!(), + RawVc::LocalCell(execution_id, _) => { + assert_execution_id(*execution_id); + current_task("RawVc::get_task_id") + } } } }