feat(cgroup): add memory extension methods

Author: akagiyuu

src/metrics.rs

@@ -13,10 +13,11 @@ pub enum Verdict {
     IdleTimeLimitExceeded,
 }
 
+#[derive(Debug)]
 pub struct Metrics {
     pub verdict: Verdict,
     pub run_time: Duration,
+    pub memory_usage: Byte,
     pub stdout: String,
     pub stderr: String,
-    pub memory: Byte,
 }

src/sandbox/cgroup.rs

@@ -1,11 +1,14 @@
 use std::time::Duration;
 
-use cgroups_rs::fs::{Cgroup, cpu::CpuController};
+use byte_unit::Byte;
+use cgroups_rs::fs::{Cgroup, cpu::CpuController, memory::MemController};
 
 const CPU_USAGE_PREFIX: &str = "usage_usec ";
 
 pub trait CgroupExt {
     fn get_cpu_time(&self) -> Duration;
+    fn get_memory_usage(&self) -> Byte;
+    fn get_memory_limit(&self) -> Byte;
 }
 
 impl CgroupExt for Cgroup {
@@ -22,4 +25,20 @@ impl CgroupExt for Cgroup {
         let usage = usage.parse().unwrap();
         Duration::from_micros(usage)
     }
+
+    fn get_memory_usage(&self) -> Byte {
+        // SAFETY: there must be memory controller for cgroup v2
+        let memory_controller: &MemController = self.controller_of().unwrap();
+        let stats = memory_controller.memory_stat();
+
+        Byte::from_u64(stats.usage_in_bytes)
+    }
+
+    fn get_memory_limit(&self) -> Byte {
+        // SAFETY: there must be memory controller for cgroup v2
+        let memory_controller: &MemController = self.controller_of().unwrap();
+        let stats = memory_controller.memory_stat();
+
+        Byte::from_u64(stats.limit_in_bytes.max(0) as u64)
+    }
 }