refactor: Extract parallel queue abstraction

crates/bevy_ecs/Cargo.toml

@@ -22,7 +22,6 @@ bevy_ecs_macros = { path = "macros", version = "0.11.0-dev" }
 
 async-channel = "1.4"
 event-listener = "2.5"
-thread_local = "1.1.4"
 fixedbitset = "0.4.2"
 rustc-hash = "1.1"
 downcast-rs = "1.2"

crates/bevy_ecs/src/system/commands/parallel_scope.rs

@@ -1,6 +1,4 @@
-use std::cell::Cell;
-
-use thread_local::ThreadLocal;
+use bevy_utils::Parallel;
 
 use crate::{
     self as bevy_ecs,
@@ -13,7 +11,7 @@ use super::{CommandQueue, Commands};
 
 #[derive(Default)]
 struct ParallelCommandQueue {
-    thread_local_storage: ThreadLocal<Cell<CommandQueue>>,
+    thread_queues: Parallel<CommandQueue>,
 }
 
 /// An alternative to [`Commands`] that can be used in parallel contexts, such as those in [`Query::par_iter`](crate::system::Query::par_iter)
@@ -55,24 +53,16 @@ impl SystemBuffer for ParallelCommandQueue {
         let _system_span =
             bevy_utils::tracing::info_span!("system_commands", name = _system_meta.name())
                 .entered();
-        for cq in &mut self.thread_local_storage {
-            cq.get_mut().apply(world);
+        for cq in self.thread_queues.iter_mut() {
+            cq.apply(world);
         }
     }
 }
 
 impl<'w, 's> ParallelCommands<'w, 's> {
     pub fn command_scope<R>(&self, f: impl FnOnce(Commands) -> R) -> R {
-        let store = &self.state.thread_local_storage;
-        let command_queue_cell = store.get_or_default();
-        let mut command_queue = command_queue_cell.take();
-
-        let r = f(Commands::new_from_entities(
-            &mut command_queue,
-            self.entities,
-        ));
-
-        command_queue_cell.set(command_queue);
-        r
+        self.state
+            .thread_queues
+            .scope(|command_queue| f(Commands::new_from_entities(command_queue, self.entities)))
     }
 }

crates/bevy_render/src/view/visibility/mod.rs

@@ -10,8 +10,7 @@ use bevy_reflect::Reflect;
 use bevy_reflect::{std_traits::ReflectDefault, FromReflect};
 use bevy_transform::components::GlobalTransform;
 use bevy_transform::TransformSystem;
-use std::cell::Cell;
-use thread_local::ThreadLocal;
+use bevy_utils::Parallel;
 
 use crate::{
     camera::{
@@ -351,7 +350,7 @@ fn propagate_recursive(
 /// [`ComputedVisibility`] of all entities, and for each view also compute the [`VisibleEntities`]
 /// for that view.
 pub fn check_visibility(
-    mut thread_queues: Local<ThreadLocal<Cell<Vec<Entity>>>>,
+    mut thread_queues: Local<Parallel<Vec<Entity>>>,
     mut view_query: Query<(&mut VisibleEntities, &Frustum, Option<&RenderLayers>), With<Camera>>,
     mut visible_aabb_query: Query<(
         Entity,
@@ -408,10 +407,9 @@ pub fn check_visibility(
                 }
 
                 computed_visibility.set_visible_in_view();
-                let cell = thread_queues.get_or_default();
-                let mut queue = cell.take();
-                queue.push(entity);
-                cell.set(queue);
+                thread_queues.scope(|queue| {
+                    queue.push(entity);
+                })
             },
         );
 
@@ -429,16 +427,13 @@ pub fn check_visibility(
                 }
 
                 computed_visibility.set_visible_in_view();
-                let cell = thread_queues.get_or_default();
-                let mut queue = cell.take();
-                queue.push(entity);
-                cell.set(queue);
+                thread_queues.scope(|queue| {
+                    queue.push(entity);
+                });
             },
         );
 
-        for cell in thread_queues.iter_mut() {
-            visible_entities.entities.append(cell.get_mut());
-        }
+        thread_queues.drain_into(&mut visible_entities.entities);
     }
 }
 

crates/bevy_utils/Cargo.toml

@@ -20,6 +20,7 @@ hashbrown = { version = "0.12", features = ["serde"] }
 bevy_utils_proc_macros = {version = "0.11.0-dev", path = "macros"}
 petgraph = "0.6"
 thiserror = "1.0"
+thread_local = "1.0"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 getrandom = {version = "0.2.0", features = ["js"]}

crates/bevy_utils/src/lib.rs

@@ -20,13 +20,15 @@ pub mod syncunsafecell;
 
 mod default;
 mod float_ord;
+mod parallel_queue;
 
 pub use ahash::AHasher;
 pub use bevy_utils_proc_macros::*;
 pub use default::default;
 pub use float_ord::*;
 pub use hashbrown;
 pub use instant::{Duration, Instant};
+pub use parallel_queue::*;
 pub use petgraph;
 pub use thiserror;
 pub use tracing;

crates/bevy_utils/src/parallel_queue.rs

@@ -0,0 +1,67 @@
+use core::cell::Cell;
+use thread_local::ThreadLocal;
+
+/// A cohesive set of thread-local values of a given type.
+///
+/// Mutable references can be fetched if `T: Default` via [`Parallel::scope`].
+#[derive(Default)]
+pub struct Parallel<T: Send> {
+    locals: ThreadLocal<Cell<T>>,
+}
+
+impl<T: Send> Parallel<T> {
+    /// Gets a mutable iterator over all of the per-thread queues.
+    pub fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut T> {
+        self.locals.iter_mut().map(|cell| cell.get_mut())
+    }
+
+    /// Clears all of the stored thread local values.
+    pub fn clear(&mut self) {
+        self.locals.clear();
+    }
+}
+
+impl<T: Default + Send> Parallel<T> {
+    /// Retrieves the thread-local value for the current thread and runs `f` on it.
+    ///
+    /// If there is no thread-local value, it will be initialized to it's default.
+    pub fn scope<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
+        let cell = self.locals.get_or_default();
+        let mut value = cell.take();
+        let ret = f(&mut value);
+        cell.set(value);
+        ret
+    }
+}
+
+impl<T, I> Parallel<I>
+where
+    I: IntoIterator<Item = T> + Default + Send + 'static,
+{
+    /// Collect all enqueued items from all threads and them into one
+    pub fn drain<B>(&mut self) -> B
+    where
+        B: FromIterator<T>,
+    {
+        self.locals
+            .iter_mut()
+            .flat_map(|item| item.take().into_iter())
+            .collect()
+    }
+}
+
+impl<T: Send> Parallel<Vec<T>> {
+    /// Collect all enqueued items from all threads and them into one
+    pub fn drain_into(&mut self, out: &mut Vec<T>) {
+        out.clear();
+        let size = self
+            .locals
+            .iter_mut()
+            .map(|queue| queue.get_mut().len())
+            .sum();
+        out.reserve(size);
+        for queue in self.locals.iter_mut() {
+            out.append(queue.get_mut());
+        }
+    }
+}