refactor(turbopack): Make invalidator flag explicit

crates/next-api/src/project.rs

@@ -423,7 +423,7 @@ impl ProjectContainer {
 
 #[turbo_tasks::value_impl]
 impl ProjectContainer {
-    #[turbo_tasks::function]
+    #[turbo_tasks::function(invalidator)]
     pub async fn project(&self) -> Result<Vc<Project>> {
         let env_map: Vc<EnvMap>;
         let next_config;

turbopack/crates/turbo-tasks-fs/src/lib.rs

@@ -551,7 +551,7 @@ impl Debug for DiskFileSystem {
 
 #[turbo_tasks::value_impl]
 impl FileSystem for DiskFileSystem {
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn read(&self, fs_path: Vc<FileSystemPath>) -> Result<Vc<FileContent>> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;
@@ -577,7 +577,7 @@ impl FileSystem for DiskFileSystem {
         Ok(content.cell())
     }
 
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn raw_read_dir(&self, fs_path: Vc<FileSystemPath>) -> Result<Vc<RawDirectoryContent>> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;
@@ -632,7 +632,7 @@ impl FileSystem for DiskFileSystem {
         Ok(RawDirectoryContent::new(entries))
     }
 
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn read_link(&self, fs_path: Vc<FileSystemPath>) -> Result<Vc<LinkContent>> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;
@@ -718,7 +718,7 @@ impl FileSystem for DiskFileSystem {
         .cell())
     }
 
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn write(&self, fs_path: Vc<FileSystemPath>, content: Vc<FileContent>) -> Result<()> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;
@@ -843,7 +843,7 @@ impl FileSystem for DiskFileSystem {
         Ok(())
     }
 
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn write_link(&self, fs_path: Vc<FileSystemPath>, target: Vc<LinkContent>) -> Result<()> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;
@@ -962,7 +962,7 @@ impl FileSystem for DiskFileSystem {
         Ok(())
     }
 
-    #[turbo_tasks::function(fs)]
+    #[turbo_tasks::function(fs, invalidator)]
     async fn metadata(&self, fs_path: Vc<FileSystemPath>) -> Result<Vc<FileMeta>> {
         mark_session_dependent();
         let full_path = self.to_sys_path(fs_path).await?;

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", "local"
+error: unexpected token, expected one of: "fs", "network", "operation", "local", "invalidator"
  --> 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", "local"
+error: unexpected token, expected one of: "fs", "network", "operation", "local", "invalidator"
   --> 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

@@ -733,6 +733,9 @@ pub struct FunctionArguments {
     /// task-local state. The function call itself will not be cached, but cells will be created on
     /// the parent task.
     pub local: Option<Span>,
+    /// If true, the function will be allowed to call `get_invalidator` . If this is false, the
+    /// `get_invalidator` function will panic on calls.
+    pub invalidator: Option<Span>,
 }
 
 impl Parse for FunctionArguments {
@@ -760,11 +763,14 @@ impl Parse for FunctionArguments {
                 ("local", Meta::Path(_)) => {
                     parsed_args.local = Some(meta.span());
                 }
+                ("invalidator", Meta::Path(_)) => {
+                    parsed_args.invalidator = Some(meta.span());
+                }
                 (_, meta) => {
                     return Err(syn::Error::new_spanned(
                         meta,
                         "unexpected token, expected one of: \"fs\", \"network\", \"operation\", \
-                         \"local\"",
+                         \"local\", \"invalidator\"",
                     ));
                 }
             }
@@ -1092,6 +1098,7 @@ pub struct NativeFn {
     pub is_self_used: bool,
     pub filter_trait_call_args: Option<FilterTraitCallArgsTokens>,
     pub local: bool,
+    pub invalidator: bool,
 }
 
 impl NativeFn {
@@ -1107,6 +1114,7 @@ impl NativeFn {
             is_self_used,
             filter_trait_call_args,
             local,
+            invalidator,
         } = self;
 
         if *is_method {
@@ -1133,6 +1141,7 @@ impl NativeFn {
                             #function_path_string.to_owned(),
                             turbo_tasks::macro_helpers::FunctionMeta {
                                 local: #local,
+                                invalidator: #invalidator,
                             },
                             #arg_filter,
                             #function_path,
@@ -1147,6 +1156,7 @@ impl NativeFn {
                             #function_path_string.to_owned(),
                             turbo_tasks::macro_helpers::FunctionMeta {
                                 local: #local,
+                                invalidator: #invalidator,
                             },
                             #arg_filter,
                             #function_path,
@@ -1162,6 +1172,7 @@ impl NativeFn {
                         #function_path_string.to_owned(),
                         turbo_tasks::macro_helpers::FunctionMeta {
                             local: #local,
+                            invalidator: #invalidator,
                         },
                         #function_path,
                     )

turbopack/crates/turbo-tasks-macros/src/function_macro.rs

@@ -42,6 +42,7 @@ pub fn function(args: TokenStream, input: TokenStream) -> TokenStream {
         .inspect_err(|err| errors.push(err.to_compile_error()))
         .unwrap_or_default();
     let local = args.local.is_some();
+    let invalidator = args.invalidator.is_some();
     let is_self_used = args.operation.is_some() || is_self_used(&block);
 
     let Some(turbo_fn) = TurboFn::new(&sig, DefinitionContext::NakedFn, args) else {
@@ -65,6 +66,7 @@ pub fn function(args: TokenStream, input: TokenStream) -> TokenStream {
         is_self_used,
         filter_trait_call_args: None, // not a trait method
         local,
+        invalidator,
     };
     let native_function_ident = get_native_function_ident(ident);
     let native_function_ty = native_fn.ty();

turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs

@@ -84,6 +84,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                     continue;
                 };
                 let local = func_args.local.is_some();
+                let invalidator = func_args.invalidator.is_some();
                 let is_self_used = func_args.operation.is_some() || is_self_used(block);
 
                 let Some(turbo_fn) =
@@ -106,6 +107,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                     is_self_used,
                     filter_trait_call_args: None, // not a trait method
                     local,
+                    invalidator,
                 };
 
                 let native_function_ident = get_inherent_impl_function_ident(ty_ident, ident);
@@ -191,6 +193,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                 };
 
                 let local = func_args.local.is_some();
+                let invalidator = func_args.invalidator.is_some();
                 let is_self_used = func_args.operation.is_some() || is_self_used(block);
 
                 let Some(turbo_fn) =
@@ -223,6 +226,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                     is_self_used,
                     filter_trait_call_args: turbo_fn.filter_trait_call_args(),
                     local,
+                    invalidator,
                 };
 
                 let native_function_ident =

turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs

@@ -139,6 +139,7 @@ pub fn value_trait(args: TokenStream, input: TokenStream) -> TokenStream {
                 //   argument. (This could be fixed)
                 // - This only makes sense when a default implementation is present.
                 local: false,
+                invalidator: func_args.invalidator.is_some(),
             };
 
             let native_function_ident = get_trait_default_impl_function_ident(trait_ident, ident);

turbopack/crates/turbo-tasks-testing/tests/read_ref_cell.rs

@@ -70,7 +70,7 @@ impl Counter {
 
 #[turbo_tasks::value_impl]
 impl Counter {
-    #[turbo_tasks::function]
+    #[turbo_tasks::function(invalidator)]
     async fn get_value(&self) -> Result<Vc<CounterValue>> {
         let mut lock = self.value.lock().unwrap();
         lock.1 = Some(get_invalidator());

turbopack/crates/turbo-tasks-testing/tests/trait_ref_cell.rs

@@ -85,7 +85,7 @@ trait CounterTrait {
 
 #[turbo_tasks::value_impl]
 impl CounterTrait for Counter {
-    #[turbo_tasks::function]
+    #[turbo_tasks::function(invalidator)]
     async fn get_value(&self) -> Result<Vc<CounterValue>> {
         let mut lock = self.value.lock().unwrap();
         lock.1 = Some(get_invalidator());

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

@@ -9,7 +9,7 @@ use std::{
 use anyhow::Result;
 use indexmap::map::Entry;
 use serde::{Deserialize, Serialize, de::Visitor};
-use tokio::runtime::Handle;
+use tokio::{runtime::Handle, task_local};
 
 use crate::{
     FxIndexMap, FxIndexSet, TaskId, TurboTasksApi,
@@ -19,9 +19,27 @@ use crate::{
     util::StaticOrArc,
 };
 
+#[cfg(debug_assertions)]
+task_local! {
+    static ALLOW_INVALIDATOR: ();
+}
+
+#[cfg(debug_assertions)]
+pub fn allow_invalidator<R>(f: impl Future<Output = R>) -> impl Future<Output = R> {
+    ALLOW_INVALIDATOR.scope((), f)
+}
+
 /// Get an [`Invalidator`] that can be used to invalidate the current task
 /// based on external events.
 pub fn get_invalidator() -> Invalidator {
+    #[cfg(debug_assertions)]
+    if ALLOW_INVALIDATOR.try_with(|_| {}).is_err() {
+        panic!(
+            "Invalidator can only be used in the turbo-tasks function that has \
+             #[turbo_tasks::function(invalidator)] attribute"
+        );
+    }
+
     let handle = Handle::current();
     Invalidator {
         task: current_task("turbo_tasks::get_invalidator()"),

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

@@ -144,6 +144,10 @@ pub struct FunctionMeta {
     /// task-local state. The function call itself will not be cached, but cells will be created on
     /// the parent task.
     pub local: bool,
+
+    /// If true, the function will be allowed to call `get_invalidator` . If this is false, the
+    /// `get_invalidator` function will panic on calls.
+    pub invalidator: bool,
 }
 
 /// A native (rust) turbo-tasks function. It's used internally by
@@ -235,7 +239,14 @@ impl NativeFunction {
     /// Executed the function
     pub fn execute(&'static self, this: Option<RawVc>, arg: &dyn MagicAny) -> NativeTaskFuture {
         match (self.implementation).functor(this, arg) {
-            Ok(functor) => functor,
+            Ok(functor) => {
+                #[cfg(debug_assertions)]
+                if self.function_meta.invalidator {
+                    return Box::pin(crate::invalidation::allow_invalidator(functor));
+                }
+
+                functor
+            }
             Err(err) => Box::pin(async { Err(err) }),
         }
     }