添加 cgroups 的 memory 限制与测量

This commit is contained in:
MeiK2333 2021-01-08 21:49:48 +08:00
parent 16047cf0f2
commit 6e6b50a745
3 changed files with 21 additions and 15 deletions

View File

@ -3,7 +3,3 @@
## 环境要求 ## 环境要求
- linux - linux
## TODO
- cgroup

View File

@ -1,6 +1,6 @@
use crate::error::{Error, Result}; use crate::error::{Error, Result};
use std::fs::{read_to_string, remove_dir, OpenOptions}; use std::fs;
use std::io::Write; use std::fs::{read_to_string, remove_dir};
use std::path::PathBuf; use std::path::PathBuf;
use tempfile::tempdir_in; use tempfile::tempdir_in;
@ -23,17 +23,13 @@ impl CGroupOptions {
/// 将指定进程加入该 cgroup 组 /// 将指定进程加入该 cgroup 组
pub fn apply(&self, pid: i32) -> Result<()> { pub fn apply(&self, pid: i32) -> Result<()> {
debug!("add {} to cgroup {:?}", pid, self.path); debug!("add {} to cgroup {:?}", pid, self.path);
self.set("tasks", &format!("{}", pid)) self.set("cgroup.procs", &format!("{}", pid))
} }
/// e.g `set("memory.limit_in_bytes", "67108864")` /// e.g `set("memory.limit_in_bytes", "67108864")`
pub fn set(&self, key: &str, value: &str) -> Result<()> { pub fn set(&self, key: &str, value: &str) -> Result<()> {
debug!("set {} values {}", key, value); debug!("set {} values {}", key, value);
let mut file = try_io!(OpenOptions::new() try_io!(fs::write(self.path.join(key), value));
.write(true)
.append(true)
.open(self.path.join(key)));
let _ = writeln!(file, "{}", value);
Ok(()) Ok(())
} }

View File

@ -122,6 +122,7 @@ impl Drop for Runner {
} }
} }
// TODO: unwrap -> Result<Error>
impl Future for Runner { impl Future for Runner {
type Output = Result<RunnerStatus>; type Output = Result<RunnerStatus>;
@ -146,15 +147,17 @@ impl Future for Runner {
}; };
debug!("pid = {}", pid); debug!("pid = {}", pid);
runner.pid = pid; runner.pid = pid;
// 设置 cgroup 限制
// 此处为父进程做的策略,所以没有与子进程的安全策略放一块
runner.cgroup_set.apply(pid).unwrap();
runner runner
.cgroup_set .cgroup_set
.memory .memory
.set( .set(
"memory.limit_in_bytes", "memory.limit_in_bytes",
&format!("{}", runner.process.memory_limit * 1024), &format!("{}", runner.process.memory_limit as i64 * 1024),
) )
.unwrap(); .unwrap();
runner.cgroup_set.apply(pid).unwrap();
let tx = runner.tx.clone(); let tx = runner.tx.clone();
// 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait以防止主流程被阻塞 // 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait以防止主流程被阻塞
@ -169,6 +172,17 @@ impl Future for Runner {
return Poll::Pending; return Poll::Pending;
} else { } else {
let status = runner.rx.lock().unwrap().recv().unwrap(); let status = runner.rx.lock().unwrap().recv().unwrap();
// TODO: 将 CGroup 的结果合并到 status 中
let mem_used = runner
.cgroup_set
.memory
.get("memory.max_usage_in_bytes")
.unwrap()
.trim()
.parse::<i32>()
.unwrap();
debug!("cgroup memory used: {} KiB", mem_used / 1024);
debug!("rusage memory used: {} KiB", status.memory_used);
return Poll::Ready(Ok(status)); return Poll::Ready(Ok(status));
} }
} }
@ -225,7 +239,7 @@ extern "C" fn runit(process: *mut libc::c_void) -> i32 {
let exec_args = ExecArgs::build(&process.cmd).unwrap(); let exec_args = ExecArgs::build(&process.cmd).unwrap();
unsafe { unsafe {
security(&process); security(&process);
fd_dup(); // fd_dup();
syscall_or_panic!(libc::execve( syscall_or_panic!(libc::execve(
exec_args.pathname, exec_args.pathname,
exec_args.argv, exec_args.argv,