添加基础的 cgroup 操作

This commit is contained in:
MeiK2333 2021-01-07 19:40:47 +08:00
parent b68bfc911b
commit 16047cf0f2
4 changed files with 186 additions and 15 deletions

71
src/cgroup.rs Normal file
View File

@ -0,0 +1,71 @@
use crate::error::{Error, Result};
use std::fs::{read_to_string, remove_dir, OpenOptions};
use std::io::Write;
use std::path::PathBuf;
use tempfile::tempdir_in;
#[cfg(test)]
use std::println as debug;
pub struct CGroupOptions {
path: PathBuf,
}
impl CGroupOptions {
pub fn new(base_path: &str) -> Result<CGroupOptions> {
let pwd = try_io!(tempdir_in(base_path));
debug!("register cgroup {:?}", pwd);
Ok(CGroupOptions {
path: pwd.path().to_path_buf(),
})
}
/// 将指定进程加入该 cgroup 组
pub fn apply(&self, pid: i32) -> Result<()> {
debug!("add {} to cgroup {:?}", pid, self.path);
self.set("tasks", &format!("{}", pid))
}
/// e.g `set("memory.limit_in_bytes", "67108864")`
pub fn set(&self, key: &str, value: &str) -> Result<()> {
debug!("set {} values {}", key, value);
let mut file = try_io!(OpenOptions::new()
.write(true)
.append(true)
.open(self.path.join(key)));
let _ = writeln!(file, "{}", value);
Ok(())
}
/// e.g `get("memory.max_usage_in_bytes")`
pub fn get(&self, key: &str) -> Result<String> {
Ok(try_io!(read_to_string(self.path.join(key))))
}
}
impl Drop for CGroupOptions {
fn drop(&mut self) {
debug!("remove cgroup {:?}", self.path);
remove_dir(&self.path).unwrap();
}
}
/// cgroup v1
pub struct CGroupSet {
// pub cpuset: CGroupOptions,
pub memory: CGroupOptions,
}
impl CGroupSet {
pub fn new() -> Result<CGroupSet> {
Ok(CGroupSet {
// cpuset: CGroupOptions::new("/sys/fs/cgroup/cpuset")?,
memory: CGroupOptions::new("/sys/fs/cgroup/memory")?,
})
}
pub fn apply(&self, pid: i32) -> Result<()> {
// self.cpuset.apply(pid)?;
self.memory.apply(pid)?;
Ok(())
}
}

View File

@ -47,6 +47,6 @@ impl Drop for ExecArgs {
fn drop(&mut self) {
// TODO: 将不安全的指针类型转换回内置类型,以便由 Rust 自动回收资源
// TODO: 优先级较低,因为目前只在子进程里进行这个操作,且操作后会很快 exec操作系统会回收这些内存
debug!("Dropping!");
debug!("TODO");
}
}

View File

@ -12,6 +12,8 @@ use tonic::{Request, Response, Status};
mod config;
mod error;
mod cgroup;
mod exec_args;
mod process;
mod seccomp;

View File

