runner: maintain list of task PIDs in context

1 files changed, 18 insertions, 8 deletions

diff --git a/crates/runner/src/lib.rs b/crates/runner/src/lib.rs
index 80c775d..01ab85d 100644
--- a/crates/runner/src/lib.rs
+++ b/crates/runner/src/lib.rs
@@ -7,6 +7,7 @@ mod task;
 mod util;
 
 use std::{
+	collections::HashSet,
 	fs::File,
 	net,
 	os::unix::{net::UnixStream, prelude::*},
@@ -24,6 +25,7 @@ use nix::{
 use uds::UnixSeqpacketConn;
 
 use common::{error::*, types::*};
+use util::Checkable;
 
 use self::{
 	jobserver::Jobserver,
@@ -38,15 +40,19 @@ pub struct Options {
 #[derive(Debug)]
 struct RunnerContext {
 	jobserver: Jobserver,
+	tasks: HashSet<Pid>,
 }
 
-fn handle_sigchld() {
+fn handle_sigchld(ctx: &mut RunnerContext) -> Result<()> {
 	loop {
 		let status = match wait::waitpid(Pid::from_raw(-1), Some(wait::WaitPidFlag::WNOHANG)) {
-			Ok(wait::WaitStatus::StillAlive) | Err(Errno::ECHILD) => return,
+			Ok(wait::WaitStatus::StillAlive) | Err(Errno::ECHILD) => return Ok(()),
 			res => res.expect("waitpid()"),
 		};
-		let _pid = status.pid().unwrap();
+		let pid = status.pid().unwrap();
+		if ctx.tasks.remove(&pid) {
+			status.check()?;
+		}
 	}
 }
 
@@ -79,9 +85,10 @@ fn handle_request(
 		wait_res.expect("waitpid()");
 	};
 
-	unsafe { clone::spawn(None, socket, child) }
-		.expect("fork()")
-		.1
+	let (pid, socket) = unsafe { clone::spawn(None, socket, child) }.expect("fork()");
+	assert!(ctx.tasks.insert(pid));
+
+	socket
 }
 
 fn handle_socket(ctx: &mut RunnerContext, socket: UnixSeqpacketConn) -> Option<UnixSeqpacketConn> {
@@ -113,7 +120,10 @@ fn runner(uid: Uid, gid: Gid, mut socket: UnixSeqpacketConn, _lockfile: File, op
 		.jobs
 		.unwrap_or_else(|| unix::nproc().expect("Failed to get number of available CPUs"));
 	let jobserver = Jobserver::new(jobs).expect("Failed to initialize jobserver pipe");
-	let mut ctx = RunnerContext { jobserver };
+	let mut ctx = RunnerContext {
+		jobserver,
+		tasks: HashSet::new(),
+	};
 
 	let mut signals = signal::SigSet::empty();
 	signals.add(signal::Signal::SIGCHLD);
@@ -134,7 +144,7 @@ fn runner(uid: Uid, gid: Gid, mut socket: UnixSeqpacketConn, _lockfile: File, op
 			.expect("Unknown events in poll() return");
 		if events.contains(poll::PollFlags::POLLIN) {
 			let _signal = sfd.read_signal().expect("read_signal()").unwrap();
-			handle_sigchld();
+			handle_sigchld(&mut ctx).expect("Task process exited abnormally");
 		} else if events.intersects(!poll::PollFlags::POLLIN) {
 			panic!("Unexpected error status for signal file descriptor");
 		}