From efa73225877d651c979288fda2b60674cad33c9b Mon Sep 17 00:00:00 2001
From: Linus Heckemann <git@sphalerite.org>
Date: Wed, 16 Feb 2022 10:10:52 +0100
Subject: [PATCH] streaming evaluator: respect eval-nodes-limit (#58)

---
 src/nix/deployment/mod.rs          | 2 ++
 src/nix/evaluator/mod.rs           | 4 ++++
 src/nix/evaluator/nix_eval_jobs.rs | 8 +++++++-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/nix/deployment/mod.rs b/src/nix/deployment/mod.rs
index 1454678..3e9f6fa 100644
--- a/src/nix/deployment/mod.rs
+++ b/src/nix/deployment/mod.rs
@@ -245,6 +245,8 @@ impl Deployment {
 
         let futures = job.run(|job| async move {
             let mut evaluator = NixEvalJobs::default();
+            let eval_limit = self.evaluation_node_limit.get_limit().unwrap_or_else(|| self.targets.len());
+            evaluator.set_eval_limit(eval_limit);
             evaluator.set_job(job.clone());
 
             // FIXME: nix-eval-jobs currently does not support IFD with builders
diff --git a/src/nix/evaluator/mod.rs b/src/nix/evaluator/mod.rs
index 70263bf..5c88ad5 100644
--- a/src/nix/evaluator/mod.rs
+++ b/src/nix/evaluator/mod.rs
@@ -60,6 +60,10 @@ pub trait DrvSetEvaluator {
     /// Evaluates an attribute set of derivation, returning results as they come in.
     async fn evaluate(&self, expression: &dyn NixExpression, options: NixOptions) -> ColmenaResult<Pin<Box<dyn Stream<Item = EvalResult>>>>;
 
+    /// Sets the maximum number of attributes to evaluate at the same time.
+    #[allow(unused_variables)]
+    fn set_eval_limit(&mut self, limit: usize) {}
+
     /// Provides a JobHandle to use during operations.
     #[allow(unused_variables)]
     fn set_job(&mut self, job: JobHandle) {}
diff --git a/src/nix/evaluator/nix_eval_jobs.rs b/src/nix/evaluator/nix_eval_jobs.rs
index 26a2419..a8a239d 100644
--- a/src/nix/evaluator/nix_eval_jobs.rs
+++ b/src/nix/evaluator/nix_eval_jobs.rs
@@ -30,6 +30,7 @@ pub const NIX_EVAL_JOBS: Option<&str> = option_env!("NIX_EVAL_JOBS");
 pub struct NixEvalJobs {
     executable: PathBuf,
     job: JobHandle,
+    workers: usize,
 }
 
 /// A line in the eval output.
@@ -82,7 +83,7 @@ impl DrvSetEvaluator for NixEvalJobs {
         let mut command = Command::new(&self.executable);
         command
             .arg("--impure")
-            .args(&["--workers", "10"]) // FIXME: Configurable
+            .arg("--workers").arg(self.workers.to_string())
             .arg(&expr_file);
 
         command.args(options.to_args());
@@ -160,6 +161,10 @@ impl DrvSetEvaluator for NixEvalJobs {
         }))
     }
 
+    fn set_eval_limit(&mut self, limit: usize) {
+        self.workers = limit;
+    }
+
     fn set_job(&mut self, job: JobHandle) {
         self.job = job;
     }
@@ -172,6 +177,7 @@ impl Default for NixEvalJobs {
         Self {
             executable: PathBuf::from(binary),
             job: null_job_handle(),
+            workers: 10,
         }
     }
 }