Use atomics to allow for concurrent updates.

Author: BGR360

src/resource.rs

@@ -2,12 +2,12 @@ use std::any::{Any, TypeId};
 
 use bevy::utils::HashMap;
 
-use crate::MemoryStats;
+use crate::stats::{MemoryStats, MemoryStatsInternal};
 
 /// Stores memory usage statistics for registered data types.
 #[derive(Debug, Default)]
 pub struct MemoryUsage {
-    pub(crate) datasizes: HashMap<TypeId, MemoryStats>,
+    pub(crate) datasizes: HashMap<TypeId, MemoryStatsInternal>,
 }
 
 impl MemoryUsage {
@@ -20,7 +20,7 @@ impl MemoryUsage {
     {
         let type_id = TypeId::of::<T>();
 
-        self.datasizes.get(&type_id).copied()
+        self.datasizes.get(&type_id).map(MemoryStatsInternal::get)
     }
 
     /// Updates the [`MemoryStats`] for the given type.
@@ -35,7 +35,26 @@ impl MemoryUsage {
             .get_mut(&type_id)
             .expect("Memory usage not tracked for this type. Did you forget to register the type?");
 
-        *entry = stats;
+        *entry = MemoryStatsInternal::from(stats);
+    }
+
+    /// Like [`update_stats`][Self::update_stats] but operates on a shared reference.
+    ///
+    /// In exchange for the slight possibility that individual fields in
+    /// [`MemoryStats`] will be inconsistent with each other, this allows making
+    /// multiple updates to the resource concurrently.
+    pub fn update_stats_fast<T>(&self, stats: MemoryStats)
+    where
+        T: Any,
+    {
+        let type_id = TypeId::of::<T>();
+
+        let entry = self
+            .datasizes
+            .get(&type_id)
+            .expect("Memory usage not tracked for this type. Did you forget to register the type?");
+
+        entry.set(stats);
     }
 
     /// Registers the given type with the usage tracker.

src/stats.rs

@@ -1,4 +1,7 @@
-use std::any::Any;
+use std::{
+    any::Any,
+    sync::atomic::{AtomicUsize, Ordering},
+};
 
 use crate::DataSize;
 
@@ -107,3 +110,38 @@ impl MemoryStats {
         Self::stack_size_of(value) + Self::heap_size_of(value)
     }
 }
+
+#[derive(Debug, Default)]
+pub(crate) struct MemoryStatsInternal {
+    count: AtomicUsize,
+    total_stack_bytes: AtomicUsize,
+    total_heap_bytes: AtomicUsize,
+}
+
+impl MemoryStatsInternal {
+    pub(crate) fn get(&self) -> MemoryStats {
+        MemoryStats {
+            count: self.count.load(Ordering::Relaxed),
+            total_stack_bytes: self.total_stack_bytes.load(Ordering::Relaxed),
+            total_heap_bytes: self.total_heap_bytes.load(Ordering::Relaxed),
+        }
+    }
+
+    pub(crate) fn set(&self, stats: MemoryStats) {
+        self.count.store(stats.count, Ordering::Relaxed);
+        self.total_stack_bytes
+            .store(stats.total_stack_bytes, Ordering::Relaxed);
+        self.total_heap_bytes
+            .store(stats.total_heap_bytes, Ordering::Relaxed);
+    }
+}
+
+impl From<MemoryStats> for MemoryStatsInternal {
+    fn from(stats: MemoryStats) -> Self {
+        Self {
+            count: AtomicUsize::new(stats.count),
+            total_stack_bytes: AtomicUsize::new(stats.total_stack_bytes),
+            total_heap_bytes: AtomicUsize::new(stats.total_heap_bytes),
+        }
+    }
+}

src/systems.rs

@@ -12,7 +12,7 @@ use crate::{MemoryConfig, MemoryStats, MemoryUsage};
 pub fn update_stats_for_component<T>(
     query: Query<&T>,
     memory_config: Res<MemoryConfig>,
-    mut memory_usage: ResMut<MemoryUsage>,
+    memory_usage: Res<MemoryUsage>,
 ) where
     T: Any + Component + DataSize,
 {
@@ -22,5 +22,5 @@ pub fn update_stats_for_component<T>(
 
     let stats = MemoryStats::from_values(query.iter());
 
-    memory_usage.update_stats::<T>(stats);
+    memory_usage.update_stats_fast::<T>(stats);
 }