From 23a62b06d697f2beb8350347de9ea4241f1430b5 Mon Sep 17 00:00:00 2001
From: Matthias Schiffer <mschiffer@universe-factory.net>
Date: Mon, 25 Oct 2021 23:29:30 +0200
Subject: Allow to set the maximum number of jobs on the command line

---
 crates/executor/src/main.rs |  7 ++++++-
 crates/runner/src/lib.rs    | 27 ++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

(limited to 'crates')

diff --git a/crates/executor/src/main.rs b/crates/executor/src/main.rs
index 4a045be..6360e69 100644
--- a/crates/executor/src/main.rs
+++ b/crates/executor/src/main.rs
@@ -13,6 +13,11 @@ use runner::Runner;
 #[derive(Parser)]
 #[clap(version = clap::crate_version!())]
 struct Opts {
+	/// Allow N jobs at once.
+	/// Defaults to the number of available CPUs
+	#[clap(short, long)]
+	jobs: Option<usize>,
+	/// The tasks to run
 	#[clap(name = "task", required = true)]
 	tasks: Vec<String>,
 }
@@ -20,7 +25,7 @@ struct Opts {
 fn main() {
 	let opts: Opts = Opts::parse();
 
-	let runner = unsafe { Runner::new() }.unwrap();
+	let runner = unsafe { Runner::new(&runner::Options { jobs: opts.jobs }) }.unwrap();
 
 	let ctx = context::Context::new(recipe::read_recipes("examples").unwrap());
 
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()")
-- 
cgit v1.2.3