Allow to set the maximum number of jobs on the command line

1 files changed, 20 insertions, 7 deletions

diff --git a/crates/runner/src/lib.rs b/crates/runner/src/lib.rs
index 45366f3..af1e2a7 100644
--- a/crates/runner/src/lib.rs
+++ b/crates/runner/src/lib.rs
@@ -26,6 +26,11 @@ use self::{
 	util::{clone, unix},
 };
 
+#[derive(Debug, Clone)]
+pub struct Options {
+	pub jobs: Option<usize>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 enum Message {
 	Request(Task, ipc::IpcSender<Result<TaskOutput>>),
@@ -61,9 +66,11 @@ fn handle_request(
 		.1
 }
 
-fn runner_loop(mut channel: ipc::IpcReceiver<Message>) {
-	let nproc = unix::nproc().expect("Failed to get number of available CPUs");
-	let mut jobserver = Jobserver::new(nproc).expect("Failed to initialize jobserver pipe");
+fn runner_loop(mut channel: ipc::IpcReceiver<Message>, options: &Options) {
+	let jobs = options
+		.jobs
+		.unwrap_or_else(|| unix::nproc().expect("Failed to get number of available CPUs"));
+	let mut jobserver = Jobserver::new(jobs).expect("Failed to initialize jobserver pipe");
 
 	while let Ok(msg) = channel.recv() {
 		match msg {
@@ -76,7 +83,13 @@ fn runner_loop(mut channel: ipc::IpcReceiver<Message>) {
 	}
 }
 
-fn runner(uid: Uid, gid: Gid, channel: ipc::IpcReceiver<Message>, _lockfile: File) {
+fn runner(
+	uid: Uid,
+	gid: Gid,
+	channel: ipc::IpcReceiver<Message>,
+	_lockfile: File,
+	options: &Options,
+) {
 	ns::mount_proc();
 	ns::setup_userns(Uid::from_raw(0), Gid::from_raw(0), uid, gid);
 
@@ -87,7 +100,7 @@ fn runner(uid: Uid, gid: Gid, channel: ipc::IpcReceiver<Message>, _lockfile: Fil
 	let msg_handler = unsafe {
 		clone::spawn(None, (), |()| {
 			signal::signal(signal::Signal::SIGCHLD, signal::SigHandler::SigIgn).unwrap();
-			runner_loop(channel);
+			runner_loop(channel, options);
 		})
 	}
 	.expect("fork()")
@@ -109,7 +122,7 @@ impl Runner {
 	/// Creates a new container runner
 	///
 	/// Unsafe: Do not call in multithreaded processes
-	pub unsafe fn new() -> Result<Self> {
+	pub unsafe fn new(options: &Options) -> Result<Self> {
 		let lockfile = Runner::lock()?;
 
 		let uid = unistd::geteuid();
@@ -122,7 +135,7 @@ impl Runner {
 			(tx, rx),
 			|(tx, rx)| {
 				drop(tx);
-				runner(uid, gid, rx, lockfile);
+				runner(uid, gid, rx, lockfile, options);
 			},
 		)
 		.expect("clone()")