diff --git a/Cargo.lock b/Cargo.lock index cfe94ffe..b99cc38f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -797,7 +797,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8599749b6667e2f0c910c1d0dff6901163ff698a52d5a39720f61b5be4b20d3" dependencies = [ "futures-core", - "prost 0.14.3", + "prost 0.14.4", "prost-types", "tonic", "tonic-prost", @@ -817,7 +817,7 @@ dependencies = [ "hdrhistogram", "humantime", "hyper-util", - "prost 0.14.3", + "prost 0.14.4", "prost-types", "serde", "serde_json", @@ -1285,11 +1285,11 @@ dependencies = [ [[package]] name = "dragonfly-api" -version = "2.2.30" +version = "2.2.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc71e0b2ce9c6315d54945d1c64de28a210a93aed009e981e87f35f2895482e7" +checksum = "cf76fa86d7330ac30c27189315fe5f4692bfd80c68a7616100f994c6d2699651" dependencies = [ - "prost 0.14.3", + "prost 0.14.4", "prost-types", "prost-wkt-types", "serde", @@ -3362,7 +3362,7 @@ dependencies = [ "http-body-util", "hyper 1.9.0", "hyper-util", - "prost 0.14.3", + "prost 0.14.4", "rand 0.9.4", "serde", "serde_json", @@ -3874,7 +3874,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-proto", "opentelemetry_sdk", - "prost 0.14.3", + "prost 0.14.4", "reqwest 0.12.28", "thiserror 2.0.18", "tokio", @@ -3889,7 +3889,7 @@ checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" dependencies = [ "opentelemetry", "opentelemetry_sdk", - "prost 0.14.3", + "prost 0.14.4", "tonic", "tonic-prost", ] @@ -4490,12 +4490,12 @@ dependencies = [ [[package]] name = "prost" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", - "prost-derive 0.14.3", + "prost-derive 0.14.4", ] [[package]] @@ -4510,7 +4510,7 @@ dependencies = [ "multimap", "petgraph", "prettyplease", - "prost 0.14.3", + "prost 0.14.4", "prost-types", "pulldown-cmark", "pulldown-cmark-to-cmark", @@ -4534,9 +4534,9 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", "itertools", @@ -4551,7 +4551,7 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" dependencies = [ - "prost 0.14.3", + "prost 0.14.4", ] [[package]] @@ -4562,7 +4562,7 @@ checksum = "cd3de5e9c9e84fcb5efa204b8e283d23e615a8bc8c777bf1d6622bb01dc61445" dependencies = [ "chrono", "inventory", - "prost 0.14.3", + "prost 0.14.4", "serde", "serde_derive", "serde_json", @@ -4576,7 +4576,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe500dc80e757a75e1e8fb7290e448d62dfba3105ece1d058579cb00b58151cd" dependencies = [ "heck", - "prost 0.14.3", + "prost 0.14.4", "prost-build", "prost-types", "quote", @@ -4589,7 +4589,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13807eaa7e15833d06e899008371926201cdcd11d74b6d490f49130cdb3f415e" dependencies = [ "chrono", - "prost 0.14.3", + "prost 0.14.4", "prost-build", "prost-types", "prost-wkt", @@ -6313,7 +6313,7 @@ version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f4ff0636fef47afb3ec02818f5bceb4377b8abb9d6a386aeade18bd6212f8eb7" dependencies = [ - "prost 0.14.3", + "prost 0.14.4", "tokio", "tokio-stream", "tonic", @@ -6327,7 +6327,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" dependencies = [ "bytes", - "prost 0.14.3", + "prost 0.14.4", "tonic", ] @@ -6353,7 +6353,7 @@ version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aaf0685a51e6d02b502ba0764002e766b7f3042aed13d9234925b6ffbfa3fca7" dependencies = [ - "prost 0.14.3", + "prost 0.14.4", "prost-types", "tokio", "tokio-stream", diff --git a/Cargo.toml b/Cargo.toml index 6de43fa7..7ec45f89 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,7 +35,7 @@ chrono = { version = "0.4.45", features = ["clock", "serde"] } clap = { version = "4.6.1", features = ["derive", "env"] } crc32fast = "1.5.0" dashmap = "6.1.0" -dragonfly-api = "=2.2.30" +dragonfly-api = "=2.2.31" dragonfly-client = { path = "dragonfly-client", version = "1.4.0" } dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.4.0" } dragonfly-client-config = { path = "dragonfly-client-config", version = "1.4.0" } diff --git a/dragonfly-client-util/src/sysinfo/disk.rs b/dragonfly-client-util/src/sysinfo/disk.rs index 7aa7478e..671acb71 100644 --- a/dragonfly-client-util/src/sysinfo/disk.rs +++ b/dragonfly-client-util/src/sysinfo/disk.rs @@ -48,6 +48,16 @@ pub struct ProcessDiskStats { pub read_bandwidth: u64, } +/// Represents disk I/O statistics from cgroup (Linux containers/resource limits). +#[derive(Debug, Clone, Default)] +pub struct CgroupDiskStats { + /// The write bandwidth of the cgroup at the block-device level in bytes per second. + pub write_bandwidth: u64, + + /// The read bandwidth of the cgroup at the block-device level in bytes per second. + pub read_bandwidth: u64, +} + /// Represents a disk interface for monitoring disk statistics. #[derive(Debug, Clone, Default)] pub struct Disk { @@ -57,8 +67,10 @@ pub struct Disk { /// Implementation of disk monitoring functionality. /// -/// Provides methods to retrieve disk space information and process-specific -/// disk I/O statistics, including read/write bandwidth measurements. +/// Provides methods to retrieve disk statistics at three different levels: +/// - Path-level: Total, free and used space of the filesystem. +/// - Process-level: I/O bandwidth of a specific process. +/// - Cgroup-level: Block-device I/O bandwidth of the cgroup (Linux only). impl Disk { /// Default interval for refreshing disk statistics. const DEFAULT_DISK_REFRESH_INTERVAL: Duration = Duration::from_secs(1); @@ -104,11 +116,12 @@ impl Disk { }) } - /// Retrieves the disk I/O statistics for a specific process. + /// Retrieves the disk I/O statistics for the process with the given PID. /// - /// This method measures disk I/O activity over a time interval - /// (DEFAULT_DISK_REFRESH_INTERVAL) to calculate read and write bandwidth - /// for the specified process. + /// This method measures the process's /proc I/O counters over a time + /// interval (DEFAULT_DISK_REFRESH_INTERVAL) to calculate read and write + /// bandwidth. Writes are accounted at page-dirtying time rather than at + /// the block device. /// /// # Arguments /// * `pid` - The process ID to monitor for disk I/O statistics. @@ -151,4 +164,245 @@ impl Disk { read_bandwidth, } } + + /// Retrieves disk I/O statistics from the cgroup (Linux only). + /// + /// This method measures the cumulative block-device I/O counters of the + /// cgroup the process belongs to over a time interval + /// (DEFAULT_DISK_REFRESH_INTERVAL) to calculate device-level read and + /// write bandwidth. cgroup v2 supports writeback attribution, so flushed + /// page cache writes are charged to the cgroup that dirtied the pages, + /// while cgroup v1 only accounts direct and synchronous I/O. + /// + /// # Arguments + /// * `pid` - Process ID used to determine which cgroup to query. + /// + /// # Returns + /// Some(CgroupDiskStats) if cgroup I/O statistics are available and accessible, + /// None otherwise or on non-Linux platforms. + #[allow(unused_variables)] + pub async fn get_cgroup_stats(&self, pid: u32) -> Option { + #[cfg(target_os = "linux")] + { + // Lock the mutex to ensure exclusive access to disk stats. + let _guard = self.mutex.lock().await; + + // Take a baseline of the cumulative block-device I/O counters. + let (baseline_read_bytes, baseline_written_bytes) = Self::get_cgroup_io_counters(pid)?; + + // Sleep to calculate the disk traffic difference over + // the DEFAULT_DISK_REFRESH_INTERVAL. + tokio::time::sleep(Self::DEFAULT_DISK_REFRESH_INTERVAL).await; + + let (read_bytes, written_bytes) = Self::get_cgroup_io_counters(pid)?; + + // Calculate the write bandwidth in bytes per second. + let write_bandwidth = (written_bytes.saturating_sub(baseline_written_bytes) as f64 + / Self::DEFAULT_DISK_REFRESH_INTERVAL.as_secs_f64()) + .round() as u64; + + // Calculate the read bandwidth in bytes per second. + let read_bandwidth = (read_bytes.saturating_sub(baseline_read_bytes) as f64 + / Self::DEFAULT_DISK_REFRESH_INTERVAL.as_secs_f64()) + .round() as u64; + + debug!( + "process {} cgroup disk write bandwidth: {} bytes/s, read bandwidth: {} bytes/s", + pid, write_bandwidth, read_bandwidth + ); + + Some(CgroupDiskStats { + write_bandwidth, + read_bandwidth, + }) + } + + #[cfg(not(target_os = "linux"))] + None + } + + /// Retrieves the cumulative bytes read from and written to the block + /// devices by the cgroup the process belongs to. + /// + /// # Arguments + /// * `pid` - The process ID used to determine which cgroup to read. + /// + /// # Returns + /// Some((read_bytes, written_bytes)) summed across all block devices if + /// the cgroup I/O statistics are accessible, None otherwise. + #[cfg(target_os = "linux")] + fn get_cgroup_io_counters(pid: u32) -> Option<(u64, u64)> { + use cgroups_rs::fs::hierarchies; + + if hierarchies::auto().v2() { + Self::get_cgroup_v2_io_counters(pid) + } else { + Self::get_cgroup_v1_io_counters(pid) + } + } + + /// Retrieves the cumulative block-device I/O counters from the cgroup v2 + /// io.stat file. + /// + /// # Arguments + /// * `pid` - The process ID used to determine which cgroup to read. + /// + /// # Returns + /// Some((read_bytes, written_bytes)) summed across all block devices, + /// None if the io.stat file is unavailable. + #[cfg(target_os = "linux")] + fn get_cgroup_v2_io_counters(pid: u32) -> Option<(u64, u64)> { + use crate::cgroups::get_cgroup_v2_path_by_pid; + use tracing::error; + + let path = match get_cgroup_v2_path_by_pid(pid) { + Ok(path) => path.join("io.stat"), + Err(err) => { + error!("failed to get cgroup v2 path for pid {}: {}", pid, err); + return None; + } + }; + + match std::fs::read_to_string(&path) { + Ok(content) => Some(Self::parse_cgroup_v2_io_stat(&content)), + Err(err) => { + error!("failed to read {}: {}", path.display(), err); + None + } + } + } + + /// Parses the content of a cgroup v2 io.stat file and sums the read and + /// written bytes across all block devices. + /// + /// Each line is `MAJ:MIN rbytes=N wbytes=N rios=N wios=N dbytes=N dios=N`. + /// + /// # Arguments + /// * `content` - The content of the io.stat file. + /// + /// # Returns + /// (read_bytes, written_bytes) summed across all block devices. + #[cfg(target_os = "linux")] + fn parse_cgroup_v2_io_stat(content: &str) -> (u64, u64) { + let (mut read_bytes, mut written_bytes) = (0u64, 0u64); + for part in content.split_whitespace() { + if let Some(value) = part.strip_prefix("rbytes=") { + read_bytes = read_bytes.saturating_add(value.parse().unwrap_or(0)); + } else if let Some(value) = part.strip_prefix("wbytes=") { + written_bytes = written_bytes.saturating_add(value.parse().unwrap_or(0)); + } + } + + (read_bytes, written_bytes) + } + + /// Retrieves the cumulative block-device I/O counters from the cgroup v1 + /// blkio controller. + /// + /// cgroup v1 charges page cache writeback to the root cgroup, so only + /// direct and synchronous I/O are accounted here. + /// + /// # Arguments + /// * `pid` - The process ID used to determine which cgroup to read. + /// + /// # Returns + /// Some((read_bytes, written_bytes)) summed across all block devices, + /// None if the blkio controller is unavailable. + #[cfg(target_os = "linux")] + fn get_cgroup_v1_io_counters(pid: u32) -> Option<(u64, u64)> { + use crate::cgroups::get_cgroup_by_pid; + use cgroups_rs::fs::blkio::BlkIoController; + use tracing::error; + + let cgroup = match get_cgroup_by_pid(pid) { + Ok(cgroup) => cgroup, + Err(err) => { + error!("failed to get cgroup for pid {}: {}", pid, err); + return None; + } + }; + + let Some(blkio_controller) = cgroup.controller_of::() else { + error!("no blkio controller found for pid {}", pid); + return None; + }; + + // blkio.throttle.io_service_bytes is maintained by any I/O scheduler, + // while blkio.io_service_bytes is only maintained by CFQ. + let blkio = blkio_controller.blkio(); + let io_service_bytes = if !blkio.throttle.io_service_bytes.is_empty() { + blkio.throttle.io_service_bytes + } else { + blkio.io_service_bytes + }; + + Some(Self::parse_cgroup_v1_io_stat(&io_service_bytes)) + } + + /// Parses the cgroup v1 blkio IoService entries and sums the read and + /// written bytes across all block devices. + /// + /// # Arguments + /// * `io_service_bytes` - The per-device IoService entries from + /// blkio.throttle.io_service_bytes or blkio.io_service_bytes. + /// + /// # Returns + /// (read_bytes, written_bytes) summed across all block devices. + #[cfg(target_os = "linux")] + fn parse_cgroup_v1_io_stat( + io_service_bytes: &[cgroups_rs::fs::blkio::IoService], + ) -> (u64, u64) { + io_service_bytes + .iter() + .fold((0u64, 0u64), |(read_bytes, written_bytes), io_service| { + ( + read_bytes.saturating_add(io_service.read), + written_bytes.saturating_add(io_service.write), + ) + }) + } +} + +#[cfg(all(test, target_os = "linux"))] +mod tests { + use super::*; + + #[test] + fn test_parse_cgroup_v2_io_stat() { + let content = "8:0 rbytes=90430464 wbytes=299008000 rios=8950 wios=1252 dbytes=50331648 dios=3021\n253:0 rbytes=1459200 wbytes=314773504 rios=192 wios=353 dbytes=0 dios=0"; + assert_eq!( + Disk::parse_cgroup_v2_io_stat(content), + (90430464 + 1459200, 299008000 + 314773504) + ); + + assert_eq!(Disk::parse_cgroup_v2_io_stat(""), (0, 0)); + } + + #[test] + fn test_parse_cgroup_v1_io_stat() { + use cgroups_rs::fs::blkio::IoService; + + let io_service_bytes = vec![ + IoService { + major: 8, + minor: 0, + read: 90430464, + write: 299008000, + ..Default::default() + }, + IoService { + major: 253, + minor: 0, + read: 1459200, + write: 314773504, + ..Default::default() + }, + ]; + assert_eq!( + Disk::parse_cgroup_v1_io_stat(&io_service_bytes), + (90430464 + 1459200, 299008000 + 314773504) + ); + + assert_eq!(Disk::parse_cgroup_v1_io_stat(&[]), (0, 0)); + } } diff --git a/dragonfly-client-util/src/sysinfo/memory.rs b/dragonfly-client-util/src/sysinfo/memory.rs index 977a1008..74221b19 100644 --- a/dragonfly-client-util/src/sysinfo/memory.rs +++ b/dragonfly-client-util/src/sysinfo/memory.rs @@ -149,24 +149,28 @@ impl Memory { Ok(cgroup) => { if let Some(memory_controller) = cgroup.controller_of::() { let memory_stats = memory_controller.memory_stat(); + // The `rss` key only exists in cgroup v1's memory.stat. In cgroup v2 the + // equivalent counter (NR_ANON_MAPPED) is exported as `anon`. + let memory_usage = if cgroup.v2() { + memory_stats.stat.raw.get("anon").copied().unwrap_or(0) + } else { + memory_stats.stat.rss + }; + let used_percent = if memory_stats.limit_in_bytes > 0 { - (memory_stats.stat.rss as f64 / memory_stats.limit_in_bytes as f64) - * 100.0 + (memory_usage as f64 / memory_stats.limit_in_bytes as f64) * 100.0 } else { - (memory_stats.stat.rss as f64 / self.get_stats().total as f64) * 100.0 + (memory_usage as f64 / self.get_stats().total as f64) * 100.0 }; debug!( "process {} cgroup memory limit: {} bytes, memory usage: {} bytes, used percent: {}%", - pid, - memory_stats.limit_in_bytes, - memory_stats.stat.rss, - used_percent, + pid, memory_stats.limit_in_bytes, memory_usage, used_percent, ); return Some(CgroupMemoryStats { limit: memory_stats.limit_in_bytes, - usage: memory_stats.stat.rss, + usage: memory_usage, used_percent: used_percent.clamp(0.0, 100.0), }); } diff --git a/dragonfly-client/src/announcer/mod.rs b/dragonfly-client/src/announcer/mod.rs index d8e97d0d..f0d0e2bd 100644 --- a/dragonfly-client/src/announcer/mod.rs +++ b/dragonfly-client/src/announcer/mod.rs @@ -15,7 +15,9 @@ */ use crate::grpc::scheduler::SchedulerClient; -use dragonfly_api::common::v2::{Build, CgroupCpu, CgroupMemory, Cpu, Disk, Host, Memory, Network}; +use dragonfly_api::common::v2::{ + Build, CgroupCpu, CgroupDisk, CgroupMemory, Cpu, Disk, Host, Memory, Network, +}; use dragonfly_api::scheduler::v2::{AnnounceHostRequest, DeleteHostRequest}; use dragonfly_client_config::{ dfdaemon::{Config, HostType}, @@ -212,13 +214,14 @@ impl SchedulerAnnouncer { .disk .get_stats(self.config.storage.dir.as_path())?; let process_disk_stats = self.system_monitor.disk.get_process_stats(pid).await; - let disk = Disk { + let mut disk = Disk { total: disk_stats.total, free: disk_stats.free, used: disk_stats.usage, used_percent: disk_stats.used_percent, write_bandwidth: process_disk_stats.write_bandwidth, read_bandwidth: process_disk_stats.read_bandwidth, + cgroup: None, // TODO: Get the disk inodes information. inodes_total: 0, @@ -227,6 +230,19 @@ impl SchedulerAnnouncer { inodes_used_percent: 0.0, }; + // Get the cgroup disk information if running in a container environment. + if self.is_running_in_container { + disk.cgroup = self + .system_monitor + .disk + .get_cgroup_stats(pid) + .await + .map(|stats| CgroupDisk { + read_bandwidth: stats.read_bandwidth, + write_bandwidth: stats.write_bandwidth, + }); + } + // Get the build information. let build = Build { git_version: CARGO_PKG_VERSION.to_string(),