Allow systems using Diagnostics to run in parallel

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,24 +227,84 @@ 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.
     pub fn add_measurement<F>(&mut self, id: DiagnosticId, value: F)
     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,15 +1,15 @@
 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)]
 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 {