refactor(sysinfo): collect network stats in a background task shared via watch channel

Cargo.toml

@@ -13,7 +13,7 @@ members = [
 ]
 
 [workspace.package]
-version = "1.5.3"
+version = "1.5.4"
 authors = ["The Dragonfly Developers"]
 edition = "2021"
 readme = "README.md"
@@ -34,14 +34,14 @@ clap = { version = "4.6.1", features = ["derive", "env"] }
 crc32fast = "1.5.0"
 dashmap = "6.1.0"
 dragonfly-api = "=2.3.6"
-dragonfly-client = { path = "dragonfly-client", version = "1.5.3" }
-dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.5.3" }
-dragonfly-client-config = { path = "dragonfly-client-config", version = "1.5.3" }
-dragonfly-client-core = { path = "dragonfly-client-core", version = "1.5.3" }
-dragonfly-client-init = { path = "dragonfly-client-init", version = "1.5.3" }
-dragonfly-client-metric = { path = "dragonfly-client-metric", version = "1.5.3" }
-dragonfly-client-storage = { path = "dragonfly-client-storage", version = "1.5.3" }
-dragonfly-client-util = { path = "dragonfly-client-util", version = "1.5.3" }
+dragonfly-client = { path = "dragonfly-client", version = "1.5.4" }
+dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.5.4" }
+dragonfly-client-config = { path = "dragonfly-client-config", version = "1.5.4" }
+dragonfly-client-core = { path = "dragonfly-client-core", version = "1.5.4" }
+dragonfly-client-init = { path = "dragonfly-client-init", version = "1.5.4" }
+dragonfly-client-metric = { path = "dragonfly-client-metric", version = "1.5.4" }
+dragonfly-client-storage = { path = "dragonfly-client-storage", version = "1.5.4" }
+dragonfly-client-util = { path = "dragonfly-client-util", version = "1.5.4" }
 fastrand = "2.4.1"
 fs2 = "0.4.3"
 futures = "0.3.32"

dragonfly-client-util/src/sysinfo/network.rs

@@ -18,10 +18,9 @@ use bytesize::ByteSize;
 use pnet::datalink::{self, NetworkInterface};
 use std::cmp::min;
 use std::net::IpAddr;
-use std::sync::Arc;
 use std::time::Duration;
 use sysinfo::Networks;
-use tokio::sync::Mutex;
+use tokio::sync::watch;
 use tracing::{debug, info, warn};
 
 /// Represents the network statistics for a specific interface.
@@ -40,28 +39,20 @@ pub struct NetworkStats {
     pub tx_bandwidth: Option<u64>,
 }
 
-/// Represents a network interface with its information.
-#[derive(Debug, Clone, Default)]
+/// Represents a network interface with its latest statistics.
+#[derive(Debug, Clone)]
 pub struct Network {
-    // The name of the network interface.
-    interface_name: String,
-
-    // The bandwidth of the network interface in bits per second (bps).
-    bandwidth: u64,
-
-    // Mutex to protect concurrent access to network statistics.
-    mutex: Arc<Mutex<()>>,
+    // The latest statistics published by the collector, None until the first collection.
+    stats: watch::Receiver<Option<NetworkStats>>,
 }
 
 /// Implementation of network monitoring functionality.
 ///
 /// Provides methods to retrieve network interface information and statistics,
 /// including bandwidth measurements and traffic monitoring.
 impl Network {
-    /// Default interval for refreshing network statistics.
-    const DEFAULT_NETWORK_REFRESH_INTERVAL: Duration = Duration::from_secs(1);
-
-    /// Creates a new Network instance based on the provided IP address and rate limit.
+    /// Creates a new Network instance based on the provided IP address and rate limit,
+    /// and spawns the statistics collector, so it must be called within a tokio runtime.
     ///
     /// # Arguments
     /// * `ip` - The IP address to identify the network interface.
@@ -76,11 +67,7 @@ impl Network {
             "can not find interface for IP address {}, network interface unknown with bandwidth {} bps",
             ip, rate_limit
         );
-            return Self {
-                interface_name: "unknown".to_string(),
-                bandwidth: rate_limit,
-                mutex: Arc::new(Mutex::new(())),
-            };
+            return Self::spawn("unknown".to_string(), rate_limit);
         };
 
         match Self::get_speed(&interface.name) {
@@ -91,80 +78,41 @@ impl Network {
                     interface.name, bandwidth
                 );
 
-                Self {
-                    interface_name: interface.name,
-                    bandwidth,
-                    mutex: Arc::new(Mutex::new(())),
-                }
+                Self::spawn(interface.name, bandwidth)
             }
             None => {
                 warn!(
                     "can not get speed, network interface {} with bandwidth {} bps",
                     interface.name, rate_limit
                 );
 
-                Self {
-                    interface_name: interface.name,
-                    bandwidth: rate_limit,
-                    mutex: Arc::new(Mutex::new(())),
-                }
+                Self::spawn(interface.name, rate_limit)
             }
         }
     }
 
