Refactor system diagnostics to use a single task

System information diagnostics now use a single task that wakes up every
time the `First` schedule runs. The task checks if enough time has
passed since the last refresh. If enough time has passed, it refreshes
the system information and sends it to a channel. The
`read_diagonstic_task` system then reads the system information from the
channel to add diagnostic data.

Author: dloukadakis

crates/bevy_diagnostic/src/diagnostic.rs

@@ -39,18 +39,18 @@ impl DiagnosticPath {
     pub fn new(path: impl Into<Cow<'static, str>>) -> DiagnosticPath {
         let path = path.into();
 
-        debug_assert!(!path.is_empty(), "diagnostic path can't be empty");
+        debug_assert!(!path.is_empty(), "diagnostic path should not be empty");
         debug_assert!(
             !path.starts_with('/'),
-            "diagnostic path can't be start with `/`"
+            "diagnostic path should not start with `/`"
         );
         debug_assert!(
             !path.ends_with('/'),
-            "diagnostic path can't be end with `/`"
+            "diagnostic path should not end with `/`"
         );
         debug_assert!(
             !path.contains("//"),
-            "diagnostic path can't contain empty components"
+            "diagnostic path should not contain empty components"
         );
 
         DiagnosticPath {

crates/bevy_diagnostic/src/system_information_diagnostics_plugin.rs

@@ -70,19 +70,25 @@ pub struct SystemInfo {
     feature = "std",
 ))]
 pub mod internal {
+    use core::{
+        pin::Pin,
+        task::{Context, Poll, Waker},
+    };
+    use std::sync::{
+        mpsc::{self, Receiver, Sender},
+        Arc, Mutex,
+    };
+
     use alloc::{
         format,
         string::{String, ToString},
-        sync::Arc,
-        vec::Vec,
     };
     use bevy_app::{App, First, Startup, Update};
     use bevy_ecs::resource::Resource;
-    use bevy_ecs::{prelude::ResMut, system::Local};
+    use bevy_ecs::{prelude::ResMut, system::Commands};
     use bevy_platform::time::Instant;
-    use bevy_tasks::{available_parallelism, block_on, poll_once, AsyncComputeTaskPool, Task};
+    use bevy_tasks::AsyncComputeTaskPool;
     use log::info;
-    use std::sync::Mutex;
     use sysinfo::{CpuRefreshKind, MemoryRefreshKind, RefreshKind, System};
 
     use crate::{Diagnostic, Diagnostics, DiagnosticsStore};
@@ -93,12 +99,20 @@ pub mod internal {
 
     pub(super) fn setup_plugin(app: &mut App) {
         app.add_systems(Startup, setup_system)
-            .add_systems(First, launch_diagnostic_tasks)
-            .add_systems(Update, read_diagnostic_tasks)
-            .init_resource::<SysinfoTasks>();
+            .add_systems(First, wake_diagnostic_task)
+            .add_systems(Update, read_diagnostic_task);
     }
 
-    fn setup_system(mut diagnostics: ResMut<DiagnosticsStore>) {
+    fn setup_system(mut diagnostics: ResMut<DiagnosticsStore>, mut commands: Commands) {
+        let (tx, rx) = mpsc::channel();
+        let diagnostic_task = DiagnosticTask::new(tx);
+        let waker = Arc::clone(&diagnostic_task.waker);
+        AsyncComputeTaskPool::get().spawn(diagnostic_task).detach();
+        commands.insert_resource(SysinfoTask {
+            receiver: Mutex::new(rx),
+            waker,
+        });
+
         diagnostics.add(
             Diagnostic::new(SystemInformationDiagnosticsPlugin::SYSTEM_CPU_USAGE).with_suffix("%"),
         );
@@ -121,78 +135,100 @@ pub mod internal {
         process_mem_usage: f64,
     }
 
-    #[derive(Resource, Default)]
-    struct SysinfoTasks {
-        tasks: Vec<Task<SysinfoRefreshData>>,
+    impl SysinfoRefreshData {
+        fn new(system: &mut System) -> Self {
+            let pid = sysinfo::get_current_pid().expect("Failed to get current process ID");
+            system.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
+
+            system.refresh_cpu_specifics(CpuRefreshKind::nothing().with_cpu_usage());
+            system.refresh_memory();
+
+            let system_cpu_usage = system.global_cpu_usage().into();
+            let total_mem = system.total_memory() as f64;
+            let used_mem = system.used_memory() as f64;
+            let system_mem_usage = used_mem / total_mem * 100.0;
+
+            let process_mem_usage = system
+                .process(pid)
+                .map(|p| p.memory() as f64 * BYTES_TO_GIB)
+                .unwrap_or(0.0);
+
+            let process_cpu_usage = system
+                .process(pid)
+                .map(|p| p.cpu_usage() as f64 / system.cpus().len() as f64)
+                .unwrap_or(0.0);
+
+            Self {
+                system_cpu_usage,
+                system_mem_usage,
+                process_cpu_usage,
+                process_mem_usage,
+            }
+        }
+    }
+
+    #[derive(Resource)]
+    struct SysinfoTask {
+        receiver: Mutex<Receiver<SysinfoRefreshData>>,
+        waker: Arc<Mutex<Option<Waker>>>,
     }
 
-    fn launch_diagnostic_tasks(
-        mut tasks: ResMut<SysinfoTasks>,
-        // TODO: Consider a fair mutex
-        mut sysinfo: Local<Option<Arc<Mutex<System>>>>,
-        // TODO: FromWorld for Instant?
-        mut last_refresh: Local<Option<Instant>>,
-    ) {
-        let sysinfo = sysinfo.get_or_insert_with(|| {
-            Arc::new(Mutex::new(System::new_with_specifics(
-                RefreshKind::nothing()
-                    .with_cpu(CpuRefreshKind::nothing().with_cpu_usage())
-                    .with_memory(MemoryRefreshKind::everything()),
-            )))
-        });
+    struct DiagnosticTask {
+        system: System,
+        last_refresh: Instant,
+        sender: Sender<SysinfoRefreshData>,
+        waker: Arc<Mutex<Option<Waker>>>,
+    }
 
-        let last_refresh = last_refresh.get_or_insert_with(Instant::now);
-
-        let thread_pool = AsyncComputeTaskPool::get();
-
-        // Only queue a new system refresh task when necessary
-        // Queuing earlier than that will not give new data
-        if last_refresh.elapsed() > sysinfo::MINIMUM_CPU_UPDATE_INTERVAL
-            // These tasks don't yield and will take up all of the task pool's
-            // threads if we don't limit their amount.
-            && tasks.tasks.len() * 2 < available_parallelism()
-        {
-            let sys = Arc::clone(sysinfo);
-            let task = thread_pool.spawn(async move {
-                let mut sys = sys.lock().unwrap();
-                let pid = sysinfo::get_current_pid().expect("Failed to get current process ID");
-                sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
-
-                sys.refresh_cpu_specifics(CpuRefreshKind::nothing().with_cpu_usage());
-                sys.refresh_memory();
-                let system_cpu_usage = sys.global_cpu_usage().into();
-                let total_mem = sys.total_memory() as f64;
-                let used_mem = sys.used_memory() as f64;
-                let system_mem_usage = used_mem / total_mem * 100.0;
-
-                let process_mem_usage = sys
-                    .process(pid)
-                    .map(|p| p.memory() as f64 * BYTES_TO_GIB)
-                    .unwrap_or(0.0);
-
-                let process_cpu_usage = sys
-                    .process(pid)
-                    .map(|p| p.cpu_usage() as f64 / sys.cpus().len() as f64)
-                    .unwrap_or(0.0);
-
-                SysinfoRefreshData {
-                    system_cpu_usage,
-                    system_mem_usage,
-                    process_cpu_usage,
-                    process_mem_usage,
-                }
-            });
-            tasks.tasks.push(task);
-            *last_refresh = Instant::now();
+    impl DiagnosticTask {
+        fn new(sender: Sender<SysinfoRefreshData>) -> Self {
+            Self {
+                system: System::new_with_specifics(
+                    RefreshKind::nothing()
+                        .with_cpu(CpuRefreshKind::nothing().with_cpu_usage())
+                        .with_memory(MemoryRefreshKind::everything()),
+                ),
+                last_refresh: Instant::now(),
+                sender,
+                waker: Arc::default(),
+            }
         }
     }
 
-    fn read_diagnostic_tasks(mut diagnostics: Diagnostics, mut tasks: ResMut<SysinfoTasks>) {
-        tasks.tasks.retain_mut(|task| {
-            let Some(data) = block_on(poll_once(task)) else {
-                return true;
-            };
+    impl DiagnosticTask {
+        fn update_waker(&mut self, cx: &mut Context<'_>) {
+            let mut waker = self.waker.lock().unwrap();
+            *waker = Some(cx.waker().clone());
+        }
+    }
+
+    impl Future for DiagnosticTask {
+        type Output = ();
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            self.update_waker(cx);
+
+            if self.last_refresh.elapsed() > sysinfo::MINIMUM_CPU_UPDATE_INTERVAL {
+                self.last_refresh = Instant::now();
+
+                let sysinfo_refresh_data = SysinfoRefreshData::new(&mut self.system);
+                self.sender.send(sysinfo_refresh_data).unwrap();
+            }
 
+            Poll::Pending
+        }
+    }
+
+    fn wake_diagnostic_task(task: ResMut<SysinfoTask>) {
+        let mut waker = task.waker.lock().unwrap();
+        if let Some(waker) = waker.take() {
+            waker.wake_by_ref();
+        }
+    }
+
+    fn read_diagnostic_task(mut diagnostics: Diagnostics, task: ResMut<SysinfoTask>) {
+        let receiver = task.receiver.lock().unwrap();
+        while let Ok(data) = receiver.try_recv() {
             diagnostics.add_measurement(
                 &SystemInformationDiagnosticsPlugin::SYSTEM_CPU_USAGE,
                 || data.system_cpu_usage,
@@ -209,8 +245,7 @@ pub mod internal {
                 &SystemInformationDiagnosticsPlugin::PROCESS_MEM_USAGE,
                 || data.process_mem_usage,
             );
-            false
-        });
+        }
     }
 
     impl Default for SystemInfo {