1 files changed, 79 insertions, 0 deletions
diff --git a/crates/rebel-runner/src/jobserver.rs b/crates/rebel-runner/src/jobserver.rs
new file mode 100644
index 0000000..7c3f2f7
--- /dev/null
+++ b/crates/rebel-runner/src/jobserver.rs
@@ -0,0 +1,79 @@
+use std::{
+	os::fd::{AsFd, AsRawFd, OwnedFd},
+	slice,
+};
+
+use nix::{errno::Errno, fcntl::OFlag, poll, unistd};
+
+use rebel_common::error::*;
+
+use super::util::unix;
+
+#[derive(Debug)]
+pub struct Jobserver {
+	tokens: usize,
+	r: OwnedFd,
+	w: OwnedFd,
+}
+
+impl Jobserver {
+	pub fn new(tokens: usize) -> Result<Jobserver> {
+		let (r, w) = unistd::pipe2(OFlag::O_CLOEXEC | OFlag::O_NONBLOCK).context("pipe()")?;
+
+		for _ in 0..tokens {
+			if unistd::write(w.as_fd(), b"+").is_err() {
+				break;
+			}
+		}
+		unix::set_blocking(&w, true)?;
+
+		Ok(Jobserver { tokens, r, w })
+	}
+
+	pub fn wait(&mut self) -> u8 {
+		loop {
+			poll::poll(
+				&mut [poll::PollFd::new(self.r.as_fd(), poll::PollFlags::POLLIN)],
+				poll::PollTimeout::NONE,
+			)
+			.expect("poll()");
+
+			let mut token = 0;
+			match unistd::read(self.r.as_raw_fd(), slice::from_mut(&mut token)) {
+				Ok(n) => {
+					assert!(n == 1);
+					return token;
+				}
+				Err(Errno::EAGAIN) => {
+					// Token was sniped by another task
+					continue;
+				}
+				error @ Err(_) => {
+					error.expect("read()");
+				}
+			}
+		}
+	}
+
+	pub fn post(&mut self, token: u8) {
+		let n = unistd::write(self.w.as_fd(), slice::from_ref(&token)).expect("write()");
+		assert!(n == 1);
+	}
+
+	pub fn set_cloexec(&mut self, cloexec: bool) -> Result<()> {
+		unix::set_cloexec(&self.r, cloexec)?;
+		unix::set_cloexec(&self.w, cloexec)?;
+		Ok(())
+	}
+
+	pub fn to_makeflags(&self) -> String {
+		format!(
+			" -j{} --jobserver-auth={},{}",
+			self.tokens,
+			self.r.as_raw_fd(),
+			self.w.as_raw_fd(),
+		)
+	}
+}
+
+// FIXME Log lost tokens on drop