-    /// Retrieves the network statistics for the interface.
-    ///
-    /// This method measures network traffic over a time interval (DEFAULT_NETWORK_REFRESH_INTERVAL)
-    /// to calculate current receive and transmit bandwidth.
+    /// Spawns the statistics collector for the interface and returns the Network instance.
+    fn spawn(interface_name: String, bandwidth: u64) -> Network {
+        let (tx, rx) = watch::channel(None);
+        tokio::spawn(
+            StatsCollector {
+                interface_name,
+                bandwidth,
+            }
+            .run(tx),
+        );
+
+        Self { stats: rx }
+    }
+
+    /// Retrieves the next network statistics not yet seen by this instance, so a fresh
+    /// clone returns the latest statistics right away, or None once the collector has stopped.
     ///
     /// # Returns
     /// NetworkStats containing maximum and current bandwidth information.
-    pub async fn get_stats(&self) -> NetworkStats {
-        // Lock the mutex to ensure exclusive access to network stats.
-        let _guard = self.mutex.lock().await;
-
-        // Initialize sysinfo network.
-        let mut networks = Networks::new_with_refreshed_list();
-
-        // Sleep to calculate the network traffic difference over
-        // the DEFAULT_NETWORK_REFRESH_INTERVAL.
-        tokio::time::sleep(Self::DEFAULT_NETWORK_REFRESH_INTERVAL).await;
-
-        // Refresh network information to get updated statistics.
-        networks.refresh(true);
-        let Some(network_stats) = networks.get(self.interface_name.as_str()) else {
-            warn!(
-                "can not find network data for interface {}",
-                self.interface_name
-            );
-            return NetworkStats {
-                max_rx_bandwidth: self.bandwidth,
-                max_tx_bandwidth: self.bandwidth,
-                ..Default::default()
-            };
-        };
-
-        // Calculate the receive bandwidth in bits per second.
-        let rx_bandwidth = (Self::bytes_to_bits(network_stats.received()) as f64
-            / Self::DEFAULT_NETWORK_REFRESH_INTERVAL.as_secs_f64())
-        .round() as u64;
-
-        // Calculate the transmit bandwidth in bits per second.
-        let tx_bandwidth = (Self::bytes_to_bits(network_stats.transmitted()) as f64
-            / Self::DEFAULT_NETWORK_REFRESH_INTERVAL.as_secs_f64())
-        .round() as u64;
-
-        debug!(
-            "network interface {} max receive bandwidth: {} bps, receive bandwidth: {} bps, max transmit bandwidth: {} bps, transmit bandwidth: {} bps",
-            self.interface_name, self.bandwidth, rx_bandwidth, self.bandwidth, tx_bandwidth
-        );
-
-        NetworkStats {
-            max_rx_bandwidth: self.bandwidth,
-            rx_bandwidth: Some(rx_bandwidth),
-            max_tx_bandwidth: self.bandwidth,
-            tx_bandwidth: Some(tx_bandwidth),
-        }
+    pub async fn get_stats(&mut self) -> Option<NetworkStats> {
+        self.stats.changed().await.ok()?;
+        self.stats.borrow_and_update().clone()
     }
 
     /// Retrieves the speed of the network interface in megabits per second (Mbps).
@@ -239,10 +187,110 @@ impl Network {
     }
 }
 
