Allow systems using Diagnostics to run in parallel (#8677)

# Objective

I was trying to add some `Diagnostics` to have a better break down of
performance but I noticed that the current implementation uses a
`ResMut` which forces the functions to all run sequentially whereas
before they could run in parallel. This created too great a performance
penalty to be usable.

## Solution

This PR reworks how the diagnostics work with a couple of breaking
changes. The idea is to change how `Diagnostics` works by changing it to
a `SystemParam`. This allows us to hold a `Deferred` buffer of
measurements that can be applied later, avoiding the need for multiple
mutable references to the hashmap. This means we can run systems that
write diagnostic measurements in parallel.

Firstly, we rename the old `Diagnostics` to `DiagnosticsStore`. This
clears up the original name for the new interface while allowing us to
preserve more closely the original API.

Then we create a new `Diagnostics` struct which implements `SystemParam`
and contains a deferred `SystemBuffer`. This can be used very similar to
the old `Diagnostics` for writing new measurements.

```rust
fn system(diagnostics: ResMut<Diagnostics>) { diagnostics.new_measurement(ID, || 10.0)}
// changes to
fn system(mut diagnostics: Diagnostics) { diagnostics.new_measurement(ID, || 10.0)}
``` 
For reading the diagnostics, the user needs to change from `Diagnostics`
to `DiagnosticsStore` but otherwise the function calls are the same.

Finally, we add a new method to the `App` for registering diagnostics.
This replaces the old method of creating a startup system and adding it
manually.

Testing it, this PR does indeed allow Diagnostic systems to be run in
parallel.

## Changelog

- Change `Diagnostics` to implement `SystemParam` which allows
diagnostic systems to run in parallel.

## Migration Guide

- Register `Diagnostic`'s using the new
`app.register_diagnostic(Diagnostic::new(DIAGNOSTIC_ID,
"diagnostic_name", 10));`
- In systems for writing new measurements, change `mut diagnostics:
ResMut<Diagnostics>` to `mut diagnostics: Diagnostics` to allow the
systems to run in parallel.
- In systems for reading measurements, change `diagnostics:
Res<Diagnostics>` to `diagnostics: Res<DiagnosticsStore>`.

crates/bevy_asset/src/diagnostic/asset_count_diagnostics_plugin.rs

@@ -1,6 +1,8 @@
 use crate::{Asset, Assets};
 use bevy_app::prelude::*;
-use bevy_diagnostic::{Diagnostic, DiagnosticId, Diagnostics, MAX_DIAGNOSTIC_NAME_WIDTH};
+use bevy_diagnostic::{
+    Diagnostic, DiagnosticId, Diagnostics, DiagnosticsStore, MAX_DIAGNOSTIC_NAME_WIDTH,
+};
 use bevy_ecs::prelude::*;
 
 /// Adds an asset count diagnostic to an [`App`] for assets of type `T`.
@@ -32,7 +34,7 @@ impl<T: Asset> AssetCountDiagnosticsPlugin<T> {
     }
 
     /// Registers the asset count diagnostic for the current application.
-    pub fn setup_system(mut diagnostics: ResMut<Diagnostics>) {
+    pub fn setup_system(mut diagnostics: ResMut<DiagnosticsStore>) {
         let asset_type_name = std::any::type_name::<T>();
         let max_length = MAX_DIAGNOSTIC_NAME_WIDTH - "asset_count ".len();
         diagnostics.add(Diagnostic::new(
@@ -52,7 +54,7 @@ impl<T: Asset> AssetCountDiagnosticsPlugin<T> {
     }
 
     /// Updates the asset count of `T` assets.
-    pub fn diagnostic_system(mut diagnostics: ResMut<Diagnostics>, assets: Res<Assets<T>>) {
+    pub fn diagnostic_system(mut diagnostics: Diagnostics, assets: Res<Assets<T>>) {
         diagnostics.add_measurement(Self::diagnostic_id(), || assets.len() as f64);
     }
 }

crates/bevy_diagnostic/src/diagnostic.rs

@@ -1,4 +1,5 @@
-use bevy_ecs::system::Resource;
+use bevy_app::App;
+use bevy_ecs::system::{Deferred, Res, Resource, SystemBuffer, SystemParam};
 use bevy_log::warn;
 use bevy_utils::{Duration, Instant, StableHashMap, Uuid};
 use std::{borrow::Cow, collections::VecDeque};
@@ -44,16 +45,14 @@ pub struct Diagnostic {
 }
 
 impl Diagnostic {
-    /// Add a new value as a [`DiagnosticMeasurement`]. Its timestamp will be [`Instant::now`].
-    pub fn add_measurement(&mut self, value: f64) {
-        let time = Instant::now();
-
+    /// Add a new value as a [`DiagnosticMeasurement`].
+    pub fn add_measurement(&mut self, measurement: DiagnosticMeasurement) {
         if let Some(previous) = self.measurement() {
-            let delta = (time - previous.time).as_secs_f64();
+            let delta = (measurement.time - previous.time).as_secs_f64();
             let alpha = (delta / self.ema_smoothing_factor).clamp(0.0, 1.0);
-            self.ema += alpha * (value - self.ema);
+            self.ema += alpha * (measurement.value - self.ema);
         } else {
-            self.ema = value;
+            self.ema = measurement.value;
         }
 
         if self.max_history_length > 1 {
@@ -63,14 +62,13 @@ impl Diagnostic {
                 }
             }
 
-            self.sum += value;
+            self.sum += measurement.value;
         } else {
             self.history.clear();
-            self.sum = value;
+            self.sum = measurement.value;
         }
 
-        self.history
-            .push_back(DiagnosticMeasurement { time, value });
+        self.history.push_back(measurement);
     }
 
     /// Create a new diagnostic with the given ID, name and maximum history.
@@ -199,14 +197,16 @@ impl Diagnostic {
 
 /// A collection of [Diagnostic]s
 #[derive(Debug, Default, Resource)]
-pub struct Diagnostics {
+pub struct DiagnosticsStore {
     // This uses a [`StableHashMap`] to ensure that the iteration order is deterministic between
     // runs when all diagnostics are inserted in the same order.
     diagnostics: StableHashMap<DiagnosticId, Diagnostic>,
 }
 
-impl Diagnostics {
+impl DiagnosticsStore {
     /// Add a new [`Diagnostic`].
+    ///
+    /// If possible, prefer calling [`App::register_diagnostic`].
     pub fn add(&mut self, diagnostic: Diagnostic) {
         self.diagnostics.insert(diagnostic.id, diagnostic);
     }
@@ -227,6 +227,20 @@ impl Diagnostics {
             .and_then(|diagnostic| diagnostic.measurement())
     }
 
+    /// Return an iterator over all [`Diagnostic`].
+    pub fn iter(&self) -> impl Iterator<Item = &Diagnostic> {
+        self.diagnostics.values()
+    }
+}
+
+/// Record new [`DiagnosticMeasurement`]'s.
+#[derive(SystemParam)]
+pub struct Diagnostics<'w, 's> {
+    store: Res<'w, DiagnosticsStore>,
+    queue: Deferred<'s, DiagnosticsBuffer>,
+}
+
+impl<'w, 's> Diagnostics<'w, 's> {
     /// Add a measurement to an enabled [`Diagnostic`]. The measurement is passed as a function so that
     /// it will be evaluated only if the [`Diagnostic`] is enabled. This can be useful if the value is
     /// costly to calculate.
@@ -234,17 +248,63 @@ impl Diagnostics {
     where
         F: FnOnce() -> f64,
     {
-        if let Some(diagnostic) = self
-            .diagnostics
-            .get_mut(&id)
+        if self
+            .store
+            .get(id)
             .filter(|diagnostic| diagnostic.is_enabled)
+            .is_some()
         {
-            diagnostic.add_measurement(value());
+            let measurement = DiagnosticMeasurement {
+                time: Instant::now(),
+                value: value(),
+            };
+            self.queue.0.insert(id, measurement);
         }
     }
+}
 
-    /// Return an iterator over all [`Diagnostic`].
-    pub fn iter(&self) -> impl Iterator<Item = &Diagnostic> {
-        self.diagnostics.values()
+#[derive(Default)]
+struct DiagnosticsBuffer(StableHashMap<DiagnosticId, DiagnosticMeasurement>);
+
+impl SystemBuffer for DiagnosticsBuffer {
+    fn apply(
+        &mut self,
+        _system_meta: &bevy_ecs::system::SystemMeta,
+        world: &mut bevy_ecs::world::World,
+    ) {
+        let mut diagnostics = world.resource_mut::<DiagnosticsStore>();
+        for (id, measurement) in self.0.drain() {
+            if let Some(diagnostic) = diagnostics.get_mut(id) {
+                diagnostic.add_measurement(measurement);
+            }
+        }
+    }
+}
+
+/// Extend [`App`] with new `register_diagnostic` function.
+pub trait RegisterDiagnostic {
+    fn register_diagnostic(&mut self, diagnostic: Diagnostic) -> &mut Self;
+}
+
+impl RegisterDiagnostic for App {
+    /// Register a new [`Diagnostic`] with an [`App`].
+    ///
+    /// ```rust
+    /// use bevy_app::App;
+    /// use bevy_diagnostic::{Diagnostic, DiagnosticsPlugin, DiagnosticId, RegisterDiagnostic};
+    ///
+    /// const UNIQUE_DIAG_ID: DiagnosticId = DiagnosticId::from_u128(42);
+    ///
+    /// App::new()
+    ///     .add_plugin(DiagnosticsPlugin)
+    ///     // Must only be called after the `DiagnosticsPlugin` has been added.
+    ///     .register_diagnostic(Diagnostic::new(UNIQUE_DIAG_ID, "example", 10))
+    ///     .run();
+    /// ```
+    fn register_diagnostic(&mut self, diagnostic: Diagnostic) -> &mut Self {
+        let mut diagnostics = self.world.resource_mut::<DiagnosticsStore>();
+        diagnostics.add(diagnostic);
+
+        self
     }
 }

crates/bevy_diagnostic/src/entity_count_diagnostics_plugin.rs

@@ -1,7 +1,7 @@
 use bevy_app::prelude::*;
-use bevy_ecs::{entity::Entities, prelude::*};
+use bevy_ecs::entity::Entities;
 
-use crate::{Diagnostic, DiagnosticId, Diagnostics};
+use crate::{Diagnostic, DiagnosticId, Diagnostics, RegisterDiagnostic};
 
 /// Adds "entity count" diagnostic to an App
 #[derive(Default)]
@@ -9,7 +9,7 @@ pub struct EntityCountDiagnosticsPlugin;
 
 impl Plugin for EntityCountDiagnosticsPlugin {
     fn build(&self, app: &mut App) {
-        app.add_systems(Startup, Self::setup_system)
+        app.register_diagnostic(Diagnostic::new(Self::ENTITY_COUNT, "entity_count", 20))
             .add_systems(Update, Self::diagnostic_system);
     }
 }
@@ -18,11 +18,7 @@ impl EntityCountDiagnosticsPlugin {
     pub const ENTITY_COUNT: DiagnosticId =
         DiagnosticId::from_u128(187513512115068938494459732780662867798);
 
-    pub fn setup_system(mut diagnostics: ResMut<Diagnostics>) {
-        diagnostics.add(Diagnostic::new(Self::ENTITY_COUNT, "entity_count", 20));
-    }
-
-    pub fn diagnostic_system(mut diagnostics: ResMut<Diagnostics>, entities: &Entities) {
+    pub fn diagnostic_system(mut diagnostics: Diagnostics, entities: &Entities) {
         diagnostics.add_measurement(Self::ENTITY_COUNT, || entities.len() as f64);
     }
 }

crates/bevy_diagnostic/src/frame_time_diagnostics_plugin.rs

@@ -1,4 +1,4 @@
-use crate::{Diagnostic, DiagnosticId, Diagnostics};
+use crate::{Diagnostic, DiagnosticId, Diagnostics, RegisterDiagnostic};
 use bevy_app::prelude::*;
 use bevy_core::FrameCount;
 use bevy_ecs::prelude::*;
@@ -10,8 +10,14 @@ pub struct FrameTimeDiagnosticsPlugin;
 
 impl Plugin for FrameTimeDiagnosticsPlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        app.add_systems(Startup, Self::setup_system)
-            .add_systems(Update, Self::diagnostic_system);
+        app.register_diagnostic(
+            Diagnostic::new(Self::FRAME_TIME, "frame_time", 20).with_suffix("ms"),
+        )
+        .register_diagnostic(Diagnostic::new(Self::FPS, "fps", 20))
+        .register_diagnostic(
+            Diagnostic::new(Self::FRAME_COUNT, "frame_count", 1).with_smoothing_factor(0.0),
+        )
+        .add_systems(Update, Self::diagnostic_system);
     }
 }
 
@@ -22,15 +28,8 @@ impl FrameTimeDiagnosticsPlugin {
     pub const FRAME_TIME: DiagnosticId =
         DiagnosticId::from_u128(73441630925388532774622109383099159699);
 
-    pub fn setup_system(mut diagnostics: ResMut<Diagnostics>) {
-        diagnostics.add(Diagnostic::new(Self::FRAME_TIME, "frame_time", 20).with_suffix("ms"));
-        diagnostics.add(Diagnostic::new(Self::FPS, "fps", 20));
-        diagnostics
-            .add(Diagnostic::new(Self::FRAME_COUNT, "frame_count", 1).with_smoothing_factor(0.0));
-    }
-
     pub fn diagnostic_system(
-        mut diagnostics: ResMut<Diagnostics>,
+        mut diagnostics: Diagnostics,
         time: Res<Time>,
         frame_count: Res<FrameCount>,
     ) {

crates/bevy_diagnostic/src/lib.rs

@@ -19,7 +19,7 @@ pub struct DiagnosticsPlugin;
 
 impl Plugin for DiagnosticsPlugin {
     fn build(&self, app: &mut App) {
-        app.init_resource::<Diagnostics>().add_systems(
+        app.init_resource::<DiagnosticsStore>().add_systems(
             Startup,
             system_information_diagnostics_plugin::internal::log_system_info,
         );

crates/bevy_diagnostic/src/log_diagnostics_plugin.rs

@@ -1,4 +1,4 @@
-use super::{Diagnostic, DiagnosticId, Diagnostics};
+use super::{Diagnostic, DiagnosticId, DiagnosticsStore};
 use bevy_app::prelude::*;
 use bevy_ecs::prelude::*;
 use bevy_log::{debug, info};
@@ -83,7 +83,7 @@ impl LogDiagnosticsPlugin {
     fn log_diagnostics_system(
         mut state: ResMut<LogDiagnosticsState>,
         time: Res<Time>,
-        diagnostics: Res<Diagnostics>,
+        diagnostics: Res<DiagnosticsStore>,
     ) {
         if state.timer.tick(time.raw_delta()).finished() {
             if let Some(ref filter) = state.filter {
@@ -108,7 +108,7 @@ impl LogDiagnosticsPlugin {
     fn log_diagnostics_debug_system(
         mut state: ResMut<LogDiagnosticsState>,
         time: Res<Time>,
-        diagnostics: Res<Diagnostics>,
+        diagnostics: Res<DiagnosticsStore>,
     ) {
         if state.timer.tick(time.raw_delta()).finished() {
             if let Some(ref filter) = state.filter {

crates/bevy_diagnostic/src/system_information_diagnostics_plugin.rs

@@ -41,11 +41,11 @@ pub mod internal {
     use bevy_log::info;
     use sysinfo::{CpuExt, CpuRefreshKind, RefreshKind, System, SystemExt};
 
-    use crate::{Diagnostic, Diagnostics};
+    use crate::{Diagnostic, Diagnostics, DiagnosticsStore};
 
     const BYTES_TO_GIB: f64 = 1.0 / 1024.0 / 1024.0 / 1024.0;
 
-    pub(crate) fn setup_system(mut diagnostics: ResMut<Diagnostics>) {
+    pub(crate) fn setup_system(mut diagnostics: ResMut<DiagnosticsStore>) {
         diagnostics.add(
             Diagnostic::new(
                 super::SystemInformationDiagnosticsPlugin::CPU_USAGE,
@@ -65,7 +65,7 @@ pub mod internal {
     }
 
     pub(crate) fn diagnostic_system(
-        mut diagnostics: ResMut<Diagnostics>,
+        mut diagnostics: Diagnostics,
         mut sysinfo: Local<Option<System>>,
     ) {
         if sysinfo.is_none() {

examples/diagnostics/custom_diagnostic.rs

@@ -1,7 +1,7 @@
 //! This example illustrates how to create a custom diagnostic.
 
 use bevy::{
-    diagnostic::{Diagnostic, DiagnosticId, Diagnostics, LogDiagnosticsPlugin},
+    diagnostic::{Diagnostic, DiagnosticId, Diagnostics, LogDiagnosticsPlugin, RegisterDiagnostic},
     prelude::*,
 };
 
@@ -11,7 +11,11 @@ fn main() {
         // The "print diagnostics" plugin is optional.
         // It just visualizes our diagnostics in the console.
         .add_plugin(LogDiagnosticsPlugin::default())
-        .add_systems(Startup, setup_diagnostic_system)
+        // Diagnostics must be initialized before measurements can be added.
+        .register_diagnostic(
+            Diagnostic::new(SYSTEM_ITERATION_COUNT, "system_iteration_count", 10)
+                .with_suffix(" iterations"),
+        )
         .add_systems(Update, my_system)
         .run();
 }
@@ -21,17 +25,7 @@ fn main() {
 pub const SYSTEM_ITERATION_COUNT: DiagnosticId =
     DiagnosticId::from_u128(337040787172757619024841343456040760896);
 
-fn setup_diagnostic_system(mut diagnostics: ResMut<Diagnostics>) {
-    // Diagnostics must be initialized before measurements can be added.
-    // In general it's a good idea to set them up in a "startup system".
-    diagnostics.add(Diagnostic::new(
-        SYSTEM_ITERATION_COUNT,
-        "system_iteration_count",
-        10,
-    ));
-}
-
-fn my_system(mut diagnostics: ResMut<Diagnostics>) {
+fn my_system(mut diagnostics: Diagnostics) {
     // Add a measurement of 10.0 for our diagnostic each time this system runs.
     diagnostics.add_measurement(SYSTEM_ITERATION_COUNT, || 10.0);
 }

examples/stress_tests/bevymark.rs

@@ -3,7 +3,7 @@
 //! Usage: spawn more entities by clicking on the screen.
 
 use bevy::{
-    diagnostic::{Diagnostics, FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
+    diagnostic::{DiagnosticsStore, FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
     prelude::*,
     window::{PresentMode, WindowResolution},
 };
@@ -244,7 +244,7 @@ fn collision_system(windows: Query<&Window>, mut bird_query: Query<(&mut Bird, &
 }
 
 fn counter_system(
-    diagnostics: Res<Diagnostics>,
+    diagnostics: Res<DiagnosticsStore>,
     counter: Res<BevyCounter>,
     mut query: Query<&mut Text, With<StatsText>>,
 ) {

examples/stress_tests/many_gizmos.rs

@@ -1,7 +1,7 @@
 use std::f32::consts::TAU;
 
 use bevy::{
-    diagnostic::{Diagnostics, FrameTimeDiagnosticsPlugin},
+    diagnostic::{DiagnosticsStore, FrameTimeDiagnosticsPlugin},
     prelude::*,
     window::PresentMode,
 };
@@ -90,7 +90,7 @@ fn setup(mut commands: Commands) {
     ));
 }
 
-fn ui_system(mut query: Query<&mut Text>, config: Res<Config>, diag: Res<Diagnostics>) {
+fn ui_system(mut query: Query<&mut Text>, config: Res<Config>, diag: Res<DiagnosticsStore>) {
     let mut text = query.single_mut();
 
     let Some(fps) = diag.get(FrameTimeDiagnosticsPlugin::FPS).and_then(|fps| fps.smoothed()) else {

examples/ui/text.rs

@@ -4,7 +4,7 @@
 //! in the bottom right. For text within a scene, please see the text2d example.
 
 use bevy::{
-    diagnostic::{Diagnostics, FrameTimeDiagnosticsPlugin},
+    diagnostic::{DiagnosticsStore, FrameTimeDiagnosticsPlugin},
     prelude::*,
 };
 
@@ -86,7 +86,10 @@ fn text_color_system(time: Res<Time>, mut query: Query<&mut Text, With<ColorText
     }
 }
 
-fn text_update_system(diagnostics: Res<Diagnostics>, mut query: Query<&mut Text, With<FpsText>>) {
+fn text_update_system(
+    diagnostics: Res<DiagnosticsStore>,
+    mut query: Query<&mut Text, With<FpsText>>,
+) {
     for mut text in &mut query {
         if let Some(fps) = diagnostics.get(FrameTimeDiagnosticsPlugin::FPS) {
             if let Some(value) = fps.smoothed() {

examples/ui/text_debug.rs

@@ -1,7 +1,7 @@
 //! Shows various text layout options.
 
 use bevy::{
-    diagnostic::{Diagnostics, FrameTimeDiagnosticsPlugin},
+    diagnostic::{DiagnosticsStore, FrameTimeDiagnosticsPlugin},
     prelude::*,
     window::{PresentMode, WindowPlugin},
 };
@@ -135,7 +135,7 @@ fn infotext_system(mut commands: Commands, asset_server: Res<AssetServer>) {
 
 fn change_text_system(
     time: Res<Time>,
-    diagnostics: Res<Diagnostics>,
+    diagnostics: Res<DiagnosticsStore>,
     mut query: Query<&mut Text, With<TextChanges>>,
 ) {
     for mut text in &mut query {