Repository: vercel/next.js
Pull request: Turbopack: allow to read State from turbo_tasks::run

### What?

Automatically switch reading from state to untracked mode when in `turbo_tasks::run` (similar to reading from cells).

Closes PACK-5492

Changed files:
filename: turbopack/crates/turbo-tasks-backend/tests/read_ref_cell.rs
patch: @@ -73,7 +73,7 @@ impl Counter {
     #[turbo_tasks::function]
     fn get_value(&self) -> Result<Vc<CounterValue>> {
         let mut lock = self.value.lock().unwrap();
-        lock.1.insert(get_invalidator());
+        lock.1.insert(get_invalidator().unwrap());
         Ok(Vc::cell(lock.0))
     }
 }

filename: turbopack/crates/turbo-tasks-backend/tests/trait_ref_cell.rs
patch: @@ -90,7 +90,7 @@ impl CounterTrait for Counter {
     #[turbo_tasks::function]
     fn get_value(&self) -> Result<Vc<CounterValue>> {
         let mut lock = self.value.lock().unwrap();
-        lock.1.insert(get_invalidator());
+        lock.1.insert(get_invalidator().unwrap());
         Ok(Vc::cell(lock.0))
     }
 

filename: turbopack/crates/turbo-tasks-fs/src/lib.rs
patch: @@ -257,10 +257,11 @@ impl DiskFileSystemInner {
     /// registers the path as an invalidator for the current task,
     /// has to be called within a turbo-tasks function
     fn register_read_invalidator(&self, path: &Path) -> Result<()> {
-        let invalidator = turbo_tasks::get_invalidator();
-        self.invalidator_map
-            .insert(path.to_owned(), invalidator, None);
-        self.watcher.ensure_watched_file(path, self.root_path())?;
+        if let Some(invalidator) = turbo_tasks::get_invalidator() {
+            self.invalidator_map
+                .insert(path.to_owned(), invalidator, None);
+            self.watcher.ensure_watched_file(path, self.root_path())?;
+        }
         Ok(())
     }
 
@@ -293,10 +294,11 @@ impl DiskFileSystemInner {
     /// registers the path as an invalidator for the current task,
     /// has to be called within a turbo-tasks function
     fn register_dir_invalidator(&self, path: &Path) -> Result<()> {
-        let invalidator = turbo_tasks::get_invalidator();
-        self.dir_invalidator_map
-            .insert(path.to_owned(), invalidator, None);
-        self.watcher.ensure_watched_dir(path, self.root_path())?;
+        if let Some(invalidator) = turbo_tasks::get_invalidator() {
+            self.dir_invalidator_map
+                .insert(path.to_owned(), invalidator, None);
+            self.watcher.ensure_watched_dir(path, self.root_path())?;
+        }
         Ok(())
     }
 
@@ -712,11 +714,16 @@ impl FileSystem for DiskFileSystem {
             let _lock = inner.lock_path(&full_path).await;
 
             // Track the file, so that we will rewrite it if it ever changes.
-            let old_invalidators = inner.register_write_invalidator(
-                &full_path,
-                invalidator,
-                WriteContent::File(content.clone()),
-            )?;
+            let old_invalidators = invalidator
+                .map(|invalidator| {
+                    inner.register_write_invalidator(
+                        &full_path,
+                        invalidator,
+                        WriteContent::File(content.clone()),
+                    )
+                })
+                .transpose()?
+                .unwrap_or_default();
 
             // We perform an untracked comparison here, so that this write is not dependent
             // on a read's Vc<FileContent> (and the memory it holds). Our untracked read can
@@ -846,11 +853,16 @@ impl FileSystem for DiskFileSystem {
 
             let _lock = inner.lock_path(&full_path).await;
 
-            let old_invalidators = inner.register_write_invalidator(
-                &full_path,
-                invalidator,
-                WriteContent::Link(content.clone()),
-            )?;
+            let old_invalidators = invalidator
+                .map(|invalidator| {
+                    inner.register_write_invalidator(
+                        &full_path,
+                        invalidator,
+                        WriteContent::Link(content.clone()),
+                    )
+                })
+                .transpose()?
+                .unwrap_or_default();
 
             // TODO(sokra) preform a untracked read here, register an invalidator and get
             // all existing invalidators

filename: turbopack/crates/turbo-tasks/src/invalidation.rs
patch: @@ -15,21 +15,26 @@ use turbo_dyn_eq_hash::{
 
 use crate::{
     FxIndexMap, FxIndexSet, TaskId, TurboTasksApi,
-    manager::{current_task, mark_invalidator, with_turbo_tasks},
+    manager::{current_task_if_available, mark_invalidator, with_turbo_tasks},
     trace::TraceRawVcs,
     util::StaticOrArc,
 };
 
 /// Get an [`Invalidator`] that can be used to invalidate the current task
 /// based on external events.
-pub fn get_invalidator() -> Invalidator {
-    mark_invalidator();
+/// Returns `None` if called outside of a task context.
+pub fn get_invalidator() -> Option<Invalidator> {
+    if let Some(task) = current_task_if_available("turbo_tasks::get_invalidator()") {
+        mark_invalidator();
 
-    let handle = Handle::current();
-    Invalidator {
-        task: current_task("turbo_tasks::get_invalidator()"),
-        turbo_tasks: with_turbo_tasks(Arc::downgrade),
-        handle,
+        let handle = Handle::current();
+        Some(Invalidator {
+            task,
+            turbo_tasks: with_turbo_tasks(Arc::downgrade),
+            handle,
+        })
+    } else {
+        None
     }
 }
 

filename: turbopack/crates/turbo-tasks/src/state.rs
patch: @@ -170,7 +170,9 @@ impl<T> State<T> {
     pub fn get(&self) -> StateRef<'_, T> {
         let invalidator = get_invalidator();
         let mut inner = self.inner.lock();
-        inner.add_invalidator(invalidator);
+        if let Some(invalidator) = invalidator {
+            inner.add_invalidator(invalidator);
+        }
         StateRef {
             serialization_invalidator: Some(&self.serialization_invalidator),
             inner,
@@ -295,7 +297,9 @@ impl<T> TransientState<T> {
         mark_session_dependent();
         let invalidator = get_invalidator();
         let mut inner = self.inner.lock();
-        inner.add_invalidator(invalidator);
+        if let Some(invalidator) = invalidator {
+            inner.add_invalidator(invalidator);
+        }
         StateRef {
             serialization_invalidator: None,
             inner,