refactor: Extract parallel queue abstraction

crates/bevy_ecs/Cargo.toml

@@ -23,7 +23,6 @@ bevy_ecs_macros = { path = "macros", version = "0.12.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)
@@ -53,8 +51,8 @@ impl SystemBuffer for ParallelCommandQueue {
     fn apply(&mut self, _system_meta: &SystemMeta, world: &mut World) {
         #[cfg(feature = "trace")]
         let _system_span = _system_meta.commands_span.enter();
-        for cq in &mut self.thread_local_storage {
-            cq.get_mut().apply(world);
+        for cq in self.thread_queues.iter_mut() {
+            cq.apply(world);
         }
     }
 }
@@ -64,16 +62,8 @@ impl<'w, 's> ParallelCommands<'w, 's> {
     ///
     /// For an example, see the type-level documentation for [`ParallelCommands`].
     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

@@ -9,8 +9,7 @@ use bevy_ecs::prelude::*;
 use bevy_hierarchy::{Children, Parent};
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use bevy_transform::{components::GlobalTransform, TransformSystem};
-use std::cell::Cell;
-use thread_local::ThreadLocal;
+use bevy_utils::Parallel;
 
 use crate::{
     camera::{
@@ -383,7 +382,7 @@ fn reset_view_visibility(mut query: Query<&mut ViewVisibility>) {
 /// [`ViewVisibility`] 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,
@@ -448,10 +447,9 @@ pub fn check_visibility(
             }
 
             view_visibility.set();
-            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);
+            });
         });
 
         visible_no_aabb_query.par_iter_mut().for_each(|query_item| {
@@ -469,15 +467,13 @@ pub fn check_visibility(
             }
 
             view_visibility.set();
-            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 &mut thread_queues {
-            visible_entities.entities.append(cell.get_mut());
-        }
+        visible_entities.entities.clear();
+        thread_queues.drain_into(&mut visible_entities.entities);
     }
 }
 

crates/bevy_utils/Cargo.toml

@@ -20,6 +20,7 @@ hashbrown = { version = "0.14", features = ["serde"] }
 bevy_utils_proc_macros = {version = "0.12.0-dev", path = "macros"}
 petgraph = "0.6"
 thiserror = "1.0"
+thread_local = "1.0"
 nonmax = "0.5"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]

crates/bevy_utils/src/lib.rs

@@ -21,6 +21,7 @@ pub mod syncunsafecell;
 mod cow_arc;
 mod default;
 mod float_ord;
+mod parallel_queue;
 
 pub use ahash::{AHasher, RandomState};
 pub use bevy_utils_proc_macros::*;
@@ -29,6 +30,7 @@ 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,70 @@
+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
+    /// 
+    /// The ordering is not guarenteed.
+    pub fn drain<B>(&mut self) -> impl Iterator<Item = T> + '_
+    where
+        B: FromIterator<T>,
+    {
+        self.locals
+            .iter_mut()
+            .flat_map(|item| item.take().into_iter())
+    }
+}
+
+impl<T: Send> Parallel<Vec<T>> {
+    /// Collect all enqueued items from all threads and appends them to the end of a 
+    /// single Vec.
+    /// 
+    /// The ordering is not guarenteed.
+    pub fn drain_into(&mut self, out: &mut Vec<T>) {
+        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());
+        }
+    }
+}