runner: use signal fd to wait for exiting child processes

Will be used to poll for new tasks and signals in the same process.

1 files changed, 21 insertions, 5 deletions

diff --git a/crates/runner/src/lib.rs b/crates/runner/src/lib.rs
index bf052ac..97d8815 100644
--- a/crates/runner/src/lib.rs
+++ b/crates/runner/src/lib.rs
@@ -15,9 +15,10 @@ use std::{
 
 use capctl::prctl;
 use nix::{
+	errno::Errno,
 	sched::CloneFlags,
-	sys::{signal, stat, wait},
-	unistd::{self, Gid, Uid},
+	sys::{signal, signalfd::SignalFd, stat, wait},
+	unistd::{self, Gid, Pid, Uid},
 };
 use uds::UnixSeqpacketConn;
 
@@ -94,6 +95,12 @@ fn runner(uid: Uid, gid: Gid, socket: UnixSeqpacketConn, _lockfile: File, option
 
 	init::init_runner().unwrap();
 
+	let mut signals = signal::SigSet::empty();
+	signals.add(signal::Signal::SIGCHLD);
+	signal::pthread_sigmask(signal::SigmaskHow::SIG_BLOCK, Some(&signals), None)
+		.expect("pthread_sigmask()");
+	let mut sfd = SignalFd::new(&signals).expect("Failed to create signal file descriptor");
+
 	let msg_handler = unsafe {
 		clone::spawn(None, (), |()| {
 			signal::signal(signal::Signal::SIGCHLD, signal::SigHandler::SigIgn).unwrap();
@@ -104,9 +111,18 @@ fn runner(uid: Uid, gid: Gid, socket: UnixSeqpacketConn, _lockfile: File, option
 	.0;
 
 	loop {
-		let status = wait::wait().expect("wait()");
-		if status.pid() == Some(msg_handler) {
-			break;
+		let _signal = sfd.read_signal().expect("read_signal()").unwrap();
+
+		loop {
+			let status = match wait::waitpid(Pid::from_raw(-1), Some(wait::WaitPidFlag::WNOHANG)) {
+				Ok(wait::WaitStatus::StillAlive) | Err(Errno::ECHILD) => break,
+				res => res.expect("waitpid()"),
+			};
+			let pid = status.pid().unwrap();
+
+			if pid == msg_handler {
+				return;
+			}
 		}
 	}
 }