refactor: extract memory resource construction

Signed-off-by: bharqav <bhargavpodapati28@gmail.com>

Author: bharqav

crates/liboci-cli/src/update.rs

@@ -43,15 +43,15 @@ pub struct Update {
     pub cpuset_mems: Option<String>,
 
     /// Set memory limit to num bytes.
-    #[arg(long)]
-    pub memory: Option<u64>,
+    #[arg(long, allow_hyphen_values = true)]
+    pub memory: Option<i64>,
 
     /// Set memory reservation (or soft limit) to num bytes.
-    #[arg(long)]
-    pub memory_reservation: Option<u64>,
+    #[arg(long, allow_hyphen_values = true)]
+    pub memory_reservation: Option<i64>,
 
     /// Set total memory + swap usage to num bytes. Use -1 to unset the limit (i.e. use unlimited swap).
-    #[arg(long)]
+    #[arg(long, allow_hyphen_values = true)]
     pub memory_swap: Option<i64>,
 
     /// Set the maximum number of processes allowed in the container

crates/youki/src/commands/update.rs

@@ -4,7 +4,9 @@ use std::{fs, io};
 use anyhow::Result;
 use libcgroups::common::{CgroupManager, ControllerOpt};
 use libcgroups::{self};
-use libcontainer::oci_spec::runtime::{LinuxPidsBuilder, LinuxResources, LinuxResourcesBuilder};
+use libcontainer::oci_spec::runtime::{
+    LinuxMemory, LinuxMemoryBuilder, LinuxPidsBuilder, LinuxResources, LinuxResourcesBuilder,
+};
 use liboci_cli::Update;
 
 use crate::commands::create_cgroup_manager;
@@ -26,6 +28,11 @@ pub fn update(args: Update, root_path: PathBuf) -> Result<()> {
         if let Some(new_pids_limit) = args.pids_limit {
             builder = builder.pids(LinuxPidsBuilder::default().limit(new_pids_limit).build()?);
         }
+
+        if let Some(memory) = build_memory(&args)? {
+            builder = builder.memory(memory);
+        }
+
         linux_res = builder.build()?;
     }
 
@@ -37,3 +44,101 @@ pub fn update(args: Update, root_path: PathBuf) -> Result<()> {
     })?;
     Ok(())
 }
+
+fn build_memory(args: &Update) -> Result<Option<LinuxMemory>> {
+    let mut mem_builder = LinuxMemoryBuilder::default();
+    let mut has_memory = false;
+    if let Some(memory) = args.memory {
+        mem_builder = mem_builder.limit(memory);
+        has_memory = true;
+    }
+    if let Some(reservation) = args.memory_reservation {
+        mem_builder = mem_builder.reservation(reservation);
+        has_memory = true;
+    }
+    if let Some(swap) = args.memory_swap {
+        mem_builder = mem_builder.swap(swap);
+        has_memory = true;
+    }
+
+    if has_memory {
+        Ok(Some(mem_builder.build()?))
+    } else {
+        Ok(None)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn default_update() -> Update {
+        Update {
+            resources: None,
+            blkio_weight: None,
+            cpu_period: None,
+            cpu_quota: None,
+            cpu_rt_period: None,
+            cpu_rt_runtime: None,
+            cpu_share: None,
+            cpuset_cpus: None,
+            cpuset_mems: None,
+            memory: None,
+            memory_reservation: None,
+            memory_swap: None,
+            pids_limit: None,
+            l3_cache_schema: None,
+            mem_bw_schema: None,
+            container_id: "test".to_string(),
+        }
+    }
+
+    #[test]
+    fn test_build_memory_none() {
+        let args = default_update();
+        let mem = build_memory(&args).unwrap();
+        assert!(mem.is_none());
+    }
+
+    #[test]
+    fn test_build_memory_limit() {
+        let mut args = default_update();
+        args.memory = Some(1024);
+        let mem = build_memory(&args).unwrap().unwrap();
+        assert_eq!(mem.limit(), Some(1024));
+        assert_eq!(mem.reservation(), None);
+        assert_eq!(mem.swap(), None);
+    }
+
+    #[test]
+    fn test_build_memory_reservation() {
+        let mut args = default_update();
+        args.memory_reservation = Some(512);
+        let mem = build_memory(&args).unwrap().unwrap();
+        assert_eq!(mem.limit(), None);
+        assert_eq!(mem.reservation(), Some(512));
+        assert_eq!(mem.swap(), None);
+    }
+
+    #[test]
+    fn test_build_memory_swap() {
+        let mut args = default_update();
+        args.memory_swap = Some(2048);
+        let mem = build_memory(&args).unwrap().unwrap();
+        assert_eq!(mem.limit(), None);
+        assert_eq!(mem.reservation(), None);
+        assert_eq!(mem.swap(), Some(2048));
+    }
+
+    #[test]
+    fn test_build_memory_all() {
+        let mut args = default_update();
+        args.memory = Some(1024);
+        args.memory_reservation = Some(512);
+        args.memory_swap = Some(2048);
+        let mem = build_memory(&args).unwrap().unwrap();
+        assert_eq!(mem.limit(), Some(1024));
+        assert_eq!(mem.reservation(), Some(512));
+        assert_eq!(mem.swap(), Some(2048));
+    }
+}