From 34f2477d2ef639a6ceb3243d78aad48df679f893 Mon Sep 17 00:00:00 2001
From: Emily <vcs@emily.moe>
Date: Thu, 18 Jul 2024 12:47:36 +0100
Subject: [PATCH] libstore: add load-limit setting to control parallelism

Closes: #7091
Closes: #6855
Closes: #8105
Co-authored-by: Alex Wied <centromere@users.noreply.github.com>
---
 doc/manual/src/release-notes/rl-next.md       |  4 +++
 src/libstore/daemon.cc                        |  2 ++
 src/libstore/globals.hh                       | 30 ++++++++++++++++++-
 .../unix/build/local-derivation-goal.cc       |  4 +++
 src/libutil/config-impl.hh                    |  1 +
 src/libutil/config.cc                         | 19 ++++++++++++
 src/nix-build/nix-build.cc                    |  2 ++
 tests/functional/load-limit.nix               |  8 +++++
 tests/functional/load-limit.sh                | 23 ++++++++++++++
 tests/functional/local.mk                     |  1 +
 10 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 doc/manual/src/release-notes/rl-next.md
 create mode 100644 tests/functional/load-limit.nix
 create mode 100644 tests/functional/load-limit.sh

diff --git a/doc/manual/src/release-notes/rl-next.md b/doc/manual/src/release-notes/rl-next.md
new file mode 100644
index 000000000..869dcb3be
--- /dev/null
+++ b/doc/manual/src/release-notes/rl-next.md
@@ -0,0 +1,4 @@
+# Release X.Y (202?-??-??)
+
+- Add a `load-limit` setting to control builder parallelism. This has
+  also been backported to the 2.18 and later release branches.
diff --git a/src/libstore/daemon.cc b/src/libstore/daemon.cc
index 6533b2f58..f94c7ae3f 100644
--- a/src/libstore/daemon.cc
+++ b/src/libstore/daemon.cc
@@ -199,6 +199,7 @@ struct ClientSettings
     time_t maxSilentTime;
     bool verboseBuild;
     unsigned int buildCores;
+    std::optional<unsigned int> loadLimit;
     bool useSubstitutes;
     StringMap overrides;
 
@@ -212,6 +213,7 @@ struct ClientSettings
         settings.maxSilentTime = maxSilentTime;
         settings.verboseBuild = verboseBuild;
         settings.buildCores = buildCores;
