mirror of
https://github.com/MeiK2333/river.git
synced 2025-11-04 14:49:40 +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) {
|
fn drop(&mut self) {
|
||||||
// TODO: 将不安全的指针类型转换回内置类型,以便由 Rust 自动回收资源
|
// TODO: 将不安全的指针类型转换回内置类型,以便由 Rust 自动回收资源
|
||||||
// TODO: 优先级较低,因为目前只在子进程里进行这个操作,且操作后会很快 exec,操作系统会回收这些内存
|
// TODO: 优先级较低,因为目前只在子进程里进行这个操作,且操作后会很快 exec,操作系统会回收这些内存
|
||||||
debug!("Dropping!");
|
debug!("TODO");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@ -12,6 +12,8 @@ use tonic::{Request, Response, Status};
|
|||||||
|
|
||||||
mod config;
|
mod config;
|
||||||
mod error;
|
mod error;
|
||||||
|
|
||||||
|
mod cgroup;
|
||||||
mod exec_args;
|
mod exec_args;
|
||||||
mod process;
|
mod process;
|
||||||
mod seccomp;
|
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::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::exec_args::ExecArgs;
|
||||||
use crate::seccomp;
|
use crate::seccomp;
|
||||||
use libc;
|
use libc;
|
||||||
@ -14,6 +15,7 @@ use std::ptr;
|
|||||||
use std::sync::{mpsc, Arc, Mutex};
|
use std::sync::{mpsc, Arc, Mutex};
|
||||||
use std::task::{Context, Poll};
|
use std::task::{Context, Poll};
|
||||||
use std::thread;
|
use std::thread;
|
||||||
|
use std::time::SystemTime;
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
use std::println as debug;
|
use std::println as debug;
|
||||||
@ -73,10 +75,11 @@ struct Runner {
|
|||||||
tx: Arc<Mutex<mpsc::Sender<RunnerStatus>>>,
|
tx: Arc<Mutex<mpsc::Sender<RunnerStatus>>>,
|
||||||
rx: Arc<Mutex<mpsc::Receiver<RunnerStatus>>>,
|
rx: Arc<Mutex<mpsc::Receiver<RunnerStatus>>>,
|
||||||
stack: *mut libc::c_void,
|
stack: *mut libc::c_void,
|
||||||
|
cgroup_set: CGroupSet,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Runner {
|
impl Runner {
|
||||||
pub fn from(process: Process) -> Runner {
|
pub fn from(process: Process) -> Result<Runner> {
|
||||||
let (tx, rx) = mpsc::channel();
|
let (tx, rx) = mpsc::channel();
|
||||||
let stack = unsafe {
|
let stack = unsafe {
|
||||||
libc::mmap(
|
libc::mmap(
|
||||||
@ -88,13 +91,14 @@ impl Runner {
|
|||||||
0,
|
0,
|
||||||
)
|
)
|
||||||
};
|
};
|
||||||
Runner {
|
Ok(Runner {
|
||||||
process: process,
|
process: process,
|
||||||
pid: -1,
|
pid: -1,
|
||||||
tx: Arc::new(Mutex::new(tx)),
|
tx: Arc::new(Mutex::new(tx)),
|
||||||
rx: Arc::new(Mutex::new(rx)),
|
rx: Arc::new(Mutex::new(rx)),
|
||||||
stack: stack,
|
stack: stack,
|
||||||
}
|
cgroup_set: CGroupSet::new()?,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -119,9 +123,9 @@ impl Drop for Runner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Future 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);
|
let runner = Pin::into_inner(self);
|
||||||
// 如果 pid == -1,则说明子进程还没有运行
|
// 如果 pid == -1,则说明子进程还没有运行
|
||||||
if runner.pid == -1 {
|
if runner.pid == -1 {
|
||||||
@ -142,27 +146,75 @@ impl Future for Runner {
|
|||||||
};
|
};
|
||||||
debug!("pid = {}", pid);
|
debug!("pid = {}", pid);
|
||||||
runner.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,以防止主流程被阻塞
|
// 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait,以防止主流程被阻塞
|
||||||
thread::spawn(move || {
|
thread::spawn(move || {
|
||||||
Runner::waitpid(pid);
|
let status = Runner::waitpid(pid);
|
||||||
// 子流程运行结束后通知主流程
|
// 子流程运行结束后通知主流程
|
||||||
|
let status_tx = tx.lock().unwrap();
|
||||||
|
status_tx.send(status).unwrap();
|
||||||
debug!("waker");
|
debug!("waker");
|
||||||
waker.wake();
|
waker.wake();
|
||||||
});
|
});
|
||||||
return Poll::Pending;
|
return Poll::Pending;
|
||||||
} else {
|
} else {
|
||||||
return Poll::Ready(0);
|
let status = runner.rx.lock().unwrap().recv().unwrap();
|
||||||
|
return Poll::Ready(Ok(status));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Runner {
|
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 {
|
unsafe {
|
||||||
let mut status: i32 = 0;
|
if libc::wait4(pid, &mut status, 0, &mut rusage) < 0 {
|
||||||
// TODO: 获取运行状态等并传递给主流程
|
return judge_system_error(String::from("wait4 failure!"));
|
||||||
libc::wait4(pid, &mut status, 0, ptr::null_mut());
|
}
|
||||||
|
}
|
||||||
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@ -360,7 +458,7 @@ mod tests {
|
|||||||
65535,
|
65535,
|
||||||
pwd.path().to_path_buf(),
|
pwd.path().to_path_buf(),
|
||||||
);
|
);
|
||||||
let result = Runner::from(process).await;
|
let result = Runner::from(process).unwrap().await.unwrap();
|
||||||
assert_eq!(result, 0);
|
assert_eq!(result.exit_code, 0);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user