Rewrite IdFactory and IdFactoryWithReuse

Cargo.toml

@@ -226,7 +226,7 @@ camino = { version = "1.1.4", features = ["serde1"] }
 chrono = "0.4.23"
 clap = "4.5.2"
 clap_complete = "4.5.1"
-concurrent-queue = "2.1.0"
+concurrent-queue = "2.5.0"
 console = "0.15.5"
 console-subscriber = "0.1.8"
 criterion = "0.4.0"

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

@@ -25,7 +25,7 @@ use turbo_tasks::{
         TransientTaskType,
     },
     event::EventListener,
-    util::{IdFactory, NoMoveVec},
+    util::{IdFactoryWithReuse, NoMoveVec},
     CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused,
 };
 
@@ -45,7 +45,7 @@ fn prehash_task_type(task_type: PersistentTaskType) -> PreHashed<PersistentTaskT
 pub struct MemoryBackend {
     memory_tasks: NoMoveVec<Task, 13>,
     backend_jobs: NoMoveVec<Job>,
-    backend_job_id_factory: IdFactory<BackendJobId>,
+    backend_job_id_factory: IdFactoryWithReuse<BackendJobId>,
     task_cache:
         DashMap<Arc<PreHashed<PersistentTaskType>>, TaskId, BuildHasherDefault<PassThroughHash>>,
     memory_limit: usize,
@@ -65,7 +65,7 @@ impl MemoryBackend {
         Self {
             memory_tasks: NoMoveVec::new(),
             backend_jobs: NoMoveVec::new(),
-            backend_job_id_factory: IdFactory::new(),
+            backend_job_id_factory: IdFactoryWithReuse::new(),
             task_cache: DashMap::with_hasher_and_shard_amount(
                 Default::default(),
                 (std::thread::available_parallelism().map_or(1, usize::from) * 32)

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

@@ -26,7 +26,7 @@ use turbo_tasks::{
         ActivateResult, DeactivateResult, PersistResult, PersistTaskState, PersistedGraph,
         PersistedGraphApi, ReadTaskState, TaskCell, TaskData,
     },
-    util::{IdFactory, NoMoveVec, SharedError},
+    util::{IdFactoryWithReuse, NoMoveVec, SharedError},
     CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused,
 };
 
@@ -131,7 +131,7 @@ pub struct MemoryBackendWithPersistedGraph<P: PersistedGraph + 'static> {
     pub pg: P,
     tasks: NoMoveVec<Task>,
     cache: DashMap<PersistentTaskType, TaskId>,
-    background_job_id_factory: IdFactory<BackendJobId>,
+    background_job_id_factory: IdFactoryWithReuse<BackendJobId>,
     background_jobs: NoMoveVec<BackgroundJob>,
     only_known_to_memory_tasks: DashSet<TaskId>,
     /// Tasks that were selected to persist
@@ -154,7 +154,7 @@ pub struct MemoryBackendWithPersistedGraph<P: PersistedGraph + 'static> {
 
 impl<P: PersistedGraph> MemoryBackendWithPersistedGraph<P> {
     pub fn new(pg: P) -> Self {
-        let background_job_id_factory = IdFactory::new();
+        let background_job_id_factory = IdFactoryWithReuse::new();
         let persist_job = background_job_id_factory.get();
         Self {
             pg,

crates/turbo-tasks/src/id.rs

@@ -1,7 +1,7 @@
 use std::{
     fmt::{Debug, Display},
     mem::transmute_copy,
-    num::NonZeroU32,
+    num::{NonZeroU32, NonZeroU64, TryFromIntError},
     ops::Deref,
 };
 
@@ -17,6 +17,8 @@ macro_rules! define_id {
         }
 
         impl $name {
+            /// Constructs a wrapper type from the numeric identifier.
+            ///
             /// # Safety
             ///
             /// The passed `id` must not be zero.
@@ -39,11 +41,23 @@ macro_rules! define_id {
             }
         }
 
+        /// Converts a numeric identifier to the wrapper type.
+        ///
+        /// Panics if the given id value is zero.
         impl From<u32> for $name {
             fn from(id: u32) -> Self {
                 Self { id: NonZeroU32::new(id).expect("Ids can only be created from non zero values") }
             }
         }
+
+        /// Converts a numeric identifier to the wrapper type.
+        impl TryFrom<NonZeroU64> for $name {
+            type Error = TryFromIntError;
+
+            fn try_from(id: NonZeroU64) -> Result<Self, Self::Error> {
+                Ok(Self { id: NonZeroU32::try_from(id)? })
+            }
+        }
     };
     ($name:ident) => {
         define_id!(internal $name);

crates/turbo-tasks/src/id_factory.rs

@@ -1,44 +1,118 @@
 use std::{
+    any::type_name,
     marker::PhantomData,
-    ops::Deref,
-    sync::atomic::{AtomicU32, Ordering},
+    num::NonZeroU64,
+    sync::atomic::{AtomicU64, Ordering},
 };
 
 use concurrent_queue::ConcurrentQueue;
-use once_cell::sync::Lazy;
 
+/// A helper for constructing id types like [`FunctionId`][crate::FunctionId].
+///
+/// For ids that may be re-used, see [`IdFactoryWithReuse`].
 pub struct IdFactory<T> {
-    next_id: AtomicU32,
-    free_ids: Lazy<ConcurrentQueue<T>>,
-    phantom_data: PhantomData<T>,
+    next_id: AtomicU64,
+    _phantom_data: PhantomData<T>,
 }
 
-impl<T: From<u32> + Deref<Target = u32>> Default for IdFactory<T> {
+impl<T> IdFactory<T> {
+    pub const fn new() -> Self {
+        Self {
+            next_id: AtomicU64::new(1),
+            _phantom_data: PhantomData,
+        }
+    }
+}
+
+impl<T> Default for IdFactory<T> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl<T: From<u32> + Deref<Target = u32>> IdFactory<T> {
+impl<T> IdFactory<T>
+where
+    T: TryFrom<NonZeroU64>,
+{
+    /// Return a unique new id.
+    ///
+    /// Panics (best-effort) if the id type overflows.
+    pub fn get(&self) -> 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)) };
+
+        // Use the extra bits of the AtomicU64 as cheap overflow detection when the
+        // value is less than 64 bits.
+        match new_id.try_into() {
+            Ok(id) => id,
+            Err(_) => panic!(
+                "Overflow detected while attempting to generate a unique {}",
+                type_name::<T>(),
+            ),
+        }
+    }
+}
+
+/// An [`IdFactory`], but extended with a free list to allow for id reuse for
+/// ids such as [`BackendJobId`][crate::backend::BackendJobId].
+pub struct IdFactoryWithReuse<T> {
+    factory: IdFactory<T>,
+    free_ids: ConcurrentQueue<T>,
+}
+
+impl<T> IdFactoryWithReuse<T> {
     pub const fn new() -> Self {
         Self {
-            next_id: AtomicU32::new(1),
-            free_ids: Lazy::new(|| ConcurrentQueue::unbounded()),
-            phantom_data: PhantomData,
+            factory: IdFactory::new(),
+            free_ids: ConcurrentQueue::unbounded(),
         }
     }
+}
+
+impl<T> Default for IdFactoryWithReuse<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
 
+impl<T> IdFactoryWithReuse<T>
+where
+    T: TryFrom<NonZeroU64>,
+{
+    /// Return a new or potentially reused id.
+    ///
+    /// Panics (best-effort) if the id type overflows.
     pub fn get(&self) -> T {
-        if let Ok(id) = self.free_ids.pop() {
-            return id;
-        }
-        self.next_id.fetch_add(1, Ordering::Relaxed).into()
+        self.free_ids.pop().unwrap_or_else(|_| self.factory.get())
     }
 
+    /// Add an id to the free list, allowing it to be re-used on a subsequent
+    /// call to [`IdFactoryWithReuse::get`].
+    ///
     /// # Safety
     ///
     /// It must be ensured that the id is no longer used
     pub unsafe fn reuse(&self, id: T) {
         let _ = self.free_ids.push(id);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::num::NonZeroU8;
+
+    use super::*;
+
+    #[test]
+    #[should_panic(expected = "Overflow detected")]
+    fn test_overflow() {
+        let factory = IdFactory::<NonZeroU8>::new();
+        assert_eq!(factory.get(), NonZeroU8::new(1).unwrap());
+        assert_eq!(factory.get(), NonZeroU8::new(2).unwrap());
+        for _i in 2..256 {
+            factory.get();
+        }
+    }
+}

crates/turbo-tasks/src/manager.rs

@@ -28,7 +28,7 @@ use crate::{
     capture_future::{self, CaptureFuture},
     event::{Event, EventListener},
     id::{BackendJobId, FunctionId, TraitTypeId},
-    id_factory::IdFactory,
+    id_factory::IdFactoryWithReuse,
     raw_vc::{CellId, RawVc},
     registry,
     trace::TraceRawVcs,
@@ -133,7 +133,7 @@ pub trait TaskIdProvider {
     fn reuse_task_id(&self, id: Unused<TaskId>);
 }
 
-impl TaskIdProvider for IdFactory<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()) }
@@ -234,7 +234,7 @@ pub struct UpdateInfo {
 pub struct TurboTasks<B: Backend + 'static> {
     this: Weak<Self>,
     backend: B,
-    task_id_factory: IdFactory<TaskId>,
+    task_id_factory: IdFactoryWithReuse<TaskId>,
     stopped: AtomicBool,
     currently_scheduled_tasks: AtomicUsize,
     currently_scheduled_foreground_jobs: AtomicUsize,
@@ -279,7 +279,7 @@ 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 = IdFactory::new();
+        let task_id_factory = IdFactoryWithReuse::new();
         backend.initialize(&task_id_factory);
         let this = Arc::new_cyclic(|this| Self {
             this: this.clone(),

crates/turbo-tasks/src/registry.rs

@@ -1,4 +1,4 @@
-use std::{fmt::Debug, hash::Hash, ops::Deref};
+use std::{fmt::Debug, hash::Hash, num::NonZeroU64, ops::Deref};
 
 use dashmap::{mapref::entry::Entry, DashMap};
 use once_cell::sync::Lazy;
@@ -30,7 +30,7 @@ static TRAIT_TYPES_BY_VALUE: Lazy<DashMap<&'static TraitType, TraitTypeId>> =
 static TRAIT_TYPES: Lazy<NoMoveVec<(&'static TraitType, &'static str)>> = Lazy::new(NoMoveVec::new);
 
 fn register_thing<
-    K: From<u32> + Deref<Target = u32> + Sync + Send + Copy,
+    K: TryFrom<NonZeroU64> + Deref<Target = u32> + Sync + Send + Copy,
     V: Clone + Hash + Ord + Eq + Sync + Send + Copy,
     const INITIAL_CAPACITY_BITS: u32,
 >(

crates/turbo-tasks/src/util.rs

@@ -14,7 +14,11 @@ use anyhow::{anyhow, Error};
 use pin_project_lite::pin_project;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 
-pub use super::{id_factory::IdFactory, no_move_vec::NoMoveVec, once_map::*};
+pub use super::{
+    id_factory::{IdFactory, IdFactoryWithReuse},
+    no_move_vec::NoMoveVec,
+    once_map::*,
+};
 
 /// A error struct that is backed by an Arc to allow cloning errors
 #[derive(Debug, Clone)]