mirror of
https://github.com/MeiK2333/river.git
synced 2025-09-26 22:49:11 +08:00
添加基础的 cgroup 操作
This commit is contained in:
parent
b68bfc911b
commit
16047cf0f2
71
src/cgroup.rs
Normal file
71
src/cgroup.rs
Normal 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(())
|
||||
}
|
||||
}
|
@ -47,6 +47,6 @@ impl Drop for ExecArgs {
|
||||
fn drop(&mut self) {
|
||||
// TODO: 将不安全的指针类型转换回内置类型,以便由 Rust 自动回收资源
|
||||
// TODO: 优先级较低,因为目前只在子进程里进行这个操作,且操作后会很快 exec,操作系统会回收这些内存
|
||||
debug!("Dropping!");
|
||||
debug!("TODO");
|
||||
}
|
||||
}
|
||||
|
@ -12,6 +12,8 @@ use tonic::{Request, Response, Status};
|
||||
|
||||
mod config;
|
||||
mod error;
|
||||
|
||||
mod cgroup;
|
||||
mod exec_args;
|
||||
mod process;
|
||||
mod seccomp;
|
||||
|
126
src/process.rs
126
src/process.rs
@ -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);
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user