add memory cgroup controller

.gitignore

@@ -1,2 +1,6 @@
 /target
 .vagrant/
+
+tags
+tags.lock
+tags.temp

Cargo.toml

@@ -19,4 +19,4 @@ mio = { version = "0.7", features = ["os-ext", "os-poll"] }
 chrono = "0.4"
 once_cell = "1.6.0"
 futures = { version = "0.3", features = ["thread-pool"] }
-regex = "1.5"
+regex = "1.5"

src/cgroups/controller_type.rs

@@ -4,6 +4,7 @@ pub enum ControllerType {
     Devices,
     HugeTlb,
     Pids,
+    Memory,
 }
 
 impl ToString for ControllerType {
@@ -12,6 +13,7 @@ impl ToString for ControllerType {
             Self::Devices => "devices".into(),
             Self::HugeTlb => "hugetlb".into(),
             Self::Pids => "pids".into(),
+            Self::Memory => "memory".into(),
         }
     }
 }

src/cgroups/memory.rs

@@ -0,0 +1,306 @@
+use std::io::{prelude::*, Write};
+use std::{
+    fs::{create_dir_all, OpenOptions},
+    path::Path,
+};
+
+use anyhow::{Result, *};
+use nix::{errno::Errno, unistd::Pid};
+
+use crate::{
+    cgroups::Controller,
+    spec::{LinuxMemory, LinuxResources},
+};
+
+const CGROUP_MEMORY_SWAP_LIMIT: &str = "memory.memsw.limit_in_bytes";
+const CGROUP_MEMORY_LIMIT: &str = "memory.limit_in_bytes";
+const CGROUP_MEMORY_USAGE: &str = "memory.usage_in_bytes";
+const CGROUP_MEMORY_MAX_USAGE: &str = "memory.max_usage_in_bytes";
+const CGROUP_MEMORY_SWAPPINESS: &str = "memory.swappiness";
+const CGROUP_MEMORY_RESERVATION: &str = "memory.soft_limit_in_bytes";
+
+const CGROUP_KERNEL_MEMORY_LIMIT: &str = "memory.kmem.limit_in_bytes";
+const CGROUP_KERNEL_TCP_MEMORY_LIMIT: &str = "memory.kmem.tcp.limit_in_bytes";
+
+pub struct Memory {}
+
+impl Controller for Memory {
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path, pid: Pid) -> Result<()> {
+        create_dir_all(&cgroup_root)?;
+        let memory = linux_resources.memory.as_ref().unwrap();
+        let reservation = memory.reservation.unwrap_or(0);
+
+        Self::set_memory_and_swap(&memory, cgroup_root)?;
+
+        if reservation != 0 {
+            Self::set_i64(reservation, &cgroup_root.join(CGROUP_MEMORY_RESERVATION))?;
+        }
+
+        if let Some(swappiness) = memory.swappiness {
+            if swappiness <= 100 {
+                Self::set_u64(swappiness, &cgroup_root.join(CGROUP_MEMORY_SWAPPINESS))?;
+            } else {
+                // invalid swappiness value
+                return Err(anyhow!(
+                    "Invalid swappiness value: {}. Valid range is 0-100",
+                    swappiness
+                ));
+            }
+        }
+
+        // NOTE: Seems as though kernel and kernelTCP are both deprecated
+        // neither are implemented by runc. Tests pass without this, but
+        // kept in per the spec.
+        if let Some(kmem) = memory.kernel {
+            Self::set_i64(kmem, &cgroup_root.join(CGROUP_KERNEL_MEMORY_LIMIT))?;
+        }
+        if let Some(tcp_mem) = memory.kernel_tcp {
+            Self::set_i64(tcp_mem, &cgroup_root.join(CGROUP_KERNEL_TCP_MEMORY_LIMIT))?;
+        }
+
+        OpenOptions::new()
+            .create(false)
+            .write(true)
+            .truncate(false)
+            .open(cgroup_root.join("cgroup.procs"))?
+            .write_all(pid.to_string().as_bytes())?;
+
+        Ok(())
+    }
+}
+
+impl Memory {
+    fn get_memory_usage(cgroup_root: &Path) -> Result<u64> {
+        let path = cgroup_root.join(CGROUP_MEMORY_USAGE);
+        let mut contents = String::new();
+        OpenOptions::new()
+            .create(false)
+            .read(true)
+            .open(path)?
+            .read_to_string(&mut contents)?;
+
+        let val = contents.parse::<u64>()?;
+        Ok(val)
+    }
+
+    fn get_memory_max_usage(cgroup_root: &Path) -> Result<u64> {
+        let path = cgroup_root.join(CGROUP_MEMORY_MAX_USAGE);
+        let mut contents = String::new();
+        OpenOptions::new()
+            .create(false)
+            .read(true)
+            .open(path)?
+            .read_to_string(&mut contents)?;
+
+        let val = contents.parse::<u64>()?;
+        Ok(val)
+    }
+
+    fn get_memory_limit(cgroup_root: &Path) -> Result<i64> {
+        let path = cgroup_root.join(CGROUP_MEMORY_LIMIT);
+        let mut contents = String::new();
+        OpenOptions::new()
+            .create(false)
+            .read(true)
+            .open(path)?
+            .read_to_string(&mut contents)?;
+
+        let val = contents.parse::<i64>()?;
+        Ok(val)
+    }
+
+    fn set_i64(val: i64, path: &Path) -> std::io::Result<()> {
+        OpenOptions::new()
+            .create(false)
+            .write(true)
+            .truncate(true)
+            .open(path)?
+            .write_all(val.to_string().as_bytes())?;
+        Ok(())
+    }
+
+    fn set_u64(val: u64, path: &Path) -> std::io::Result<()> {
+        OpenOptions::new()
+            .create(false)
+            .write(true)
+            .truncate(true)
+            .open(path)?
+            .write_all(val.to_string().as_bytes())?;
+        Ok(())
+    }
+
+    fn set_memory(val: i64, cgroup_root: &Path) -> Result<()> {
+        let path = cgroup_root.join(CGROUP_MEMORY_LIMIT);
+
+        match Self::set_i64(val, &path) {
+            Ok(_) => Ok(()),
+            Err(e) => {
+                // we need to look into the raw OS error for an EBUSY status
+                if let Some(code) = e.raw_os_error() {
+                    // the nix crate has a handy enum for these, lets use that
+                    let errno = Errno::from_i32(code);
+                    // if the error is EBUSY
+                    if let Errno::EBUSY = errno {
+                        let usage = Self::get_memory_usage(cgroup_root)?;
+                        let max_usage = Self::get_memory_max_usage(cgroup_root)?;
+                        Err(anyhow!(
+                            "unable to set memory limit to {} (current usage: {}, peak usage: {})",
+                            val,
+                            usage,
+                            max_usage,
+                        ))
+                    } else {
+                        Err(anyhow!(e))
+                    }
+                } else {
+                    Err(anyhow!(e))
+                }
+            }
+        }
+    }
+
+    fn set_swap(val: i64, cgroup_root: &Path) -> Result<()> {
+        if val == 0 {
+            return Ok(());
+        }
+
+        let path = cgroup_root.join(CGROUP_MEMORY_SWAP_LIMIT);
+
+        Self::set_i64(val, &path)?;
+
+        Ok(())
+    }
+
+    fn set_memory_and_swap(resource: &LinuxMemory, cgroup_root: &Path) -> Result<()> {
+        let limit = resource.limit.unwrap_or(0);
+        let mut swap = resource.swap.unwrap_or(0);
+
+        if limit == -1 && swap == 0 {
+            if cgroup_root.join(CGROUP_MEMORY_SWAP_LIMIT).exists() {
+                swap = -1;
+            }
+        }
+
+        // According to runc we need to change the write sequence of
+        // limit and swap so it won't fail, because the new and old
+        // values don't fit in the kernel's validation
+        // see:
+        // https://github.com/opencontainers/runc/blob/master/libcontainer/cgroups/fs/memory.go#L89
+        if limit != 0 && swap != 0 {
+            let current_limit = Self::get_memory_limit(cgroup_root)?;
+
+            if swap == -1 || current_limit < swap {
+                Self::set_swap(swap, cgroup_root)?;
+                Self::set_memory(limit, cgroup_root)?;
+                return Ok(());
+            }
+        }
+
+        Self::set_memory(limit, cgroup_root)?;
+        Self::set_swap(swap, cgroup_root)?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::spec::LinuxMemory;
+
+    fn set_fixture(temp_dir: &std::path::Path, filename: &str, val: &str) -> Result<()> {
+        std::fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(true)
+            .open(temp_dir.join(filename))?
+            .write_all(val.as_bytes())?;
+
+        Ok(())
+    }
+
+    fn create_temp_dir(test_name: &str) -> Result<std::path::PathBuf> {
+        std::fs::create_dir_all(std::env::temp_dir().join(test_name))?;
+        Ok(std::env::temp_dir().join(test_name))
+    }
+
+    #[test]
+    fn test_set_memory() {
+        let limit = 1024;
+        let tmp = create_temp_dir("test_set_memory").expect("create temp directory for test");
+        set_fixture(&tmp, CGROUP_MEMORY_USAGE, "0").expect("Set fixure for memory usage");
+        set_fixture(&tmp, CGROUP_MEMORY_MAX_USAGE, "0").expect("Set fixure for max memory usage");
+        set_fixture(&tmp, CGROUP_MEMORY_LIMIT, "0").expect("Set fixure for memory limit");
+        Memory::set_memory(limit, &tmp).expect("Set memory limit");
+        let content =
+            std::fs::read_to_string(tmp.join(CGROUP_MEMORY_LIMIT)).expect("Read to string");
+        assert_eq!(limit.to_string(), content)
+    }
+
+    #[test]
+    fn test_set_swap() {
+        let limit = 512;
+        let tmp = create_temp_dir("test_set_swap").expect("create temp directory for test");
+        set_fixture(&tmp, CGROUP_MEMORY_SWAP_LIMIT, "0").expect("Set fixure for swap limit");
+        Memory::set_swap(limit, &tmp).expect("Set swap limit");
+        let content =
+            std::fs::read_to_string(tmp.join(CGROUP_MEMORY_SWAP_LIMIT)).expect("Read to string");
+        assert_eq!(limit.to_string(), content)
+    }
+
+    #[test]
+    fn test_set_memory_and_swap() {
+        let tmp =
+            create_temp_dir("test_set_memory_and_swap").expect("create temp directory for test");
+        set_fixture(&tmp, CGROUP_MEMORY_USAGE, "0").expect("Set fixure for memory usage");
+        set_fixture(&tmp, CGROUP_MEMORY_MAX_USAGE, "0").expect("Set fixure for max memory usage");
+        set_fixture(&tmp, CGROUP_MEMORY_LIMIT, "0").expect("Set fixure for memory limit");
+        set_fixture(&tmp, CGROUP_MEMORY_SWAP_LIMIT, "0").expect("Set fixure for swap limit");
+
+        // test unlimited memory with no set swap
+        {
+            let limit = -1;
+            let linux_memory = &LinuxMemory {
+                limit: Some(limit),
+                swap: None, // Some(0) gives the same outcome
+                reservation: None,
+                kernel: None,
+                kernel_tcp: None,
+                swappiness: None,
+            };
+            Memory::set_memory_and_swap(linux_memory, &tmp).expect("Set memory and swap");
+
+            let limit_content =
+                std::fs::read_to_string(tmp.join(CGROUP_MEMORY_LIMIT)).expect("Read to string");
+            assert_eq!(limit.to_string(), limit_content);
+
+            let swap_content = std::fs::read_to_string(tmp.join(CGROUP_MEMORY_SWAP_LIMIT))
+                .expect("Read to string");
+            // swap should be set to -1 also
+            assert_eq!(limit.to_string(), swap_content);
+        }
+
+        // test setting swap and memory to arbitrary values
+        {
+            let limit = 1024 * 1024 * 1024;
+            let swap = 1024;
+            let linux_memory = &LinuxMemory {
+                limit: Some(limit),
+                swap: Some(swap),
+                reservation: None,
+                kernel: None,
+                kernel_tcp: None,
+                swappiness: None,
+            };
+            Memory::set_memory_and_swap(linux_memory, &tmp).expect("Set memory and swap");
+
+            let limit_content =
+                std::fs::read_to_string(tmp.join(CGROUP_MEMORY_LIMIT)).expect("Read to string");
+            assert_eq!(limit.to_string(), limit_content);
+
+            let swap_content = std::fs::read_to_string(tmp.join(CGROUP_MEMORY_SWAP_LIMIT))
+                .expect("Read to string");
+            assert_eq!(swap.to_string(), swap_content);
+        }
+    }
+}

src/cgroups/mod.rs

@@ -3,6 +3,7 @@ mod controller_type;
 mod devices;
 mod hugetlb;
 mod manager;
+mod memory;
 mod pids;
 pub use controller::Controller;
 pub use controller_type::ControllerType;