From e96a0d7b900f62642884b5b8f9943eb129acab60 Mon Sep 17 00:00:00 2001
From: Maksymilian Mozolewski <makspl17@gmail.com>
Date: Sat, 22 Feb 2025 12:58:57 +0000
Subject: [PATCH] feat: add allocator diagnostics (#305)

---
 .../src/bindings/allocator.rs                 | 45 +++++++++++++++++--
 examples/game_of_life.rs                      | 12 +++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/crates/bevy_mod_scripting_core/src/bindings/allocator.rs b/crates/bevy_mod_scripting_core/src/bindings/allocator.rs
index 0f094338..93830d0b 100644
--- a/crates/bevy_mod_scripting_core/src/bindings/allocator.rs
+++ b/crates/bevy_mod_scripting_core/src/bindings/allocator.rs
@@ -1,6 +1,12 @@
 //! An allocator used to control the lifetime of allocations
 
-use bevy::{ecs::system::Resource, prelude::ResMut, reflect::PartialReflect};
+use bevy::{
+    app::{App, Plugin, PostUpdate},
+    diagnostic::{Diagnostic, DiagnosticPath, Diagnostics, RegisterDiagnostic},
+    ecs::system::{Res, Resource},
+    prelude::ResMut,
+    reflect::PartialReflect,
+};
 use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use std::{
     cell::UnsafeCell,
@@ -11,6 +17,14 @@ use std::{
     sync::{atomic::AtomicU64, Arc},
 };
 
+/// The path used for the total number of allocations diagnostic
+pub const ALLOCATOR_TOTAL_DIAG_PATH: DiagnosticPath =
+    DiagnosticPath::const_new("scripting_allocator_total");
+
+/// The path used for the total number of deallocated allocations diagnostic
+pub const ALLOCATOR_TOTAL_COLLECTED_DIAG_PATH: DiagnosticPath =
+    DiagnosticPath::const_new("scripting_allocator_total_collected");
+
 #[derive(Clone, Debug)]
 /// Unique identifier for an allocation
 pub struct ReflectAllocationId(pub(crate) Arc<u64>);
@@ -214,9 +228,34 @@ impl ReflectAllocator {
 
 /// Cleans up dangling script allocations
 #[profiling::function]
-pub fn garbage_collector(allocator: ResMut<AppReflectAllocator>) {
+pub fn garbage_collector(allocator: ResMut<AppReflectAllocator>, mut diagnostics: Diagnostics) {
     let mut allocator = allocator.write();
-    allocator.clean_garbage_allocations()
+    let before = allocator.allocations.len();
+    allocator.clean_garbage_allocations();
+    let after = allocator.allocations.len();
+    diagnostics.add_measurement(&ALLOCATOR_TOTAL_DIAG_PATH, || after as f64);
+    diagnostics.add_measurement(&ALLOCATOR_TOTAL_COLLECTED_DIAG_PATH, || {
+        (before - after) as f64
+    });
+}
+
+/// Measures the number of allocations in the allocator and other diagnostics when enabled
+pub fn measure_allocations(allocator: Res<AppReflectAllocator>, mut diagnostics: Diagnostics) {
+    let allocator = allocator.read();
+    let allocations_count = allocator.allocations.len();
+    diagnostics.add_measurement(&ALLOCATOR_TOTAL_DIAG_PATH, || allocations_count as f64);
+}
+
+/// A plugin which registers various allocator diagnostics
+pub struct AllocatorDiagnosticPlugin;
+impl Plugin for AllocatorDiagnosticPlugin {
+    fn build(&self, app: &mut App) {
+        app.register_diagnostic(Diagnostic::new(ALLOCATOR_TOTAL_DIAG_PATH).with_suffix(" allocs"))
+            .register_diagnostic(
+                Diagnostic::new(ALLOCATOR_TOTAL_COLLECTED_DIAG_PATH).with_suffix(" deallocs"),
+            )
+            .add_systems(PostUpdate, measure_allocations);
+    }
 }
 
 #[cfg(test)]
diff --git a/examples/game_of_life.rs b/examples/game_of_life.rs
index c36e283f..247b9847 100644
--- a/examples/game_of_life.rs
+++ b/examples/game_of_life.rs
@@ -1,6 +1,9 @@
 #![allow(deprecated)]
 
+use std::time::Duration;
+
 use bevy::{
+    diagnostic::LogDiagnosticsPlugin,
     image::ImageSampler,
     log::LogPlugin,
     prelude::*,
@@ -18,6 +21,7 @@ use bevy_mod_scripting_core::{
     bindings::{
         function::namespace::{GlobalNamespace, NamespaceBuilder},
         script_value::ScriptValue,
+        AllocatorDiagnosticPlugin,
     },
     callback_labels,
     commands::AddStaticScript,
@@ -325,6 +329,14 @@ fn main() -> std::io::Result<()> {
     console_app(&mut app);
     game_of_life_app(&mut app);
 
+    app.add_plugins((
+        AllocatorDiagnosticPlugin,
+        LogDiagnosticsPlugin {
+            wait_duration: Duration::from_secs(60),
+            ..Default::default()
+        },
+    ));
+
     app.run();
 
     Ok(())