+/// Collects the traffic statistics of a network interface.
+#[derive(Debug)]
+struct StatsCollector {
+    /// The name of the network interface.
+    interface_name: String,
+
+    /// The bandwidth of the network interface in bits per second (bps).
+    bandwidth: u64,
+}
+
+impl StatsCollector {
+    /// Default interval for refreshing network statistics.
+    const DEFAULT_NETWORK_REFRESH_INTERVAL: Duration = Duration::from_secs(1);
+
+    /// Collects the statistics continuously and publishes them until every receiver is dropped.
+    async fn run(self, tx: watch::Sender<Option<NetworkStats>>) {
+        loop {
+            let stats = self.collect().await;
+            if tx.send(Some(stats)).is_err() {
+                return;
+            }
+        }
+    }
+
+    /// Measures network traffic over DEFAULT_NETWORK_REFRESH_INTERVAL to calculate
+    /// current receive and transmit bandwidth.
+    async fn collect(&self) -> NetworkStats {
+        // Initialize sysinfo network.
+        let mut networks = Networks::new_with_refreshed_list();
+
+        // Sleep to calculate the network traffic difference over
+        // the DEFAULT_NETWORK_REFRESH_INTERVAL.
+        tokio::time::sleep(Self::DEFAULT_NETWORK_REFRESH_INTERVAL).await;
+
+        // Refresh network information to get updated statistics.
+        networks.refresh(true);
+        let Some(network_stats) = networks.get(self.interface_name.as_str()) else {
+            warn!(
+                "can not find network data for interface {}",
+                self.interface_name
+            );
+            return NetworkStats {
+                max_rx_bandwidth: self.bandwidth,
+                max_tx_bandwidth: self.bandwidth,
+                ..Default::default()
+            };
+        };
+
+        // Calculate the receive bandwidth in bits per second.
+        let rx_bandwidth = (Network::bytes_to_bits(network_stats.received()) as f64
+            / Self::DEFAULT_NETWORK_REFRESH_INTERVAL.as_secs_f64())
+        .round() as u64;
+
+        // Calculate the transmit bandwidth in bits per second.
+        let tx_bandwidth = (Network::bytes_to_bits(network_stats.transmitted()) as f64
+            / Self::DEFAULT_NETWORK_REFRESH_INTERVAL.as_secs_f64())
+        .round() as u64;
+
+        debug!(
+            "network interface {} max receive bandwidth: {} bps, receive bandwidth: {} bps, max transmit bandwidth: {} bps, transmit bandwidth: {} bps",
+            self.interface_name, self.bandwidth, rx_bandwidth, self.bandwidth, tx_bandwidth
+        );
+
+        NetworkStats {
+            max_rx_bandwidth: self.bandwidth,
+            rx_bandwidth: Some(rx_bandwidth),
+            max_tx_bandwidth: self.bandwidth,
+            tx_bandwidth: Some(tx_bandwidth),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use bytesize::ByteSize;
+    use std::net::Ipv4Addr;
+    use std::time::Instant;
+    use tokio::task::JoinSet;
+
+    #[tokio::test]
+    async fn test_get_stats() {
+        let mut network = Network::new(IpAddr::V4(Ipv4Addr::LOCALHOST), ByteSize::mb(100));
+
+        let start = Instant::now();
+        let mut join_set = JoinSet::new();
+        for _ in 0..10 {
+            let mut network = network.clone();
+            join_set.spawn(async move { network.get_stats().await });
+        }
+        while let Some(stats) = join_set.join_next().await {
+            assert!(stats.unwrap().is_some());
+        }
+        assert!(start.elapsed() < StatsCollector::DEFAULT_NETWORK_REFRESH_INTERVAL * 2);
+
+        let start = Instant::now();
+        assert!(network.get_stats().await.is_some());
+        assert!(start.elapsed() < StatsCollector::DEFAULT_NETWORK_REFRESH_INTERVAL / 2);
+
+        let start = Instant::now();
+        assert!(network.get_stats().await.is_some());
+        assert!(start.elapsed() >= StatsCollector::DEFAULT_NETWORK_REFRESH_INTERVAL / 2);
+        assert!(start.elapsed() < StatsCollector::DEFAULT_NETWORK_REFRESH_INTERVAL * 2);
+    }
 
     #[test]
     fn test_byte_size_to_bits() {