diff --git a/Cargo.lock b/Cargo.lock index ce2edb177..6f4c92f0e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1288,7 +1288,7 @@ dependencies = [ [[package]] name = "dragonfly-client" -version = "1.5.3" +version = "1.5.4" dependencies = [ "anyhow", "async-trait", @@ -1365,7 +1365,7 @@ dependencies = [ [[package]] name = "dragonfly-client-backend" -version = "1.5.3" +version = "1.5.4" dependencies = [ "async-trait", "bytes", @@ -1403,7 +1403,7 @@ dependencies = [ [[package]] name = "dragonfly-client-config" -version = "1.5.3" +version = "1.5.4" dependencies = [ "bytesize", "bytesize-serde", @@ -1433,7 +1433,7 @@ dependencies = [ [[package]] name = "dragonfly-client-core" -version = "1.5.3" +version = "1.5.4" dependencies = [ "cgroups-rs", "http", @@ -1453,7 +1453,7 @@ dependencies = [ [[package]] name = "dragonfly-client-init" -version = "1.5.3" +version = "1.5.4" dependencies = [ "anyhow", "clap", @@ -1470,7 +1470,7 @@ dependencies = [ [[package]] name = "dragonfly-client-metric" -version = "1.5.3" +version = "1.5.4" dependencies = [ "bytes", "dragonfly-api", @@ -1522,7 +1522,7 @@ dependencies = [ [[package]] name = "dragonfly-client-storage" -version = "1.5.3" +version = "1.5.4" dependencies = [ "bincode", "bytes", @@ -1557,7 +1557,7 @@ dependencies = [ [[package]] name = "dragonfly-client-util" -version = "1.5.3" +version = "1.5.4" dependencies = [ "async-trait", "base64 0.22.1", @@ -2092,7 +2092,7 @@ dependencies = [ [[package]] name = "hdfs" -version = "1.5.3" +version = "1.5.4" dependencies = [ "async-trait", "dragonfly-client-backend", diff --git a/Cargo.toml b/Cargo.toml index 85177b93d..0b2965968 100644 --- a/Cargo.toml +++ b/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" diff --git a/dragonfly-client-util/src/sysinfo/network.rs b/dragonfly-client-util/src/sysinfo/network.rs index 255cb2da3..ad32f275e 100644 --- a/dragonfly-client-util/src/sysinfo/network.rs +++ b/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,17 +39,11 @@ pub struct NetworkStats { pub tx_bandwidth: Option, } -/// 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>, + // The latest statistics published by the collector, None until the first collection. + stats: watch::Receiver>, } /// Implementation of network monitoring functionality. @@ -58,10 +51,8 @@ pub struct Network { /// 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,11 +78,7 @@ impl Network { interface.name, bandwidth ); - Self { - interface_name: interface.name, - bandwidth, - mutex: Arc::new(Mutex::new(())), - } + Self::spawn(interface.name, bandwidth) } None => { warn!( @@ -103,68 +86,33 @@ impl Network { 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 { + 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>) { + 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() { diff --git a/dragonfly-client/src/announcer/mod.rs b/dragonfly-client/src/announcer/mod.rs index f0d0e2bdb..cfbd6eafb 100644 --- a/dragonfly-client/src/announcer/mod.rs +++ b/dragonfly-client/src/announcer/mod.rs @@ -24,7 +24,7 @@ use dragonfly_client_config::{ CARGO_PKG_RUSTC_VERSION, CARGO_PKG_VERSION, GIT_COMMIT_SHORT_HASH, INSTANCE_NAME, }; use dragonfly_client_core::error::{ErrorType, OrErr}; -use dragonfly_client_core::Result; +use dragonfly_client_core::{Error, Result}; use dragonfly_client_util::{container::is_running_in_container, shutdown, sysinfo::SystemMonitor}; use std::env; use std::sync::Arc; @@ -188,7 +188,13 @@ impl SchedulerAnnouncer { } // Wait for getting the network data. - let network_stats = self.system_monitor.network.get_stats().await; + let network_stats = self + .system_monitor + .network + .clone() + .get_stats() + .await + .ok_or_else(|| Error::Unknown("network stats collector stopped".to_string()))?; debug!( "network data: rx bandwidth {}/{} bps, tx bandwidth {}/{} bps", network_stats.rx_bandwidth.unwrap_or(0), diff --git a/dragonfly-client/src/grpc/dfdaemon_upload.rs b/dragonfly-client/src/grpc/dfdaemon_upload.rs index 777a87109..deb89ae56 100644 --- a/dragonfly-client/src/grpc/dfdaemon_upload.rs +++ b/dragonfly-client/src/grpc/dfdaemon_upload.rs @@ -1176,7 +1176,7 @@ impl DfdaemonUpload for DfdaemonUploadServerHandler { info!("sync host in upload server"); // Clone the network monitor. - let network = self.system_monitor.network.clone(); + let mut network = self.system_monitor.network.clone(); // Initialize stream channel. let (out_stream_tx, out_stream_rx) = mpsc::channel(128); @@ -1184,8 +1184,11 @@ impl DfdaemonUpload for DfdaemonUploadServerHandler { async move { // Start the host info update loop. loop { - // Wait for getting the network stats. - let network_stats = network.get_stats().await; + // Wait for the next network stats. + let Some(network_stats) = network.get_stats().await else { + return; + }; + debug!( "network data: rx bandwidth {}/{} bps, tx bandwidth {}/{} bps", network_stats.rx_bandwidth.unwrap_or(0),