P1: address review (fail-open hardening)

lix/legacy/build-remote.cc

@@ -1,8 +1,21 @@
 #include <algorithm>
 #include <chrono>
 #include <set>
+#include <map>
 #include <memory>
+#include <optional>
 #include <tuple>
+#include <fstream>
+#include <sstream>
+#include <cstring>
+#include <cerrno>
+#include <sys/socket.h>
+#include <netdb.h>
+#include <netinet/in.h>
+#include <poll.h>
+#include <fcntl.h>
+#include <unistd.h>
+#include <nlohmann/json.hpp>
 #if __APPLE__
 #include <sys/time.h>
 #endif
@@ -53,6 +66,220 @@ static bool allSupportedLocally(Store & store, const std::set<std::string>& requ
     return true;
 }
 
+/* --------------------------------------------------------------------------
+ * P1: load- and memory-aware adaptive remote-build selection.
+ *
+ * All state below is populated ONCE from out-of-band, env-driven config
+ * (never from the derivation). Every helper FAILS OPEN: if config is unset,
+ * the metrics socket is unreachable/slow/malformed, or the storeUri is
+ * unknown, the helpers behave exactly like unpatched Lix
+ * (machineHasRoom -> true, liveLoadPenalty -> 0).
+ * ------------------------------------------------------------------------ */
+
+// A drv that does not match the heavy-crate table is treated as "light":
+// we have NO confident signal that it is memory-heavy, so adaptiveEstPeakRSS
+// returns nullopt and machineHasRoom never filters on account of it. (There is
+// deliberately no numeric light default - an unmatched drv must always permit,
+// so keying it on a free-RAM threshold would wrongly filter light drvs.)
+
+// name-substring -> estimated peak RSS in MiB (from LIX_ADAPTIVE_RSS_TABLE).
+static std::map<std::string, uint64_t> adaptiveRssTable;
+// machine storeUri -> "host:port" metrics endpoint (from LIX_ADAPTIVE_METRICS_MAP).
+static std::map<std::string, std::string> adaptiveMetricsMap;
+
+struct AdaptiveProbe {
+    bool ok = false;
+    uint64_t memAvailKb = 0;
+    double psiMem = 0, psiIo = 0, psiCpu = 0, load1 = 0, nproc = 0;
+};
+
+// In-process TTL cache keyed by storeUri, so selection probes each machine
+// at most once every ~2s regardless of how many drvs stream through.
+static std::map<std::string, std::pair<std::chrono::steady_clock::time_point, AdaptiveProbe>> adaptiveProbeCache;
+
+/* Parse the two env-driven config sources once. Any error leaves the tables
+ * empty, which degrades to unpatched behavior. */
+static void adaptiveLoadConfig()
+{
+    // LIX_ADAPTIVE_RSS_TABLE is a PATH to a JSON object {substring: MiB}.
+    try {
+        if (auto p = getEnv("LIX_ADAPTIVE_RSS_TABLE")) {
+            std::ifstream f(*p);
+            if (f) {
+                nlohmann::json j;
+                f >> j;
+                if (j.is_object())
+                    for (auto & [k, v] : j.items())
+                        // Per-entry guard: one bad value skips only that entry,
+                        // it does not discard the whole (otherwise valid) table.
+                        try {
+                            if (v.is_number_unsigned() || (v.is_number_integer() && v.get<int64_t>() >= 0))
+                                adaptiveRssTable[k] = v.get<uint64_t>();
+                        } catch (...) { continue; }
+            }
+        }
+    } catch (...) { adaptiveRssTable.clear(); }
+
+    // LIX_ADAPTIVE_METRICS_MAP is an inline JSON object {storeUri: "host:port"}.
+    try {
+        if (auto m = getEnv("LIX_ADAPTIVE_METRICS_MAP")) {
+            auto j = nlohmann::json::parse(*m);
+            if (j.is_object())
+                for (auto & [k, v] : j.items())
+                    // Per-entry guard: skip one bad value, keep the rest.
+                    try {
+                        if (v.is_string())
+                            adaptiveMetricsMap[k] = v.get<std::string>();
+                    } catch (...) { continue; }
+        }
+    } catch (...) { adaptiveMetricsMap.clear(); }
+}
+
+/* Estimated peak RSS (MiB) for a drv, or nullopt when the drv does not match
+ * the heavy-crate table. nullopt == "no confident heavy signal". Keyed on the
+ * store-path NAME, which is available before readDerivation and never mutates
+ * the drv. */
+static std::optional<uint64_t> adaptiveEstPeakRSS(const StorePath & drvPath)
+{
+    if (adaptiveRssTable.empty()) return std::nullopt;
+    std::string_view name = drvPath.name();
+    std::optional<uint64_t> best;
+    for (auto & [sub, mib] : adaptiveRssTable)
+        if (!sub.empty() && name.find(sub) != std::string_view::npos)
+            best = std::max(best.value_or(0), mib);
+    return best;
+}
+
+/* TCP-connect the metrics endpoint and read one line:
+ *   "MemAvail_kB psi_mem psi_io psi_cpu load1 nproc"
+ * A single ~500ms wall-clock deadline bounds the WHOLE probe (resolve +
+ * connect + read) so selection NEVER hangs, regardless of a slow or
+ * byte-dribbling peer. The endpoint MUST be a numeric IP:port - resolution is
+ * pinned to AI_NUMERICHOST|AI_NUMERICSERV so getaddrinfo never does network
+ * I/O (a hostname simply fails fast -> fail-open). Any failure returns an
+ * AdaptiveProbe with ok=false. */
+static AdaptiveProbe adaptiveProbeEndpoint(const std::string & hostport)
+{
+    AdaptiveProbe r;
+    auto colon = hostport.rfind(':');
+    if (colon == std::string::npos || colon == 0 || colon + 1 >= hostport.size())
+        return r;
+    std::string host = hostport.substr(0, colon);
+    std::string port = hostport.substr(colon + 1);
+
+    // Single wall-clock budget for the entire probe.
+    auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500);
+    auto remainingMs = [&]() -> int {
+        auto d = std::chrono::duration_cast<std::chrono::milliseconds>(
+            deadline - std::chrono::steady_clock::now()).count();
+        return d <= 0 ? 0 : (int) d;
+    };
+
+    struct addrinfo hints;
+    memset(&hints, 0, sizeof hints);
+    hints.ai_family = AF_UNSPEC;
+    hints.ai_socktype = SOCK_STREAM;
+    // Numeric-only: no DNS, no resolver blocking. Non-IP endpoint -> fail-open.
+    hints.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV;
+    struct addrinfo * res = nullptr;
+    if (getaddrinfo(host.c_str(), port.c_str(), &hints, &res) != 0 || !res)
+        return r;
+
+    int fd = socket(res->ai_family, res->ai_socktype | SOCK_NONBLOCK, res->ai_protocol);
+    if (fd < 0) { freeaddrinfo(res); return r; }
+
+    int cr = connect(fd, res->ai_addr, res->ai_addrlen);
+    if (cr < 0 && errno == EINPROGRESS) {
+        struct pollfd pfd;
+        pfd.fd = fd;
+        pfd.events = POLLOUT;
+        if (poll(&pfd, 1, remainingMs()) <= 0) { close(fd); freeaddrinfo(res); return r; }
+        int soerr = 0;
+        socklen_t sl = sizeof soerr;
+        if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &soerr, &sl) < 0 || soerr != 0) {
+            close(fd); freeaddrinfo(res); return r;
+        }
+    } else if (cr < 0) {
+        close(fd); freeaddrinfo(res); return r;
+    }
+    freeaddrinfo(res);
+
+    /* Read one short line. Keep the socket non-blocking and gate every recv on
+       poll(POLLIN) against the shared deadline, so the total read time is
+       bounded even if the peer drips one byte at a time. A valid reply is tiny,
+       so also cap the number of reads. */
+    std::string line;
+    char buf[512];
+    for (int iter = 0; iter < 16 && line.size() < 4096; ++iter) {
+        int rem = remainingMs();
+        if (rem == 0) break;
+        struct pollfd pfd;
+        pfd.fd = fd;
+        pfd.events = POLLIN;
+        int pr = poll(&pfd, 1, rem);
+        if (pr <= 0) break;                 // timeout or error -> fail-open
+        if (!(pfd.revents & POLLIN)) break; // POLLHUP/POLLERR with no data
+        ssize_t n = recv(fd, buf, sizeof buf, 0);
+        if (n < 0) {
+            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
+            break;
+        }
+        if (n == 0) break;                  // peer closed
+        line.append(buf, n);
+        if (line.find('\n') != std::string::npos) break;
+    }
+    close(fd);
+
+    std::istringstream ss(line);
+    AdaptiveProbe tmp;
+    if (ss >> tmp.memAvailKb >> tmp.psiMem >> tmp.psiIo >> tmp.psiCpu >> tmp.load1 >> tmp.nproc) {
+        tmp.ok = true;
+        return tmp;
+    }
+    return r;
+}
+
+/* Cached probe for a machine. Unknown storeUri -> ok=false (fail-open). */
+static AdaptiveProbe adaptiveProbe(const Machine & m)
+{
+    auto now = std::chrono::steady_clock::now();
+    auto it = adaptiveProbeCache.find(m.storeUri);
+    if (it != adaptiveProbeCache.end() && now - it->second.first < std::chrono::seconds(2))
+        return it->second.second;
+
+    AdaptiveProbe r;
+    auto mit = adaptiveMetricsMap.find(m.storeUri);
+    if (mit != adaptiveMetricsMap.end())
+        r = adaptiveProbeEndpoint(mit->second);
+
+    adaptiveProbeCache[m.storeUri] = { now, r };
+    return r;
+}
+
+/* OOM guard. Returns TRUE (permit as a candidate) UNLESS we have a confident
+ * signal that the drv is heavy AND the machine's free RAM is below the drv's
+ * estimated peak RSS. No env, dead socket, or unknown machine -> permit. */
+static bool machineHasRoom(const Machine & m, const StorePath & drvPath)
+{
+    auto est = adaptiveEstPeakRSS(drvPath);
+    if (!est) return true;              // no confident heavy signal
+    auto p = adaptiveProbe(m);
+    if (!p.ok) return true;             // no live signal -> fail open
+    uint64_t freeMib = p.memAvailKb / 1024;
+    return freeMib >= *est;
+}
+
+/* Extra ranking cost from live pressure on a machine; 0 when no signal. */
+static double liveLoadPenalty(const Machine & m)
+{
+    auto p = adaptiveProbe(m);
+    if (!p.ok) return 0.0;
+    double penalty = 0.0;
+    penalty += p.psiIo / 10.0;                       // io-PSI (0..100) -> up to 10
+    if (p.nproc > 0) penalty += p.load1 / p.nproc;   // load normalized by cores
+    return penalty;
+}
+
 static int main_build_remote(AsyncIoRoot & aio, std::string programName, Strings argv)
 {
     {
@@ -108,6 +335,9 @@ static int main_build_remote(AsyncIoRoot & aio, std::string programName, Strings
         std::optional<StorePath> drvPath;
         std::string storeUri;
 
+        /* P1: parse out-of-band adaptive config once (fail-open on any error). */
+        adaptiveLoadConfig();
+
         while (true) {
 
             try {
@@ -140,14 +370,15 @@ static int main_build_remote(AsyncIoRoot & aio, std::string programName, Strings
                 bool rightType = false;
 
                 Machine * bestMachine = nullptr;
-                uint64_t bestLoad = 0;
+                double bestCost = 0;
                 for (auto & m : machines) {
                     debug("considering building on remote machine '%s'", m.storeUri);
 
                     if (m.enabled &&
                         m.systemSupported(neededSystem) &&
                         m.allSupported(requiredFeatures) &&
-                        m.mandatoryMet(requiredFeatures))
+                        m.mandatoryMet(requiredFeatures) &&
+                        machineHasRoom(m, *drvPath))
                     {
                         rightType = true;
                         AutoCloseFD free;
@@ -165,22 +396,21 @@ static int main_build_remote(AsyncIoRoot & aio, std::string programName, Strings
                         if (!free) {
                             continue;
                         }
+                        /* P1: ranking cost folds in live pressure (0 when no
+                           signal, so this reduces to load / speedFactor). */
+                        double cost = (double(load) + liveLoadPenalty(m)) / m.speedFactor;
                         bool best = false;
                         if (!bestSlotLock) {
                             best = true;
-                        } else if (load / m.speedFactor < bestLoad / bestMachine->speedFactor) {
+                        } else if (cost < bestCost) {
                             best = true;
-                        } else if (load / m.speedFactor == bestLoad / bestMachine->speedFactor) {
+                        } else if (cost == bestCost) {
                             if (m.speedFactor > bestMachine->speedFactor) {
                                 best = true;
-                            } else if (m.speedFactor == bestMachine->speedFactor) {
-                                if (load < bestLoad) {
-                                    best = true;
-                                }
                             }
                         }
                         if (best) {
-                            bestLoad = load;
+                            bestCost = cost;
                             bestSlotLock = std::move(free);
                             bestMachine = &m;
                         }