runner: move jobserver token handling into task process

We can avoid holding onto a token while waiting on a lock this way.
Abnormal exits that would make us lose tokens should not happen here; if
they do, the runner PID1 will stop the whole build, so a lost token does
not matter.

2 files changed, 8 insertions, 5 deletions

diff --git a/crates/runner/src/lib.rs b/crates/runner/src/lib.rs
index 01ab85d..a300446 100644
--- a/crates/runner/src/lib.rs
+++ b/crates/runner/src/lib.rs
@@ -62,7 +62,6 @@ fn handle_request(
 	request_socket: UnixStream,
 ) -> UnixSeqpacketConn {
 	let child = |socket| {
-		let token = ctx.jobserver.wait();
 		let pid = unsafe {
 			clone::spawn(None, (), |()| {
 				drop(socket);
@@ -80,9 +79,10 @@ fn handle_request(
 		}
 		.expect("fork()")
 		.0;
-		let wait_res = wait::waitpid(pid, None);
-		ctx.jobserver.post(token);
-		wait_res.expect("waitpid()");
+		wait::waitpid(pid, None)
+			.expect("waidpid()")
+			.check()
+			.unwrap();
 	};
 
 	let (pid, socket) = unsafe { clone::spawn(None, socket, child) }.expect("fork()");
diff --git a/crates/runner/src/task.rs b/crates/runner/src/task.rs
index 6c52be2..e27fd49 100644
--- a/crates/runner/src/task.rs
+++ b/crates/runner/src/task.rs
@@ -433,7 +433,6 @@ fn save_cached(input_hash: &InputHash, output: &TaskOutput) -> Result<()> {
 pub fn handle(task: Task, jobserver: &mut Jobserver) -> Result<TaskOutput> {
 	let input_hash = input_hash(&task);
 
-	// TODO: We should not hold a job token while waiting on the lock
 	let _lock = unix::lock(paths::task_lock_filename(&input_hash), true, true)
 		.context("Failed to get task lock")?;
 
@@ -441,12 +440,16 @@ pub fn handle(task: Task, jobserver: &mut Jobserver) -> Result<TaskOutput> {
 		return Ok(task_output);
 	}
 
+	let token = jobserver.wait();
+
 	let start_time = Instant::now();
 	println!("Starting task {} ({})", task.label, input_hash);
 
 	let task_ret = run_and_hash_task(&input_hash, &task, jobserver);
 	let cleanup_ret = cleanup_task(&input_hash);
 
+	jobserver.post(token);
+
 	let task_output = task_ret?;
 	cleanup_ret.context("Failed to clean up after task")?;