use crate::json::frame;
use crate::subagent::protocol::{AgentMsg, ControlMsg, SUBAGENT_ENV, SpawnPayload};
use crate::supervisor::kill::kill_group;
use crate::supervisor::tree::NodeId;
use std::io;
use std::path::Path;
use std::process::{Child, ChildStdin, Command, Stdio};
use std::thread::JoinHandle;
use std::time::Duration;
pub struct Subagent {
pub node: NodeId,
child: Child,
writer: ChildStdin,
pgid: i32,
reaped: bool,
reader: Option<JoinHandle<()>>,
_cgroup: Option<crate::supervisor::cgroup::CgroupGuard>,
}
pub type FrameSink = std::sync::Arc<dyn Fn(NodeId, AgentMsg) -> bool + Send + Sync>;
pub fn spawn(
exe: &Path,
payload: &SpawnPayload,
node: NodeId,
events: FrameSink,
) -> io::Result<Subagent> {
let mut cmd = Command::new(exe);
cmd.env(SUBAGENT_ENV, "1")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit());
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
let mem = payload.limits.memory_bytes;
let cpu = payload.limits.cpu_seconds;
let nice = payload.limits.nice;
unsafe {
cmd.pre_exec(move || {
libc::setpgid(0, 0);
if let Some(bytes) = mem {
let lim = libc::rlimit {
rlim_cur: bytes as libc::rlim_t,
rlim_max: bytes as libc::rlim_t,
};
if libc::setrlimit(libc::RLIMIT_AS, &lim) != 0 {
return Err(std::io::Error::last_os_error());
}
}
if let Some(secs) = cpu {
let lim = libc::rlimit {
rlim_cur: secs as libc::rlim_t,
rlim_max: secs.saturating_add(5) as libc::rlim_t,
};
if libc::setrlimit(libc::RLIMIT_CPU, &lim) != 0 {
return Err(std::io::Error::last_os_error());
}
}
if let Some(n) = nice {
let _ = libc::setpriority(libc::PRIO_PROCESS, 0, n);
}
Ok(())
});
}
}
let mut child = {
let mut attempt = 0u32;
loop {
match cmd.spawn() {
Ok(c) => break c,
Err(e)
if attempt < 10
&& (e.raw_os_error() == Some(libc::EAGAIN)
|| e.kind() == io::ErrorKind::WouldBlock) =>
{
attempt += 1;
std::thread::sleep(Duration::from_millis(u64::from(20 * attempt)));
}
Err(e) => return Err(e),
}
}
};
let pgid = child.id() as i32;
let cgroup = crate::supervisor::cgroup::CgroupGuard::for_run().inspect(|g| {
g.place(pgid);
});
let mut writer = child
.stdin
.take()
.ok_or_else(|| io::Error::other("no child stdin"))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| io::Error::other("no child stdout"))?;
frame::write_frame(&mut writer, &ControlMsg::Spawn(Box::new(payload.clone())))?;
let reader = std::thread::Builder::new()
.name(format!("subagent-events:{}", node.0))
.spawn(move || {
let mut r = io::BufReader::new(stdout);
while let Ok(Some(bytes)) = frame::read_frame(&mut r) {
match serde_json::from_slice::<AgentMsg>(&bytes) {
Ok(msg) => {
if !events(node, msg) {
break; }
}
Err(_) => { }
}
}
})?;
Ok(Subagent {
node,
child,
writer,
pgid,
reaped: false,
reader: Some(reader),
_cgroup: cgroup,
})
}
impl Subagent {
pub fn pid(&self) -> i32 {
self.child.id() as i32
}
pub fn join_reader(&mut self) {
if let Some(h) = self.reader.take() {
let _ = h.join();
}
}
pub fn pgid(&self) -> i32 {
self.pgid
}
pub fn send(&mut self, msg: &ControlMsg) -> io::Result<()> {
frame::write_frame(&mut self.writer, msg)
}
pub fn mark_reaped(&mut self) {
self.reaped = true;
}
pub fn kill(&mut self) {
if !self.reaped {
crate::supervisor::reaper::deregister(self.pid());
kill_group(self.pgid);
let _ = self.child.kill();
let _ = self.child.wait();
self.reaped = true;
}
}
}
impl Drop for Subagent {
fn drop(&mut self) {
if !self.reaped {
crate::supervisor::reaper::deregister(self.pid());
kill_group(self.pgid);
let _ = self.child.kill();
let _ = self.child.wait();
}
}
}