joule_profiler_source_cgroup/
lib.rs

1//! cgroup metric source for Joule Profiler.
2//!
3//! This module implements a [`MetricReader`] and uses Linux cgroup v2 to
4//! collect per-process and global system metrics using kernel cgroup files.
5//!
6//! An asynchronous tokio task runs to poll non monotonic metrics.
7
8use futures::StreamExt;
9use joule_profiler_core::sensor::{Sensor, Sensors};
10use joule_profiler_core::source::MetricReader;
11use joule_profiler_core::time::get_timestamp_micros;
12use joule_profiler_core::types::{Metric, MetricValue, Metrics};
13use joule_profiler_core::unit::{MetricUnit, Unit, UnitPrefix};
14use log::{debug, trace, warn};
15use std::sync::{Arc, Mutex};
16use std::time::Duration;
17use tokio::task::JoinHandle;
18use tokio_timerfd::Interval;
19use tokio_util::sync::CancellationToken;
20
21use crate::cgroup::{CgroupBackend, ChildCgroup, RootCgroup, SysFsBackend};
22use crate::counters::{Counters, CpuCounters, IoCounters, MemoryCounters};
23use crate::error::CgroupError;
24
25mod cgroup;
26mod config;
27mod counters;
28mod error;
29mod snapshot;
30mod util;
31
32pub use config::CgroupConfig;
33
34const SOURCE_NAME: &str = "cgroup";
35
36const BYTE_UNIT: MetricUnit = MetricUnit {
37    prefix: UnitPrefix::None,
38    unit: Unit::Byte,
39};
40
41const MICRO_SECOND_UNIT: MetricUnit = MetricUnit {
42    prefix: UnitPrefix::Micro,
43    unit: Unit::Second,
44};
45
46const COUNT_UNIT: MetricUnit = MetricUnit {
47    prefix: UnitPrefix::None,
48    unit: Unit::Count,
49};
50
51pub(crate) type Result<T> = std::result::Result<T, CgroupError>;
52pub(crate) type WorkerHandle = (CancellationToken, JoinHandle<Result<()>>);
53
54/// cgroup metrics source.
55///
56/// Owns the process cgroup and root cgroup handles, and maintains
57/// internal counters for both process-level and system-wide metrics.
58pub struct Cgroup<B = SysFsBackend>
59where
60    B: CgroupBackend,
61{
62    config: CgroupConfig,
63    handle: Option<WorkerHandle>,
64    pub proc_cgroup: Arc<ChildCgroup<B>>,
65    pub root_cgroup: Arc<RootCgroup<B>>,
66    proc_memory_counters: Arc<Mutex<MemoryCounters>>,
67    global_memory_counters: Arc<Mutex<MemoryCounters>>,
68    proc_cpu_counters: CpuCounters,
69    global_cpu_counters: CpuCounters,
70    proc_io_counters: IoCounters,
71    global_io_counters: IoCounters,
72    begin_timestamp: u128,
73    end_timestamp: u128,
74}
75
76impl<B: CgroupBackend> Cgroup<B> {
77    /// Creates a background worker that periodically samples memory usage.
78    ///
79    /// This worker updates process and global memory counters at a fixed interval.
80    /// It can be cancelled using the returned `CancellationToken`.
81    pub fn create_worker(
82        proc_cgroup: Arc<ChildCgroup<B>>,
83        root_cgroup: Arc<RootCgroup<B>>,
84        proc_memory_counters: Arc<Mutex<MemoryCounters>>,
85        global_memory_counters: Arc<Mutex<MemoryCounters>>,
86        poll_interval: Duration,
87    ) -> Result<WorkerHandle> {
88        let mut ticker = Interval::new_interval(poll_interval).map_err(CgroupError::Timer)?;
89
90        let cancellation_token = CancellationToken::new();
91        let cancellation_token_clone = cancellation_token.clone();
92
93        let handle = tokio::spawn(async move {
94            debug!("Starting cgroup memory polling.");
95
96            loop {
97                tokio::select! {
98                    _ = ticker.next() => {
99                        trace!("Polled cgroup source.");
100                        {
101                            let snapshot = proc_cgroup.memory()?;
102                            let mut lock = proc_memory_counters
103                                .lock()
104                                .map_err(|_| CgroupError::MutexPoisoned)?;
105                            lock.update(&snapshot);
106                        }
107                        {
108                            let snapshot = root_cgroup.memory()?;
109                            let mut lock = global_memory_counters
110                                .lock()
111                                .map_err(|_| CgroupError::MutexPoisoned)?;
112                            lock.update(&snapshot);
113                        }
114                    }
115
116                    () = cancellation_token.cancelled() => {
117                        debug!("Cgroup worker stopped.");
118                        break;
119                    }
120                }
121            }
122
123            Ok(())
124        });
125
126        Ok((cancellation_token_clone, handle))
127    }
128}
129
130impl<B: CgroupBackend> MetricReader for Cgroup<B> {
131    type Type = Counters;
132    type Error = CgroupError;
133    type Config = CgroupConfig;
134
135    /// Builds the cgroup handles.
136    fn from_config(config: Self::Config) -> Result<Self> {
137        let (root_cgroup, proc_cgroup) =
138            RootCgroup::build(config.cgroup_root.clone(), &config.cgroup_name);
139
140        Ok(Self {
141            config,
142            handle: None,
143            root_cgroup: Arc::new(root_cgroup),
144            proc_cgroup: Arc::new(proc_cgroup),
145            proc_memory_counters: Arc::default(),
146            global_memory_counters: Arc::default(),
147            proc_cpu_counters: CpuCounters::default(),
148            global_cpu_counters: CpuCounters::default(),
149            proc_io_counters: IoCounters::default(),
150            global_io_counters: IoCounters::default(),
151            begin_timestamp: 0,
152            end_timestamp: 0,
153        })
154    }
155
156    /// Creates the cgroup if configured, then checks that it is usable.
157    async fn pre_init(&mut self) -> Result<()> {
158        self.root_cgroup.verify()?;
159
160        if self.config.create_cgroup {
161            if self.proc_cgroup.verify().is_ok() {
162                warn!(
163                    "Cgroup {} already mounted, if a process has already been spawned inside the cgroup, then the peak metrics might be wrong.",
164                    self.config.cgroup_name
165                );
166            }
167            self.proc_cgroup.create()?;
168        }
169
170        self.proc_cgroup.verify()?;
171
172        Ok(())
173    }
174
175    /// Initializes the cgroup source with the given pid.
176    async fn init(&mut self, pid: i32) -> Result<()> {
177        if self.config.attach_pid {
178            self.proc_cgroup.attach_pid(pid)?;
179        }
180
181        self.handle = Some(Self::create_worker(
182            self.proc_cgroup.clone(),
183            self.root_cgroup.clone(),
184            self.proc_memory_counters.clone(),
185            self.global_memory_counters.clone(),
186            self.config.poll_interval,
187        )?);
188
189        self.begin_timestamp = get_timestamp_micros();
190
191        debug!("Cgroup source initialized for pid {pid}.");
192        Ok(())
193    }
194
195    /// Performs a measurement of all available metrics.
196    ///
197    /// Updates internal CPU, memory, and I/O counters.
198    async fn measure(&mut self) -> Result<()> {
199        debug!("Measure cgroup source.");
200
201        self.end_timestamp = get_timestamp_micros();
202        {
203            let snapshot = self.proc_cgroup.memory()?;
204            let mut lock = self
205                .proc_memory_counters
206                .lock()
207                .map_err(|_| CgroupError::MutexPoisoned)?;
208            lock.update(&snapshot);
209        }
210        self.proc_cpu_counters.update(&self.proc_cgroup.cpu()?);
211        self.proc_io_counters.update(&self.proc_cgroup.io()?);
212
213        {
214            let snapshot = self.root_cgroup.memory()?;
215            let mut lock = self
216                .global_memory_counters
217                .lock()
218                .map_err(|_| CgroupError::MutexPoisoned)?;
219            lock.update(&snapshot);
220        }
221        self.global_cpu_counters.update(&self.root_cgroup.cpu()?);
222        self.global_io_counters.update(&self.root_cgroup.io()?);
223
224        Ok(())
225    }
226
227    /// Returns collected metrics and resets per-phase counters.
228    async fn retrieve(&mut self) -> Result<Self::Type> {
229        debug!("Retrieving cgroup counters.");
230
231        let proc_memory = {
232            let mut lock = self
233                .proc_memory_counters
234                .lock()
235                .map_err(|_| CgroupError::MutexPoisoned)?;
236            let counters = *lock;
237            lock.reset();
238            counters
239        };
240
241        let proc_cpu = self.proc_cpu_counters;
242        self.proc_cpu_counters.new_phase();
243
244        let proc_io = self.proc_io_counters;
245        self.proc_io_counters.new_phase();
246
247        let global_memory = {
248            let mut lock = self
249                .global_memory_counters
250                .lock()
251                .map_err(|_| CgroupError::MutexPoisoned)?;
252            let counters = *lock;
253            lock.reset();
254            counters
255        };
256
257        let global_cpu = self.global_cpu_counters;
258        self.global_cpu_counters.new_phase();
259
260        let global_io = self.global_io_counters;
261        self.global_io_counters.new_phase();
262
263        let begin_timestamp = self.begin_timestamp;
264        self.begin_timestamp = self.end_timestamp;
265
266        Ok(Counters {
267            proc_memory,
268            proc_cpu,
269            proc_io,
270            global_memory,
271            global_cpu,
272            global_io,
273            begin_timestamp,
274            end_timestamp: self.end_timestamp,
275        })
276    }
277
278    /// Stops background worker and cleans up the cgroup.
279    async fn join(&mut self) -> Result<()> {
280        if let Some((cancellation_token, handle)) = self.handle.take() {
281            debug!("Joining cgroup source polling task.");
282            cancellation_token.cancel();
283            handle.await??;
284        }
285        if self.config.create_cgroup {
286            self.proc_cgroup.cleanup()?;
287        }
288        Ok(())
289    }
290
291    /// Returns the list of exported sensors.
292    fn get_sensors(&self) -> Result<Sensors> {
293        debug!("Retrieving cgroup source sensors.");
294        Ok(vec![
295            Sensor::new("usage_usec", MICRO_SECOND_UNIT, SOURCE_NAME),
296            Sensor::new("user_usec", MICRO_SECOND_UNIT, SOURCE_NAME),
297            Sensor::new("system_usec", MICRO_SECOND_UNIT, SOURCE_NAME),
298            Sensor::new("nr_periods", COUNT_UNIT, SOURCE_NAME),
299            Sensor::new("nr_throttled", COUNT_UNIT, SOURCE_NAME),
300            Sensor::new("throttled_usec", MICRO_SECOND_UNIT, SOURCE_NAME),
301            Sensor::new("nr_bursts", COUNT_UNIT, SOURCE_NAME),
302            Sensor::new("burst_usec", MICRO_SECOND_UNIT, SOURCE_NAME),
303            Sensor::new("anon_min", BYTE_UNIT, SOURCE_NAME),
304            Sensor::new("anon_max", BYTE_UNIT, SOURCE_NAME),
305            Sensor::new("file_min", BYTE_UNIT, SOURCE_NAME),
306            Sensor::new("file_max", BYTE_UNIT, SOURCE_NAME),
307            Sensor::new("kernel_min", BYTE_UNIT, SOURCE_NAME),
308            Sensor::new("kernel_max", BYTE_UNIT, SOURCE_NAME),
309            Sensor::new("kernel_stack_min", BYTE_UNIT, SOURCE_NAME),
310            Sensor::new("kernel_stack_max", BYTE_UNIT, SOURCE_NAME),
311            Sensor::new("peak_min", BYTE_UNIT, SOURCE_NAME),
312            Sensor::new("peak_max", BYTE_UNIT, SOURCE_NAME),
313            Sensor::new("shmem_min", BYTE_UNIT, SOURCE_NAME),
314            Sensor::new("shmem_max", BYTE_UNIT, SOURCE_NAME),
315            Sensor::new("slab_min", BYTE_UNIT, SOURCE_NAME),
316            Sensor::new("slab_max", BYTE_UNIT, SOURCE_NAME),
317            Sensor::new("swap_current_min", BYTE_UNIT, SOURCE_NAME),
318            Sensor::new("swap_current_max", BYTE_UNIT, SOURCE_NAME),
319            Sensor::new("read_bytes", BYTE_UNIT, SOURCE_NAME),
320            Sensor::new("write_bytes", BYTE_UNIT, SOURCE_NAME),
321        ])
322    }
323
324    #[allow(clippy::cast_precision_loss)]
325    /// Converts counters into metrics.
326    fn to_metrics(&self, counters: Self::Type) -> Result<Metrics> {
327        let phase_duration: u64 = counters
328            .end_timestamp
329            .saturating_sub(counters.begin_timestamp)
330            .try_into()
331            .unwrap_or_default();
332
333        let proc_cpu_usage_usec = counters.proc_cpu.usage_usec.diff();
334
335        let proc_usage_ratio = if phase_duration == 0 {
336            0.0
337        } else {
338            proc_cpu_usage_usec as f64 / phase_duration as f64 * 100.0
339        };
340
341        let proc_cpu_usage = Metric::new(
342            "proc_cpu_usage",
343            MetricValue::Float(proc_usage_ratio, Some(2)),
344            MetricUnit {
345                prefix: UnitPrefix::None,
346                unit: Unit::Percent,
347            },
348            Self::get_name(),
349        );
350
351        let global_cpu_usage_usec = counters.global_cpu.usage_usec.diff();
352
353        let global_usage_ratio = if phase_duration == 0 {
354            0.0
355        } else {
356            global_cpu_usage_usec as f64 / phase_duration as f64 * 100.0
357        };
358
359        let global_cpu_usage = Metric::new(
360            "global_cpu_usage",
361            MetricValue::Float(global_usage_ratio, Some(2)),
362            MetricUnit {
363                prefix: UnitPrefix::None,
364                unit: Unit::Percent,
365            },
366            Self::get_name(),
367        );
368
369        let metrics = vec![proc_cpu_usage, global_cpu_usage];
370
371        Ok(metrics
372            .into_iter()
373            .chain(to_metrics(
374                &counters.proc_memory,
375                &counters.proc_cpu,
376                &counters.proc_io,
377                "proc",
378            ))
379            .chain(to_metrics(
380                &counters.global_memory,
381                &counters.global_cpu,
382                &counters.global_io,
383                "global",
384            ))
385            .collect())
386    }
387
388    fn get_name() -> &'static str {
389        SOURCE_NAME
390    }
391
392    fn get_id() -> &'static str {
393        SOURCE_NAME
394    }
395}
396
397fn to_metrics(
398    memory: &MemoryCounters,
399    cpu: &CpuCounters,
400    io: &IoCounters,
401    prefix: &str,
402) -> Metrics {
403    macro_rules! push {
404        ($metrics:expr, $name:expr, $value:expr, $unit:expr) => {
405            $metrics.push(Metric::new(
406                format!("{prefix}_{}", $name),
407                $value,
408                $unit,
409                Cgroup::<SysFsBackend>::get_name(),
410            ));
411        };
412    }
413
414    macro_rules! push_minmax {
415        ($metrics:expr, $field:expr, $name:expr, $unit:expr) => {
416            if let Some(mm) = $field {
417                push!(
418                    $metrics,
419                    concat!($name, "_min"),
420                    mm.min().unwrap_or_default(),
421                    $unit
422                );
423                push!(
424                    $metrics,
425                    concat!($name, "_max"),
426                    mm.max().unwrap_or_default(),
427                    $unit
428                );
429            }
430        };
431    }
432
433    macro_rules! push_begin_end {
434        ($metrics:expr, $field:expr, $name:expr, $unit:expr) => {
435            if let Some(be) = $field {
436                push!($metrics, $name, be.diff(), $unit);
437            }
438        };
439    }
440
441    let mut metrics = Vec::new();
442
443    push_minmax!(metrics, memory.anon, "anon", BYTE_UNIT);
444    push_minmax!(metrics, memory.file, "file", BYTE_UNIT);
445    push_minmax!(metrics, memory.kernel, "kernel", BYTE_UNIT);
446    push_minmax!(metrics, memory.kernel_stack, "kernel_stack", BYTE_UNIT);
447    push_minmax!(metrics, memory.peak, "peak", BYTE_UNIT);
448    push_minmax!(metrics, memory.shmem, "shmem", BYTE_UNIT);
449    push_minmax!(metrics, memory.slab, "slab", BYTE_UNIT);
450    push_minmax!(metrics, memory.swap_current, "swap_current", BYTE_UNIT);
451
452    push!(
453        metrics,
454        "usage_usec",
455        cpu.usage_usec.diff(),
456        MICRO_SECOND_UNIT
457    );
458    push!(
459        metrics,
460        "user_usec",
461        cpu.user_usec.diff(),
462        MICRO_SECOND_UNIT
463    );
464    push!(
465        metrics,
466        "system_usec",
467        cpu.system_usec.diff(),
468        MICRO_SECOND_UNIT
469    );
470
471    push_begin_end!(metrics, cpu.nr_periods, "nr_periods", COUNT_UNIT);
472    push_begin_end!(metrics, cpu.nr_throttled, "nr_throttled", COUNT_UNIT);
473    push_begin_end!(
474        metrics,
475        cpu.throttled_usec,
476        "throttled_usec",
477        MICRO_SECOND_UNIT
478    );
479    push_begin_end!(metrics, cpu.nr_bursts, "nr_bursts", COUNT_UNIT);
480    push_begin_end!(metrics, cpu.burst_usec, "burst_usec", MICRO_SECOND_UNIT);
481
482    push_begin_end!(metrics, io.rbytes, "read_bytes", BYTE_UNIT);
483    push_begin_end!(metrics, io.wbytes, "write_bytes", BYTE_UNIT);
484
485    metrics
486}
487#[cfg(test)]
488mod tests {
489    use super::*;
490    use crate::snapshot::{CpuSnapshot, IoSnapshot, MemorySnapshot};
491    use std::path::{Path, PathBuf};
492    use std::sync::Arc;
493    use tokio::time::{Duration, sleep};
494
495    #[derive(Default, Clone)]
496    struct MockCgroupBackend {
497        memory: Arc<std::sync::Mutex<MemorySnapshot>>,
498        cpu: Arc<std::sync::Mutex<CpuSnapshot>>,
499        io: Arc<std::sync::Mutex<IoSnapshot>>,
500    }
501
502    impl CgroupBackend for MockCgroupBackend {
503        fn memory(&self, _path: &Path) -> Result<MemorySnapshot> {
504            Ok(*self.memory.lock().unwrap())
505        }
506
507        fn cpu(&self, _path: &Path) -> Result<CpuSnapshot> {
508            Ok(*self.cpu.lock().unwrap())
509        }
510
511        fn io(&self, _path: &Path) -> Result<IoSnapshot> {
512            Ok(*self.io.lock().unwrap())
513        }
514
515        fn cleanup(&self, _path: &Path, _root: &Path) -> Result<()> {
516            Ok(())
517        }
518
519        fn create(&self, _path: &Path) -> Result<()> {
520            Ok(())
521        }
522
523        fn attach_pid(&self, _path: &Path, _pid: i32) -> Result<()> {
524            Ok(())
525        }
526    }
527
528    fn setup_source() -> (Cgroup<MockCgroupBackend>, MockCgroupBackend) {
529        let backend = MockCgroupBackend::default();
530
531        let root = Arc::new(RootCgroup::new(PathBuf::from("/tmp/root"), backend.clone()));
532
533        let proc = Arc::new(ChildCgroup::new(
534            PathBuf::from("/tmp/cgroup"),
535            PathBuf::from("/tmp/root"),
536            backend.clone(),
537        ));
538
539        let source = Cgroup {
540            config: CgroupConfig::default(),
541            handle: None,
542            root_cgroup: root,
543            proc_cgroup: proc,
544            proc_memory_counters: Arc::default(),
545            global_memory_counters: Arc::default(),
546            proc_cpu_counters: CpuCounters::default(),
547            global_cpu_counters: CpuCounters::default(),
548            proc_io_counters: IoCounters::default(),
549            global_io_counters: IoCounters::default(),
550            begin_timestamp: 0,
551            end_timestamp: 0,
552        };
553
554        (source, backend)
555    }
556
557    #[tokio::test]
558    async fn test_measure_updates_counters() {
559        let (mut src, backend) = setup_source();
560
561        {
562            let mem = src.proc_memory_counters.lock().unwrap();
563            assert!(mem.anon.is_none());
564        }
565
566        {
567            backend.memory.lock().unwrap().anon = Some(200);
568        }
569
570        src.measure().await.unwrap();
571
572        {
573            let mem = src.proc_memory_counters.lock().unwrap();
574            let anon = mem.anon.unwrap();
575
576            assert_eq!(anon.min(), Some(200));
577            assert_eq!(anon.max(), Some(200));
578        }
579    }
580
581    #[tokio::test]
582    async fn test_measure_tracks_min_max() {
583        let (mut src, backend) = setup_source();
584
585        {
586            let mut memory = backend.memory.lock().unwrap();
587            memory.anon = Some(200);
588            memory.current = Some(200);
589            memory.peak = Some(200);
590        }
591
592        src.measure().await.unwrap();
593
594        {
595            let mut memory = backend.memory.lock().unwrap();
596            memory.anon = Some(400);
597            memory.current = Some(400);
598            memory.peak = Some(400);
599        }
600
601        src.measure().await.unwrap();
602
603        let mem = src.proc_memory_counters.lock().unwrap();
604        let anon = mem.anon.unwrap();
605        let current = mem.current.unwrap();
606        let peak = mem.peak.unwrap();
607
608        assert_eq!(anon.min(), Some(200));
609        assert_eq!(anon.max(), Some(400));
610
611        assert_eq!(current.min(), Some(200));
612        assert_eq!(current.max(), Some(400));
613
614        assert_eq!(peak.min(), Some(200));
615        assert_eq!(peak.max(), Some(400));
616    }
617
618    #[tokio::test]
619    async fn test_retrieve_compute_diffs() {
620        let (mut src, backend) = setup_source();
621
622        {
623            let mut cpu = backend.cpu.lock().unwrap();
624
625            cpu.usage_usec = 1000;
626            cpu.user_usec = 500;
627            cpu.system_usec = 500;
628            cpu.nr_periods = Some(2);
629        }
630
631        src.measure().await.unwrap();
632
633        {
634            let mut cpu = backend.cpu.lock().unwrap();
635
636            cpu.usage_usec = 2000;
637            cpu.user_usec = 1000;
638            cpu.system_usec = 1000;
639            cpu.nr_periods = Some(4);
640        }
641
642        src.measure().await.unwrap();
643
644        let counters = src.retrieve().await.unwrap();
645
646        assert_eq!(counters.proc_cpu.usage_usec.diff(), 1000);
647        assert_eq!(counters.proc_cpu.user_usec.diff(), 500);
648        assert_eq!(counters.proc_cpu.system_usec.diff(), 500);
649        assert_eq!(counters.proc_cpu.nr_periods.unwrap().diff(), 2);
650    }
651
652    #[tokio::test]
653    async fn test_worker_updates_counters() {
654        let (src, backend) = setup_source();
655
656        let (token, handle) = Cgroup::create_worker(
657            src.proc_cgroup.clone(),
658            src.root_cgroup.clone(),
659            src.proc_memory_counters.clone(),
660            src.global_memory_counters.clone(),
661            Duration::from_millis(10),
662        )
663        .unwrap();
664
665        {
666            let mem = src.proc_memory_counters.lock().unwrap();
667            assert!(mem.anon.is_none());
668        }
669
670        {
671            backend.memory.lock().unwrap().anon = Some(200);
672        }
673
674        sleep(Duration::from_millis(10)).await;
675
676        {
677            let mem = src.proc_memory_counters.lock().unwrap();
678            assert_eq!(mem.anon.unwrap().max().unwrap(), 200);
679            assert_eq!(mem.anon.unwrap().min().unwrap(), 200);
680        }
681
682        {
683            backend.memory.lock().unwrap().anon = Some(100);
684        }
685
686        sleep(Duration::from_millis(10)).await;
687
688        {
689            let mem = src.proc_memory_counters.lock().unwrap();
690            assert_eq!(mem.anon.unwrap().max().unwrap(), 200);
691            assert_eq!(mem.anon.unwrap().min().unwrap(), 100);
692        }
693
694        token.cancel();
695        handle.await.unwrap().unwrap();
696    }
697}
joule_profiler_source_cgroup/lib.rs

joule_profiler_source_cgroup/
lib.rs
