use different task ids for persistent and transient tasks

crates/turbo-tasks-memory/src/memory_backend.rs

@@ -192,7 +192,7 @@ impl MemoryBackend {
                 unsafe {
                     self.memory_tasks.remove(*new_id as usize);
                     let new_id = Unused::new_unchecked(new_id);
-                    turbo_tasks.reuse_task_id(new_id);
+                    turbo_tasks.reuse_persistent_task_id(new_id);
                 }
                 task_id
             }
@@ -612,7 +612,7 @@ impl Backend for MemoryBackend {
             let (task_type_hash, task_type) = PreHashed::into_parts(task_type);
             let task_type = Arc::new(PreHashed::new(task_type_hash, task_type));
             // slow pass with key lock
-            let id = turbo_tasks.get_fresh_task_id();
+            let id = turbo_tasks.get_fresh_persistent_task_id();
             let task = Task::new_persistent(
                 // Safety: That task will hold the value, but we are still in
                 // control of the task
@@ -652,7 +652,7 @@ impl Backend for MemoryBackend {
         task_type: TransientTaskType,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) -> TaskId {
-        let id = turbo_tasks.get_fresh_task_id();
+        let id = turbo_tasks.get_fresh_transient_task_id();
         let id = id.into();
         match task_type {
             TransientTaskType::Root(f) => {

crates/turbo-tasks/src/backend.rs

@@ -23,8 +23,8 @@ use crate::{
     raw_vc::CellId,
     registry,
     trait_helpers::{get_trait_method, has_trait, traits},
-    FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdProvider, TaskIdSet, TraitRef,
-    TraitTypeId, ValueTypeId, VcValueTrait, VcValueType,
+    FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdSet, TraitRef, TraitTypeId,
+    ValueTypeId, VcValueTrait, VcValueType,
 };
 
 pub enum TaskType {
@@ -383,7 +383,7 @@ pub type TaskCollectiblesMap = AutoMap<RawVc, i32, BuildHasherDefault<FxHasher>,
 
 pub trait Backend: Sync + Send {
     #[allow(unused_variables)]
-    fn initialize(&mut self, task_id_provider: &dyn TaskIdProvider) {}
+    fn initialize(&mut self) {}
 
     #[allow(unused_variables)]
     fn startup(&self, turbo_tasks: &dyn TurboTasksBackendApi<Self>) {}

crates/turbo-tasks/src/id.rs

@@ -77,6 +77,14 @@ impl Debug for TaskId {
     }
 }
 
+pub const TRANSIENT_TASK_BIT: u32 = 0x8000_0000;
+
+impl TaskId {
+    pub fn is_transient(&self) -> bool {
+        **self & TRANSIENT_TASK_BIT != 0
+    }
+}
+
 macro_rules! make_serializable {
     ($ty:ty, $get_global_name:path, $get_id:path, $visitor_name:ident) => {
         impl Serialize for $ty {

crates/turbo-tasks/src/id_factory.rs

@@ -12,13 +12,19 @@ use concurrent_queue::ConcurrentQueue;
 /// For ids that may be re-used, see [`IdFactoryWithReuse`].
 pub struct IdFactory<T> {
     next_id: AtomicU64,
+    max_id: u64,
     _phantom_data: PhantomData<T>,
 }
 
 impl<T> IdFactory<T> {
     pub const fn new() -> Self {
+        Self::new_with_range(1, u32::MAX as u64)
+    }
+
+    pub const fn new_with_range(start: u64, max: u64) -> Self {
         Self {
-            next_id: AtomicU64::new(1),
+            next_id: AtomicU64::new(start as u64),
+            max_id: max as u64,
             _phantom_data: PhantomData,
         }
     }
@@ -38,10 +44,18 @@ where
     ///
     /// Panics (best-effort) if the id type overflows.
     pub fn get(&self) -> T {
+        let new_id = self.next_id.fetch_add(1, Ordering::Relaxed);
+
+        if new_id > self.max_id {
+            panic!(
+                "Max id limit hit while attempting to generate a unique {}",
+                type_name::<T>(),
+            )
+        }
+
         // Safety: u64 will not overflow. This is *very* unlikely to happen (would take
         // decades).
-        let new_id =
-            unsafe { NonZeroU64::new_unchecked(self.next_id.fetch_add(1, Ordering::Relaxed)) };
+        let new_id = unsafe { NonZeroU64::new_unchecked(new_id) };
 
         // Use the extra bits of the AtomicU64 as cheap overflow detection when the
         // value is less than 64 bits.
@@ -69,6 +83,13 @@ impl<T> IdFactoryWithReuse<T> {
             free_ids: ConcurrentQueue::unbounded(),
         }
     }
+
+    pub const fn new_with_range(start: u64, max: u64) -> Self {
+        Self {
+            factory: IdFactory::new_with_range(start, max),
+            free_ids: ConcurrentQueue::unbounded(),
+        }
+    }
 }
 
 impl<T> Default for IdFactoryWithReuse<T> {
@@ -93,7 +114,8 @@ where
     ///
     /// # Safety
     ///
-    /// It must be ensured that the id is no longer used
+    /// It must be ensured that the id is no longer used. Id must be a valid id
+    /// that was previously returned by `get`.
     pub unsafe fn reuse(&self, id: T) {
         let _ = self.free_ids.push(id);
     }

crates/turbo-tasks/src/lib.rs

@@ -90,8 +90,8 @@ pub use magic_any::MagicAny;
 pub use manager::{
     dynamic_call, dynamic_this_call, emit, get_invalidator, mark_finished, mark_stateful,
     prevent_gc, run_once, run_once_with_reason, spawn_blocking, spawn_thread, trait_call,
-    turbo_tasks, CurrentCellRef, Invalidator, TaskIdProvider, TurboTasks, TurboTasksApi,
-    TurboTasksBackendApi, TurboTasksCallApi, Unused, UpdateInfo,
+    turbo_tasks, CurrentCellRef, Invalidator, TurboTasks, TurboTasksApi, TurboTasksBackendApi,
+    TurboTasksCallApi, Unused, UpdateInfo,
 };
 pub use native_function::NativeFunction;
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};

crates/turbo-tasks/src/manager.rs

@@ -30,7 +30,7 @@ use crate::{
     },
     capture_future::{self, CaptureFuture},
     event::{Event, EventListener},
-    id::{BackendJobId, FunctionId, TraitTypeId},
+    id::{BackendJobId, FunctionId, TraitTypeId, TRANSIENT_TASK_BIT},
     id_factory::IdFactoryWithReuse,
     magic_any::MagicAny,
     raw_vc::{CellId, RawVc},
@@ -136,22 +136,6 @@ pub trait TurboTasksApi: TurboTasksCallApi + Sync + Send {
     ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>;
 }
 
-pub trait TaskIdProvider {
-    fn get_fresh_task_id(&self) -> Unused<TaskId>;
-    fn reuse_task_id(&self, id: Unused<TaskId>);
-}
-
-impl TaskIdProvider for IdFactoryWithReuse<TaskId> {
-    fn get_fresh_task_id(&self) -> Unused<TaskId> {
-        // Safety: This is a fresh id from the factory
-        unsafe { Unused::new_unchecked(self.get()) }
-    }
-
-    fn reuse_task_id(&self, id: Unused<TaskId>) {
-        unsafe { self.reuse(id.into()) }
-    }
-}
-
 /// A wrapper around a value that is unused.
 pub struct Unused<T> {
     inner: T,
@@ -182,11 +166,14 @@ impl<T> Unused<T> {
     }
 }
 
-pub trait TurboTasksBackendApi<B: Backend + 'static>:
-    TaskIdProvider + TurboTasksCallApi + Sync + Send
-{
+pub trait TurboTasksBackendApi<B: Backend + 'static>: TurboTasksCallApi + Sync + Send {
     fn pin(&self) -> Arc<dyn TurboTasksBackendApi<B>>;
 
+    fn get_fresh_persistent_task_id(&self) -> Unused<TaskId>;
+    fn get_fresh_transient_task_id(&self) -> Unused<TaskId>;
+    unsafe fn reuse_persistent_task_id(&self, id: Unused<TaskId>);
+    unsafe fn reuse_transient_task_id(&self, id: Unused<TaskId>);
+
     fn schedule(&self, task: TaskId);
     fn schedule_backend_background_job(&self, id: BackendJobId);
     fn schedule_backend_foreground_job(&self, id: BackendJobId);
@@ -210,26 +197,6 @@ pub trait TurboTasksBackendApi<B: Backend + 'static>:
     fn backend(&self) -> &B;
 }
 
-impl<B: Backend + 'static> TaskIdProvider for &dyn TurboTasksBackendApi<B> {
-    fn get_fresh_task_id(&self) -> Unused<TaskId> {
-        (*self).get_fresh_task_id()
-    }
-
-    fn reuse_task_id(&self, id: Unused<TaskId>) {
-        (*self).reuse_task_id(id)
-    }
-}
-
-impl TaskIdProvider for &dyn TaskIdProvider {
-    fn get_fresh_task_id(&self) -> Unused<TaskId> {
-        (*self).get_fresh_task_id()
-    }
-
-    fn reuse_task_id(&self, id: Unused<TaskId>) {
-        (*self).reuse_task_id(id)
-    }
-}
-
 #[allow(clippy::manual_non_exhaustive)]
 pub struct UpdateInfo {
     pub duration: Duration,
@@ -243,6 +210,7 @@ pub struct TurboTasks<B: Backend + 'static> {
     this: Weak<Self>,
     backend: B,
     task_id_factory: IdFactoryWithReuse<TaskId>,
+    transient_task_id_factory: IdFactoryWithReuse<TaskId>,
     stopped: AtomicBool,
     currently_scheduled_tasks: AtomicUsize,
     currently_scheduled_foreground_jobs: AtomicUsize,
@@ -287,12 +255,16 @@ impl<B: Backend + 'static> TurboTasks<B> {
     // so we probably want to make sure that all tasks are joined
     // when trying to drop turbo tasks
     pub fn new(mut backend: B) -> Arc<Self> {
-        let task_id_factory = IdFactoryWithReuse::new();
-        backend.initialize(&task_id_factory);
+        let task_id_factory =
+            IdFactoryWithReuse::new_with_range(1, (TRANSIENT_TASK_BIT - 1) as u64);
+        let transient_task_id_factory =
+            IdFactoryWithReuse::new_with_range(TRANSIENT_TASK_BIT as u64, u32::MAX as u64);
+        backend.initialize();
         let this = Arc::new_cyclic(|this| Self {
             this: this.clone(),
             backend,
             task_id_factory,
+            transient_task_id_factory,
             stopped: AtomicBool::new(false),
             currently_scheduled_tasks: AtomicUsize::new(0),
             currently_scheduled_background_jobs: AtomicUsize::new(0),
@@ -1085,6 +1057,7 @@ impl<B: Backend + 'static> TurboTasksBackendApi<B> for TurboTasks<B> {
     fn backend(&self) -> &B {
         &self.backend
     }
+
     #[track_caller]
     fn schedule_backend_background_job(&self, id: BackendJobId) {
         self.schedule_background_job(move |this| async move {
@@ -1172,17 +1145,24 @@ impl<B: Backend + 'static> TurboTasksBackendApi<B> for TurboTasks<B> {
     fn program_duration_until(&self, instant: Instant) -> Duration {
         instant - self.program_start
     }
-}
 
-impl<B: Backend + 'static> TaskIdProvider for TurboTasks<B> {
-    fn get_fresh_task_id(&self) -> Unused<TaskId> {
-        // Safety: This is a fresh id from the factory
+    fn get_fresh_persistent_task_id(&self) -> Unused<TaskId> {
+        // SAFETY: This is a fresh id from the factory
         unsafe { Unused::new_unchecked(self.task_id_factory.get()) }
     }
 
-    fn reuse_task_id(&self, id: Unused<TaskId>) {
+    fn get_fresh_transient_task_id(&self) -> Unused<TaskId> {
+        // SAFETY: This is a fresh id from the factory
+        unsafe { Unused::new_unchecked(self.transient_task_id_factory.get()) }
+    }
+
+    unsafe fn reuse_persistent_task_id(&self, id: Unused<TaskId>) {
         unsafe { self.task_id_factory.reuse(id.into()) }
     }
+
+    unsafe fn reuse_transient_task_id(&self, id: Unused<TaskId>) {
+        unsafe { self.transient_task_id_factory.reuse(id.into()) }
+    }
 }
 
 pub(crate) fn current_task(from: &str) -> TaskId {