perf(turbo-tasks): Add a 'local' option to #[turbo_tasks::function(..)] (#75259)

Allows tasks to be marked as local using `#[turbo_tasks::function(local)]`.

Local tasks are cached only for the lifetime of the nearest non-Local parent caller.

Local tasks do not have a unique task id, and are not shared with the backend. Instead they use the parent task's id and store data in the backend-agnostic manager.

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.

Local tasks are already used for argument resolution, this allows other arbitrary functions to be opted-in.

Author: bgw

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

@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/local_tasks.rs

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)]

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)]

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<Span>,
+    /// 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<Span>,
     /// 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,

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::<FunctionArguments>(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();

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);

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! {
                     <Box<dyn #trait_ident> 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();

turbopack/crates/turbo-tasks-memory/tests/local_tasks.rs

@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/local_tasks.rs

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(&REGISTRATION, || 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<u32> {
+    Vc::cell(*current_task_for_testing())
+}
+
+#[turbo_tasks::function]
+fn get_non_local_task_id() -> Vc<u32> {
+    Vc::cell(*current_task_for_testing())
+}

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

@@ -313,19 +313,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, Debug, Eq, PartialEq)]
@@ -588,19 +584,21 @@ impl<B: Backend + 'static> TurboTasks<B> {
         arg: Box<dyn MagicAny>,
         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"),
                     self,
                 ))
             }
             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"),