From cdadc739cd267e736d30c5bee111bbd0326c974b 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.cc                       | 19 ++++++++
 src/libstore/globals.hh                       | 48 ++++++++++++++++++-
 .../unix/build/local-derivation-goal.cc       |  4 ++
 src/nix-build/nix-build.cc                    |  2 +
 tests/functional/load-limit.nix               |  8 ++++
 tests/functional/load-limit.sh                | 21 ++++++++
 tests/functional/local.mk                     |  1 +
 9 files changed, 108 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 000000000000..869dcb3befc9
--- /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 6533b2f58c46..f94c7ae3f61c 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.cc b/src/libstore/globals.cc
index fa4c0ba7fad6..f62e86ba0a07 100644
--- a/src/libstore/globals.cc
+++ b/src/libstore/globals.cc
@@ -333,6 +333,25 @@ unsigned int MaxBuildJobsSetting::parse(const std::string & str) const
     }
 }
 
+std::optional<unsigned int> LoadLimitSetting::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);
+    }
+}
+
+std::string LoadLimitSetting::to_string() const
+{
+    if (value)
+        return std::to_string(value.value());
+    else
+        return "none";
+}
+
 
 static void preloadNSS()
 {
diff --git a/src/libstore/globals.hh b/src/libstore/globals.hh
index 30d7537bd52b..e3394eebee24 100644
--- a/src/libstore/globals.hh
+++ b/src/libstore/globals.hh
@@ -31,6 +31,24 @@ struct MaxBuildJobsSetting : public BaseSetting<unsigned int>
     unsigned int parse(const std::string & str) const override;
 };
 
+struct LoadLimitSetting : public Setting<std::optional<unsigned int>>
+{
+    LoadLimitSetting(Config * options,
+        const std::optional<unsigned int> & def,
+        const std::string & name,
+        const std::string & description,
+        const std::set<std::string> & aliases = {},
+        const bool documentDefault = true,
+        std::optional<ExperimentalFeature> experimentalFeature = std::nullopt)
+        : Setting<std::optional<unsigned int>>(
+            options, def, name, description, aliases, documentDefault, std::move(experimentalFeature))
+    {}
+
+    std::optional<unsigned int> parse(const std::string & str) const override;
+    std::string to_string() const override;
+};
+
+
 const uint32_t maxIdsPerBuild =
     #if __linux__
     1 << 16
@@ -166,7 +184,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 +193,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 +203,33 @@ public:
         // Don't document the machine-specific default value
         false};
 
+    LoadLimitSetting 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 0dd102200a5f..417577a35d8c 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/nix-build/nix-build.cc b/src/nix-build/nix-build.cc
index 0ce987d8a5c6..5b397d3091bb 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 000000000000..33c81ea5595f
--- /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 000000000000..040949b6c6f9
--- /dev/null
+++ b/tests/functional/load-limit.sh
@@ -0,0 +1,21 @@
+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 49ee31284bc2..8fee9e97fa27 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 \