From 13185133bcdff751274e55bf29b80b4e600cd973 Mon Sep 17 00:00:00 2001
From: Matej Urbas <matej.urbas@gmail.com>
Date: Mon, 8 May 2023 19:45:46 +0100
Subject: [PATCH] introduces `Goal::jobCategory`

---
 src/libstore/build/derivation-goal.hh              |  2 ++
 src/libstore/build/drv-output-substitution-goal.hh |  4 +++-
 src/libstore/build/goal.hh                         | 13 +++++++++++++
 src/libstore/build/substitution-goal.hh            |  2 ++
 src/libstore/build/worker.cc                       |  6 +++---
 5 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/libstore/build/derivation-goal.hh b/src/libstore/build/derivation-goal.hh
index 7033b7a58..223be06f2 100644
--- a/src/libstore/build/derivation-goal.hh
+++ b/src/libstore/build/derivation-goal.hh
@@ -335,6 +335,8 @@ struct DerivationGoal : public Goal
     void waiteeDone(GoalPtr waitee, ExitCode result) override;
 
     StorePathSet exportReferences(const StorePathSet & storePaths);
+
+    JobCategory jobCategory() override { return JobCategory::Build; };
 };
 
 MakeError(NotDeterministic, BuildError);
diff --git a/src/libstore/build/drv-output-substitution-goal.hh b/src/libstore/build/drv-output-substitution-goal.hh
index 697ddb283..5d1253a71 100644
--- a/src/libstore/build/drv-output-substitution-goal.hh
+++ b/src/libstore/build/drv-output-substitution-goal.hh
@@ -21,7 +21,7 @@ class Worker;
 class DrvOutputSubstitutionGoal : public Goal {
 
     /**
-     * The drv output we're trying to substitue
+     * The drv output we're trying to substitute
      */
     DrvOutput id;
 
@@ -72,6 +72,8 @@ public:
 
     void work() override;
     void handleEOF(int fd) override;
+
+    JobCategory jobCategory() override { return JobCategory::Substitution; };
 };
 
 }
diff --git a/src/libstore/build/goal.hh b/src/libstore/build/goal.hh
index c0e12a2ed..a313bf22c 100644
--- a/src/libstore/build/goal.hh
+++ b/src/libstore/build/goal.hh
@@ -34,6 +34,17 @@ typedef std::set<WeakGoalPtr, std::owner_less<WeakGoalPtr>> WeakGoals;
  */
 typedef std::map<StorePath, WeakGoalPtr> WeakGoalMap;
 
+/**
+ * Used as a hint to the worker on how to schedule a particular goal. For example,
+ * builds are typically CPU- and memory-bound, while substitutions are I/O bound.
+ * Using this information, the worker might decide to schedule more or fewer goals
+ * of each category in parallel.
+ */
+enum struct JobCategory {
+    Build,
+    Substitution,
+};
+
 struct Goal : public std::enable_shared_from_this<Goal>
 {
     typedef enum {ecBusy, ecSuccess, ecFailed, ecNoSubstituters, ecIncompleteClosure} ExitCode;
@@ -150,6 +161,8 @@ public:
     void amDone(ExitCode result, std::optional<Error> ex = {});
 
     virtual void cleanup() { }
+
+    virtual JobCategory jobCategory() = 0;
 };
 
 void addToWeakGoals(WeakGoals & goals, GoalPtr p);
diff --git a/src/libstore/build/substitution-goal.hh b/src/libstore/build/substitution-goal.hh
index c2b7fc95a..9fc041920 100644
--- a/src/libstore/build/substitution-goal.hh
+++ b/src/libstore/build/substitution-goal.hh
@@ -115,6 +115,8 @@ public:
     void handleEOF(int fd) override;
 
     void cleanup() override;
+
+    JobCategory jobCategory() override { return JobCategory::Substitution; };
 };
 
 }
diff --git a/src/libstore/build/worker.cc b/src/libstore/build/worker.cc
index 85024e4c4..ee334d54a 100644
--- a/src/libstore/build/worker.cc
+++ b/src/libstore/build/worker.cc
@@ -195,7 +195,7 @@ void Worker::childStarted(GoalPtr goal, const std::set<int> & fds,
     child.respectTimeouts = respectTimeouts;
     children.emplace_back(child);
     if (inBuildSlot) {
-        if (dynamic_cast<PathSubstitutionGoal *>(child.goal2)) nrSubstitutions++;
+        if (goal->jobCategory() == JobCategory::Substitution) nrSubstitutions++;
         else nrLocalBuilds++;
     }
 }
@@ -208,7 +208,7 @@ void Worker::childTerminated(Goal * goal, bool wakeSleepers)
     if (i == children.end()) return;
 
     if (i->inBuildSlot) {
-        if (dynamic_cast<PathSubstitutionGoal *>(goal)) {
+        if (goal->jobCategory() == JobCategory::Substitution) {
             assert(nrSubstitutions > 0);
             nrSubstitutions--;
         } else {
@@ -235,7 +235,7 @@ void Worker::childTerminated(Goal * goal, bool wakeSleepers)
 void Worker::waitForBuildSlot(GoalPtr goal)
 {
     debug("wait for build slot");
-    bool isSubstitutionGoal = dynamic_cast<PathSubstitutionGoal *>(goal.get());
+    bool isSubstitutionGoal = goal->jobCategory() == JobCategory::Substitution;
     if ((!isSubstitutionGoal && getNrLocalBuilds() < settings.maxBuildJobs) ||
         (isSubstitutionGoal && getNrSubstitutions() < settings.maxSubstitutionJobs))
         wakeUp(goal); /* we can do it right away */