@ -1,5 +1,6 @@
use crate::cgroup::CGroupSet;
use crate::config::{STDERR_FILENAME, STDIN_FILENAME, STDOUT_FILENAME};
use crate::error::errno_str;
use crate::error::{errno_str, Result};
use crate::exec_args::ExecArgs;
use crate::seccomp;
use libc;
@ -14,6 +15,7 @@ use std::ptr;
use std::sync::{mpsc, Arc, Mutex};
use std::task::{Context, Poll};
use std::thread;
use std::time::SystemTime;
#[cfg(test)]
use std::println as debug;
@ -73,10 +75,11 @@ struct Runner {
tx: Arc<Mutex<mpsc::Sender<RunnerStatus>>>,
rx: Arc<Mutex<mpsc::Receiver<RunnerStatus>>>,
stack: *mut libc::c_void,
cgroup_set: CGroupSet,
}
impl Runner {
pub fn from(process: Process) -> Runner {
pub fn from(process: Process) -> Result<Runner> {
let (tx, rx) = mpsc::channel();
let stack = unsafe {
libc::mmap(
@ -88,13 +91,14 @@ impl Runner {
0,
)
};
Runner {
Ok(Runner {
process: process,
pid: -1,
tx: Arc::new(Mutex::new(tx)),
rx: Arc::new(Mutex::new(rx)),
stack: stack,
}
cgroup_set: CGroupSet::new()?,
})
}
}
@ -119,9 +123,9 @@ impl Drop for Runner {
}
impl Future for Runner {
type Output = i32;
type Output = Result<RunnerStatus>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<i32> {
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<RunnerStatus>> {
let runner = Pin::into_inner(self);
// 如果 pid == -1则说明子进程还没有运行
if runner.pid == -1 {
@ -142,27 +146,75 @@ impl Future for Runner {
};
debug!("pid = {}", pid);
runner.pid = pid;
runner
.cgroup_set
.memory
.set(
"memory.limit_in_bytes",
&format!("{}", runner.process.memory_limit * 1024),
)
.unwrap();
runner.cgroup_set.apply(pid).unwrap();
let tx = runner.tx.clone();
// 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait以防止主流程被阻塞
thread::spawn(move || {
Runner::waitpid(pid);
let status = Runner::waitpid(pid);
// 子流程运行结束后通知主流程
let status_tx = tx.lock().unwrap();
status_tx.send(status).unwrap();
debug!("waker");
waker.wake();
});
return Poll::Pending;
} else {
return Poll::Ready(0);
let status = runner.rx.lock().unwrap().recv().unwrap();
return Poll::Ready(Ok(status));
}
}
}
impl Runner {
fn waitpid(pid: i32) {
fn waitpid(pid: i32) -> RunnerStatus {
let mut status: i32 = 0;
let mut rusage = new_rusage();
let start = SystemTime::now();
unsafe {
let mut status: i32 = 0;
// TODO: 获取运行状态等并传递给主流程
libc::wait4(pid, &mut status, 0, ptr::null_mut());
if libc::wait4(pid, &mut status, 0, &mut rusage) < 0 {
return judge_system_error(String::from("wait4 failure!"));
}
}
let real_time_used = match start.elapsed() {
Ok(elapsed) => elapsed.as_millis(),
Err(_) => return judge_system_error(String::from("real time elapsed failure!")),
};
let mut exit_code = 0;
let exited = libc::WIFEXITED(status);
if exited {
exit_code = libc::WEXITSTATUS(status);
}
let signal = if libc::WIFSIGNALED(status) {
libc::WTERMSIG(status)
} else if libc::WIFSTOPPED(status) {
libc::WSTOPSIG(status)
} else {
0
};
let time_used = rusage.ru_utime.tv_sec * 1000
+ i64::from(rusage.ru_utime.tv_usec) / 1000
+ rusage.ru_stime.tv_sec * 1000
+ i64::from(rusage.ru_stime.tv_usec) / 1000;
let memory_used = rusage.ru_maxrss;
RunnerStatus {
errmsg: String::from(""),
memory_used: memory_used,
time_used: time_used,
real_time_used: real_time_used,
exit_code: exit_code,
signal: signal,
status: status,
rusage: rusage,
}
}
}
@ -346,6 +398,52 @@ fn deny_syscall(syscall_number: i64) -> seccomp::SyscallRuleSet {
)
}
/// 一个全为 `0` 的 `rusage`
#[inline(always)]
fn new_rusage() -> libc::rusage {
libc::rusage {
ru_utime: libc::timeval {
tv_sec: 0 as libc::time_t,
tv_usec: 0 as libc::suseconds_t,
},
ru_stime: libc::timeval {
tv_sec: 0 as libc::time_t,
tv_usec: 0 as libc::suseconds_t,
},
ru_maxrss: 0 as libc::c_long,
ru_ixrss: 0 as libc::c_long,
ru_idrss: 0 as libc::c_long,
ru_isrss: 0 as libc::c_long,
ru_minflt: 0 as libc::c_long,
ru_majflt: 0 as libc::c_long,
ru_nswap: 0 as libc::c_long,
ru_inblock: 0 as libc::c_long,
ru_oublock: 0 as libc::c_long,
ru_msgsnd: 0 as libc::c_long,
ru_msgrcv: 0 as libc::c_long,
ru_nsignals: 0 as libc::c_long,
ru_nvcsw: 0 as libc::c_long,
ru_nivcsw: 0 as libc::c_long,
}
}
/// 由于评测系统本身异常而产生的错误
///
/// 因为正常程序返回 `signal` 不能为负数,因此此处使用负数的 `signal` 标识系统错误
#[inline(always)]
fn judge_system_error(errmsg: String) -> RunnerStatus {
RunnerStatus {
rusage: new_rusage(),
exit_code: -1,
status: -1,
signal: -1,
time_used: -1,
memory_used: -1,
real_time_used: 0,
errmsg: errmsg,
}
}
#[cfg(test)]
mod tests {
use super::*;
@ -360,7 +458,7 @@ mod tests {
65535,
pwd.path().to_path_buf(),
);
let result = Runner::from(process).await;
assert_eq!(result, 0);
let result = Runner::from(process).unwrap().await.unwrap();
assert_eq!(result.exit_code, 0);
}
}