+        settings.loadLimit.assign(loadLimit);
         settings.useSubstitutes = useSubstitutes;
 
         for (auto & i : overrides) {
diff --git a/src/libstore/globals.hh b/src/libstore/globals.hh
index 30d7537bd..2a38920d8 100644
--- a/src/libstore/globals.hh
+++ b/src/libstore/globals.hh
@@ -166,7 +166,6 @@ public:
         R"(
           Sets the value of the `NIX_BUILD_CORES` environment variable in the [invocation of the `builder` executable](@docroot@/language/derivations.md#builder-execution) of a derivation.
           The `builder` executable can use this variable to control its own maximum amount of parallelism.
-
           <!--
           FIXME(@fricklerhandwerk): I don't think this should even be mentioned here.
           A very generic example using `derivation` and `xargs` may be more appropriate to explain the mechanism.
@@ -176,6 +175,8 @@ public:
 
           The value `0` means that the `builder` should use all available CPU cores in the system.
 
+          The [`load-limit`](#conf-load-limit) setting can be used to limit the total amount of build parallelism based on system load average.
+
           > **Note**
           >
           > The number of parallel local Nix build jobs is independently controlled with the [`max-jobs`](#conf-max-jobs) setting.
@@ -184,6 +185,33 @@ public:
         // Don't document the machine-specific default value
         false};
 
+    Setting<std::optional<unsigned int>> loadLimit{
+      this,
+      { getDefaultCores() },
+      "load-limit",
+      R"(
+        Sets the value of the `NIX_LOAD_LIMIT` environment variable in the
+        invocation of builders. Builders can use this value at their discretion
+        to dynamically control the amount of parallelism with respect to the
+        machine's load average.
+
+        For instance, a builder could use the value to set the `-l` flag to GNU
+        Make. In this case, if the load average of the machine exceeds
+        `NIX_LOAD_LIMIT`, the amount of parallelism will be dynamically
+        reduced.
+
+        By default, it is set to the number of cores on the machine.
+
+        On busy machines where Nix co-exists with other workloads, or where
+        build throughput is paramount and memory usage is not a bottleneck, the
+        default value may not work as intended. In this case, `load-limit`
+        should be set to a higher value, or to `none` to prevent the
+        `NIX_LOAD_LIMIT` variable being set at all.
+      )",
+      {},
+      // Don't document the machine-specific default value
+      false};
+
     /**
      * Read-only mode.  Don't copy stuff to the store, don't change
      * the database.
diff --git a/src/libstore/unix/build/local-derivation-goal.cc b/src/libstore/unix/build/local-derivation-goal.cc
index 0dd102200..417577a35 100644
--- a/src/libstore/unix/build/local-derivation-goal.cc
+++ b/src/libstore/unix/build/local-derivation-goal.cc
@@ -1156,6 +1156,10 @@ void LocalDerivationGoal::initEnv()
     /* The maximum number of cores to utilize for parallel building. */
     env["NIX_BUILD_CORES"] = fmt("%d", settings.buildCores);
 
+    /* Provide a load average limit for build tools to throttle jobs. */
+    if (settings.loadLimit.get())
+        env["NIX_LOAD_LIMIT"] = fmt("%d", settings.loadLimit.get().value());
+
     initTmpDir();
 
     /* Compatibility hack with Nix <= 0.7: if this is a fixed-output
diff --git a/src/libutil/config-impl.hh b/src/libutil/config-impl.hh
index 1d349fab5..9497e1d83 100644
--- a/src/libutil/config-impl.hh
+++ b/src/libutil/config-impl.hh
@@ -105,6 +105,7 @@ void BaseSetting<T>::convertToArg(Args & args, const std::string & category)
 
 DECLARE_CONFIG_SERIALISER(std::string)
 DECLARE_CONFIG_SERIALISER(std::optional<std::string>)
+DECLARE_CONFIG_SERIALISER(std::optional<unsigned int>)
 DECLARE_CONFIG_SERIALISER(bool)
 DECLARE_CONFIG_SERIALISER(Strings)
 DECLARE_CONFIG_SERIALISER(StringSet)
diff --git a/src/libutil/config.cc b/src/libutil/config.cc
index 726e5091e..239121b54 100644
--- a/src/libutil/config.cc
+++ b/src/libutil/config.cc
@@ -273,6 +273,25 @@ template<> std::string BaseSetting<std::optional<std::string>>::to_string() cons
     return value ? *value : "";
 }
 
+template<> std::optional<unsigned int> BaseSetting<std::optional<unsigned int>>::parse(const std::string & str) const
+{
+    if (str == "none") return std::nullopt;
+    else {
+        if (auto n = string2Int<unsigned int>(str))
+            return { *n };
+        else
+            throw UsageError("configuration setting '%s' should be 'none' or an integer", name);
+    }
+}
+
+template<> std::string BaseSetting<std::optional<unsigned int>>::to_string() const
+{
+    if (value)
+        return std::to_string(value.value());
+    else
+        return "none";
+}
+
 template<> bool BaseSetting<bool>::parse(const std::string & str) const
 {
     if (str == "true" || str == "yes" || str == "1")
diff --git a/src/nix-build/nix-build.cc b/src/nix-build/nix-build.cc
index 0ce987d8a..5b397d309 100644
--- a/src/nix-build/nix-build.cc
+++ b/src/nix-build/nix-build.cc
@@ -541,6 +541,8 @@ static void main_nix_build(int argc, char * * argv)
         env["NIX_BUILD_TOP"] = env["TMPDIR"] = env["TEMPDIR"] = env["TMP"] = env["TEMP"] = tmp;
         env["NIX_STORE"] = store->storeDir;
         env["NIX_BUILD_CORES"] = std::to_string(settings.buildCores);
+        if (settings.loadLimit.get())
+            env["NIX_LOAD_LIMIT"] = std::to_string(settings.loadLimit.get().value());
 
         auto passAsFile = tokenizeString<StringSet>(getOr(drv.env, "passAsFile", ""));
 
diff --git a/tests/functional/load-limit.nix b/tests/functional/load-limit.nix
new file mode 100644
index 000000000..33c81ea55
--- /dev/null
+++ b/tests/functional/load-limit.nix
@@ -0,0 +1,8 @@
+with import ./config.nix;
+
+mkDerivation {
+  name = "load-limit";
+  buildCommand = ''
+    printf '%s' "''${NIX_LOAD_LIMIT-unset}" > "$out"
+  '';
+}
diff --git a/tests/functional/load-limit.sh b/tests/functional/load-limit.sh
new file mode 100644
index 000000000..5711033fd
--- /dev/null
+++ b/tests/functional/load-limit.sh
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+source common.sh
+
+TODO_NixOS
+
+clearStore
+
+outPath=$(nix-build --load-limit 100 load-limit.nix --no-out-link)
+text=$(cat "$outPath")
+if test "$text" != "100"; then exit 1; fi
+
+clearStore
+
+outPath=$(nix-build --load-limit 0 load-limit.nix --no-out-link)
+text=$(cat "$outPath")
+if test "$text" != "0"; then exit 1; fi
+
+clearStore
+
+outPath=$(nix-build --load-limit none load-limit.nix --no-out-link)
+text=$(cat "$outPath")
+if test "$text" != "unset"; then exit 1; fi
diff --git a/tests/functional/local.mk b/tests/functional/local.mk
index 49ee31284..8fee9e97f 100644
--- a/tests/functional/local.mk
+++ b/tests/functional/local.mk
@@ -114,6 +114,7 @@ nix_tests = \
   ssh-relay.sh \
   build.sh \
   build-delete.sh \
+  load-limit.sh \
   output-normalization.sh \
   selfref-gc.sh \
   db-migration.sh \