From 0151e49af5de338e4c1794b8565f1481e164d291 Mon Sep 17 00:00:00 2001 From: Benjamin Woodruff Date: Thu, 23 Jan 2025 15:44:58 -0800 Subject: [PATCH] perf(turbo-tasks): Add a 'local' option to `#[turbo_tasks::function(..)]` --- .../turbo-tasks-backend/tests/local_tasks.rs | 1 + .../fail_attribute_invalid_args.stderr | 2 +- ...ttribute_invalid_args_inherent_impl.stderr | 2 +- .../crates/turbo-tasks-macros/src/func.rs | 57 +++++++++---------- .../turbo-tasks-macros/src/function_macro.rs | 12 ++-- .../src/value_impl_macro.rs | 24 ++++---- .../src/value_trait_macro.rs | 17 +++--- .../turbo-tasks-memory/tests/local_tasks.rs | 1 + .../turbo-tasks-testing/tests/local_tasks.rs | 34 +++++++++++ turbopack/crates/turbo-tasks/src/manager.rs | 16 +++--- .../crates/turbo-tasks/src/native_function.rs | 10 +++- .../crates/turbo-tasks/src/task/local_task.rs | 6 +- 12 files changed, 113 insertions(+), 69 deletions(-) create mode 120000 turbopack/crates/turbo-tasks-backend/tests/local_tasks.rs create mode 120000 turbopack/crates/turbo-tasks-memory/tests/local_tasks.rs create mode 100644 turbopack/crates/turbo-tasks-testing/tests/local_tasks.rs diff --git a/turbopack/crates/turbo-tasks-backend/tests/local_tasks.rs b/turbopack/crates/turbo-tasks-backend/tests/local_tasks.rs new file mode 120000 index 00000000000000..29d9f86f455662 --- /dev/null +++ b/turbopack/crates/turbo-tasks-backend/tests/local_tasks.rs @@ -0,0 +1 @@ +../../turbo-tasks-testing/tests/local_tasks.rs \ No newline at end of file diff --git a/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args.stderr b/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args.stderr index 32bcaa2595082d..fefa9ab1b8efd4 100644 --- a/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args.stderr +++ b/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args.stderr @@ -1,4 +1,4 @@ -error: unexpected token, expected one of: "fs", "network", "operation", or "local_cells" +error: unexpected token, expected one of: "fs", "network", "operation", "local", or "local_cells" --> tests/function/fail_attribute_invalid_args.rs:9:25 | 9 | #[turbo_tasks::function(invalid_argument)] diff --git a/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args_inherent_impl.stderr b/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args_inherent_impl.stderr index 46acba8b7bc688..0b7129e2d1b9cc 100644 --- a/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args_inherent_impl.stderr +++ b/turbopack/crates/turbo-tasks-macros-tests/tests/function/fail_attribute_invalid_args_inherent_impl.stderr @@ -1,4 +1,4 @@ -error: unexpected token, expected one of: "fs", "network", "operation", or "local_cells" +error: unexpected token, expected one of: "fs", "network", "operation", "local", or "local_cells" --> tests/function/fail_attribute_invalid_args_inherent_impl.rs:14:29 | 14 | #[turbo_tasks::function(invalid_argument)] diff --git a/turbopack/crates/turbo-tasks-macros/src/func.rs b/turbopack/crates/turbo-tasks-macros/src/func.rs index 52d0e24f3d44d5..2d088ad5076360 100644 --- a/turbopack/crates/turbo-tasks-macros/src/func.rs +++ b/turbopack/crates/turbo-tasks-macros/src/func.rs @@ -31,7 +31,7 @@ pub struct TurboFn<'a> { /// Should we return `OperationVc` and require that all arguments are `NonLocalValue`s? operation: bool, /// Should this function use `TaskPersistence::LocalCells`? - local_cells: bool, + local: bool, } #[derive(Debug)] @@ -274,7 +274,7 @@ impl TurboFn<'_> { this, inputs, operation: args.operation.is_some(), - local_cells: args.local_cells.is_some(), + local: args.local.is_some(), inline_ident, }) } @@ -479,9 +479,9 @@ impl TurboFn<'_> { } pub fn persistence(&self) -> impl ToTokens { - if self.local_cells { + if self.local { quote! { - turbo_tasks::TaskPersistence::LocalCells + turbo_tasks::TaskPersistence::Local } } else { quote! { @@ -491,9 +491,9 @@ impl TurboFn<'_> { } pub fn persistence_with_this(&self) -> impl ToTokens { - if self.local_cells { + if self.local { quote! { - turbo_tasks::TaskPersistence::LocalCells + turbo_tasks::TaskPersistence::Local } } else { quote! { @@ -703,6 +703,10 @@ pub struct FunctionArguments { /// /// If there is an error due to this option being set, it should be reported to this span. operation: Option, + /// Does not run the function as a real task, and instead runs it inside the parent task using + /// task-local state. The function call itself will not be cached, but cells will be created on + /// the parent task. + pub local: Option, /// Changes the behavior of `Vc::cell` to create local cells that are not cached across task /// executions. Cells can be converted to their non-local versions by calling `Vc::resolve`. /// @@ -732,23 +736,29 @@ impl Parse for FunctionArguments { ("operation", Meta::Path(_)) => { parsed_args.operation = Some(meta.span()); } + ("local", Meta::Path(_)) => { + parsed_args.local = Some(meta.span()); + } ("local_cells", Meta::Path(_)) => { - let span = Some(meta.span()); - parsed_args.local_cells = span; + parsed_args.local_cells = Some(meta.span()); } (_, meta) => { return Err(syn::Error::new_spanned( meta, "unexpected token, expected one of: \"fs\", \"network\", \"operation\", \ - or \"local_cells\"", + \"local\", or \"local_cells\"", )) } } } - if let (Some(_), Some(span)) = (parsed_args.local_cells, parsed_args.operation) { + if let (Some(_), Some(span)) = ( + parsed_args.local.or(parsed_args.local_cells), + parsed_args.operation, + ) { return Err(syn::Error::new( span, - "\"operation\" is mutually exclusive with the \"local_cells\" option", + "\"operation\" is mutually exclusive with the \"local\" and \"local_cells\" \ + options", )); } Ok(parsed_args) @@ -1023,27 +1033,14 @@ impl DefinitionContext { #[derive(Debug)] pub struct NativeFn { - function_path_string: String, - function_path: ExprPath, - is_method: bool, - local_cells: bool, + pub function_path_string: String, + pub function_path: ExprPath, + pub is_method: bool, + pub local: bool, + pub local_cells: bool, } impl NativeFn { - pub fn new( - function_path_string: &str, - function_path: &ExprPath, - is_method: bool, - local_cells: bool, - ) -> NativeFn { - NativeFn { - function_path_string: function_path_string.to_owned(), - function_path: function_path.clone(), - is_method, - local_cells, - } - } - pub fn ty(&self) -> Type { parse_quote! { turbo_tasks::NativeFunction } } @@ -1053,6 +1050,7 @@ impl NativeFn { function_path_string, function_path, is_method, + local, local_cells, } = self; @@ -1068,6 +1066,7 @@ impl NativeFn { turbo_tasks::NativeFunction::#constructor( #function_path_string.to_owned(), turbo_tasks::FunctionMeta { + local: #local, local_cells: #local_cells, }, #function_path, diff --git a/turbopack/crates/turbo-tasks-macros/src/function_macro.rs b/turbopack/crates/turbo-tasks-macros/src/function_macro.rs index 33599492b8213c..0c5f5064b71f25 100644 --- a/turbopack/crates/turbo-tasks-macros/src/function_macro.rs +++ b/turbopack/crates/turbo-tasks-macros/src/function_macro.rs @@ -39,6 +39,7 @@ pub fn function(args: TokenStream, input: TokenStream) -> TokenStream { let args = syn::parse::(args) .inspect_err(|err| errors.push(err.to_compile_error())) .unwrap_or_default(); + let local = args.local.is_some(); let local_cells = args.local_cells.is_some(); let Some(turbo_fn) = TurboFn::new(&sig, DefinitionContext::NakedFn, args) else { @@ -54,12 +55,13 @@ pub fn function(args: TokenStream, input: TokenStream) -> TokenStream { let (inline_signature, inline_block) = turbo_fn.inline_signature_and_block(&block); let inline_attrs = filter_inline_attributes(&attrs[..]); - let native_fn = NativeFn::new( - &ident.to_string(), - &parse_quote! { #inline_function_ident }, - turbo_fn.is_method(), + let native_fn = NativeFn { + function_path_string: ident.to_string(), + function_path: parse_quote! { #inline_function_ident }, + is_method: turbo_fn.is_method(), + local, local_cells, - ); + }; let native_function_ident = get_native_function_ident(ident); let native_function_ty = native_fn.ty(); let native_function_def = native_fn.definition(); diff --git a/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs b/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs index 6663306886b650..c471ebc1572509 100644 --- a/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs +++ b/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs @@ -122,6 +122,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream { let func_args = func_args .inspect_err(|err| errors.push(err.to_compile_error())) .unwrap_or_default(); + let local = func_args.local.is_some(); let local_cells = func_args.local_cells.is_some(); let Some(turbo_fn) = @@ -135,12 +136,13 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream { let (inline_signature, inline_block) = turbo_fn.inline_signature_and_block(block); let inline_attrs = filter_inline_attributes(attrs.iter().copied()); - let native_fn = NativeFn::new( - &format!("{ty}::{ident}", ty = ty.to_token_stream()), - &parse_quote! { <#ty>::#inline_function_ident }, - turbo_fn.is_method(), + let native_fn = NativeFn { + function_path_string: format!("{ty}::{ident}", ty = ty.to_token_stream()), + function_path: parse_quote! { <#ty>::#inline_function_ident }, + is_method: turbo_fn.is_method(), + local, local_cells, - ); + }; let native_function_ident = get_inherent_impl_function_ident(ty_ident, ident); let native_function_ty = native_fn.ty(); @@ -221,6 +223,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream { let func_args = func_args .inspect_err(|err| errors.push(err.to_compile_error())) .unwrap_or_default(); + let local = func_args.local.is_some(); let local_cells = func_args.local_cells.is_some(); let Some(turbo_fn) = @@ -239,18 +242,19 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream { let (inline_signature, inline_block) = turbo_fn.inline_signature_and_block(block); let inline_attrs = filter_inline_attributes(attrs.iter().copied()); - let native_fn = NativeFn::new( - &format!( + let native_fn = NativeFn { + function_path_string: format!( "<{ty} as {trait_path}>::{ident}", ty = ty.to_token_stream(), trait_path = trait_path.to_token_stream() ), - &parse_quote! { + function_path: parse_quote! { <#ty as #inline_extension_trait_ident>::#inline_function_ident }, - turbo_fn.is_method(), + is_method: turbo_fn.is_method(), + local, local_cells, - ); + }; let native_function_ident = get_trait_impl_function_ident(ty_ident, &trait_ident, ident); diff --git a/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs b/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs index 6fbbdb6f690d28..8d749582475ab0 100644 --- a/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs +++ b/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs @@ -114,18 +114,19 @@ pub fn value_trait(args: TokenStream, input: TokenStream) -> TokenStream { let (inline_signature, inline_block) = turbo_fn.inline_signature_and_block(default); let inline_attrs = filter_inline_attributes(&attrs[..]); - let native_function = NativeFn::new( - &format!("{trait_ident}::{ident}"), - &parse_quote! { + let native_function = NativeFn { + function_path_string: format!("{trait_ident}::{ident}"), + function_path: parse_quote! { as #inline_extension_trait_ident>::#inline_function_ident }, - turbo_fn.is_method(), - // `inline_cells` is currently unsupported here because: + is_method: turbo_fn.is_method(), + // `local` and `local_cells` are currently unsupported here because: // - The `#[turbo_tasks::function]` macro needs to be present for us to read this - // argument. + // argument. (This could be fixed) // - This only makes sense when a default implementation is present. - false, - ); + local: false, + local_cells: false, + }; let native_function_ident = get_trait_default_impl_function_ident(trait_ident, ident); let native_function_ty = native_function.ty(); diff --git a/turbopack/crates/turbo-tasks-memory/tests/local_tasks.rs b/turbopack/crates/turbo-tasks-memory/tests/local_tasks.rs new file mode 120000 index 00000000000000..29d9f86f455662 --- /dev/null +++ b/turbopack/crates/turbo-tasks-memory/tests/local_tasks.rs @@ -0,0 +1 @@ +../../turbo-tasks-testing/tests/local_tasks.rs \ No newline at end of file diff --git a/turbopack/crates/turbo-tasks-testing/tests/local_tasks.rs b/turbopack/crates/turbo-tasks-testing/tests/local_tasks.rs new file mode 100644 index 00000000000000..59fec1342db808 --- /dev/null +++ b/turbopack/crates/turbo-tasks-testing/tests/local_tasks.rs @@ -0,0 +1,34 @@ +#![feature(arbitrary_self_types)] +#![feature(arbitrary_self_types_pointers)] +#![allow(clippy::needless_return)] // tokio macro-generated code doesn't respect this + +use anyhow::Result; +use turbo_tasks::{test_helpers::current_task_for_testing, Vc}; +use turbo_tasks_testing::{register, run, Registration}; + +static REGISTRATION: Registration = register!(); + +#[tokio::test] +async fn test_local_task_id() -> Result<()> { + run(®ISTRATION, || async { + let local_vc = get_local_task_id(); + assert!(local_vc.is_local()); + assert_eq!(*local_vc.await.unwrap(), *current_task_for_testing()); + + let non_local_vc = get_non_local_task_id(); + assert!(!non_local_vc.is_local()); + assert_ne!(*non_local_vc.await.unwrap(), *current_task_for_testing()); + Ok(()) + }) + .await +} + +#[turbo_tasks::function(local)] +fn get_local_task_id() -> Vc { + Vc::cell(*current_task_for_testing()) +} + +#[turbo_tasks::function] +fn get_non_local_task_id() -> Vc { + Vc::cell(*current_task_for_testing()) +} diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs index 039fd8c04aa24c..9c79c8e28eddd0 100644 --- a/turbopack/crates/turbo-tasks/src/manager.rs +++ b/turbopack/crates/turbo-tasks/src/manager.rs @@ -320,19 +320,15 @@ pub enum TaskPersistence { /// [`TransientInstance`][crate::value::TransientInstance]. Transient, - /// Tasks that are persisted only for the lifetime of the nearest non-`LocalCells` parent - /// caller. + /// Tasks that are persisted only for the lifetime of the nearest non-`Local` parent caller. /// /// This task does not have a unique task id, and is not shared with the backend. Instead it /// uses the parent task's id. /// - /// Cells are allocated onto a temporary arena by default. Resolved cells inside a local task - /// are allocated into the parent task's cells. - /// /// This is useful for functions that have a low cache hit rate. Those functions could be /// converted to non-task functions, but that would break their function signature. This /// provides a mechanism for skipping caching without changing the function signature. - LocalCells, + Local, } #[derive(Clone, Copy, Eq, PartialEq)] @@ -595,12 +591,13 @@ impl TurboTasks { arg: Box, persistence: TaskPersistence, ) -> RawVc { - let task_type = CachedTaskType { fn_type, this, arg }; match persistence { - TaskPersistence::LocalCells => { - todo!("bgw: local tasks"); + TaskPersistence::Local => { + let task_type = LocalTaskType::Native { fn_type, this, arg }; + self.schedule_local_task(task_type, persistence) } TaskPersistence::Transient => { + let task_type = CachedTaskType { fn_type, this, arg }; RawVc::TaskOutput(self.backend.get_or_create_transient_task( task_type, current_task("turbo_function calls"), @@ -608,6 +605,7 @@ impl TurboTasks { )) } TaskPersistence::Persistent => { + let task_type = CachedTaskType { fn_type, this, arg }; RawVc::TaskOutput(self.backend.get_or_create_persistent_task( task_type, current_task("turbo_function calls"), diff --git a/turbopack/crates/turbo-tasks/src/native_function.rs b/turbopack/crates/turbo-tasks/src/native_function.rs index 70960b6137c5e6..d676cd780c3f2e 100644 --- a/turbopack/crates/turbo-tasks/src/native_function.rs +++ b/turbopack/crates/turbo-tasks/src/native_function.rs @@ -86,6 +86,10 @@ impl ArgMeta { #[derive(Debug)] pub struct FunctionMeta { + /// Does not run the function as a task, and instead runs it inside the parent task using + /// task-local state. The function call itself will not be cached, but cells will be created on + /// the parent task. + pub local: bool, /// Changes the behavior of `Vc::cell` to create local cells that are not /// cached across task executions. Cells can be converted to their non-local /// versions by calling `Vc::resolve`. @@ -175,7 +179,7 @@ impl NativeFunction { transient = true, ) } - TaskPersistence::LocalCells => { + TaskPersistence::Local => { tracing::trace_span!( "turbo_tasks::function", name = self.name.as_str(), @@ -197,11 +201,11 @@ impl NativeFunction { transient = true, ) } - TaskPersistence::LocalCells => { + TaskPersistence::Local => { tracing::trace_span!( "turbo_tasks::resolve_call", name = self.name.as_str(), - local_cells = true, + local = true, ) } } diff --git a/turbopack/crates/turbo-tasks/src/task/local_task.rs b/turbopack/crates/turbo-tasks/src/task/local_task.rs index 74d34291b2ce09..4fc6bb0c79685d 100644 --- a/turbopack/crates/turbo-tasks/src/task/local_task.rs +++ b/turbopack/crates/turbo-tasks/src/task/local_task.rs @@ -29,9 +29,9 @@ pub fn get_local_task_execution_spec<'a>( this, arg, } => { - debug_assert_eq!(persistence, TaskPersistence::LocalCells); + debug_assert_eq!(persistence, TaskPersistence::Local); let func = registry::get_function(*native_fn_id); - let span = func.span(TaskPersistence::LocalCells); + let span = func.span(TaskPersistence::Local); let entered = span.enter(); let future = func.execute(*this, &**arg); drop(entered); @@ -43,7 +43,7 @@ pub fn get_local_task_execution_spec<'a>( arg, } => { let func = registry::get_function(*native_fn_id); - let span = func.resolve_span(TaskPersistence::LocalCells); + let span = func.resolve_span(TaskPersistence::Local); let entered = span.enter(); let future = Box::pin(LocalTaskType::run_resolve_native( *native_fn_id,