From a99842e0af7cb8a1b6adb2cddbbe6a63b1e9ee4d Mon Sep 17 00:00:00 2001 From: Alois Wohlschlager Date: Mon, 13 Oct 2025 11:32:46 +0200 Subject: [PATCH] nix-eval-jobs: asyncify the coordinator Change-Id: I62007ed4d43029a4c793c0152c06490d6a6a6964 --- subprojects/nix-eval-jobs/src/buffered-io.cc | 35 +++ subprojects/nix-eval-jobs/src/buffered-io.hh | 14 + .../nix-eval-jobs/src/nix-eval-jobs.cc | 284 +++++++++--------- 3 files changed, 196 insertions(+), 137 deletions(-) diff --git a/subprojects/nix-eval-jobs/src/buffered-io.cc b/subprojects/nix-eval-jobs/src/buffered-io.cc index 469dbe23f..9a48fe78a 100644 --- a/subprojects/nix-eval-jobs/src/buffered-io.cc +++ b/subprojects/nix-eval-jobs/src/buffered-io.cc @@ -56,3 +56,38 @@ LineReader::LineReader(LineReader &&other) { // Remove trailing newline return std::string_view(buffer, read - 1); } + +AsyncLineReader::AsyncLineReader(nix::AutoCloseFD fd) + : stream{std::move(fd)}, readBuffer{kj::heapArray(4096)} {} + +kj::Promise>> AsyncLineReader::readLine() +try { + auto pos = buffer.find('\n'); + if (pos != std::string::npos) { + std::string result = buffer.substr(0, pos); + memmove(buffer.data(), buffer.data() + pos + 1, buffer.size() - pos - 1); + buffer.resize(buffer.size() - pos - 1); + co_return result; + } + + // No full line was buffered, read until we have one. + while (true) { + auto nRead = LIX_TRY_AWAIT(stream.read(readBuffer.begin(), readBuffer.size())); + if (!nRead && !buffer.empty()) { + // File has ended, but not everything has been read out of the buffer yet. + co_return std::move(buffer); + } else if (!nRead) { + co_return std::nullopt; + } + std::string_view readStr{readBuffer.begin(), *nRead}; + + auto pos = readStr.find('\n'); + if (pos != std::string_view::npos) { + buffer.append(readStr.substr(0, pos)); + co_return std::exchange(buffer, readStr.substr(pos + 1)); + } + buffer.append(readStr); + } +} catch (...) { + co_return nix::result::current_exception(); +} diff --git a/subprojects/nix-eval-jobs/src/buffered-io.hh b/subprojects/nix-eval-jobs/src/buffered-io.hh index e1f068500..47573ecb3 100644 --- a/subprojects/nix-eval-jobs/src/buffered-io.hh +++ b/subprojects/nix-eval-jobs/src/buffered-io.hh @@ -1,4 +1,6 @@ #pragma once +#include +#include #include #include #include @@ -18,3 +20,15 @@ class LineReader { char *buffer = nullptr; size_t len = 0; }; + +class AsyncLineReader { +public: + AsyncLineReader(nix::AutoCloseFD fd); + + kj::Promise>> readLine(); + +private: + nix::AsyncFdIoStream stream; + std::string buffer; + kj::Array readBuffer; +}; diff --git a/subprojects/nix-eval-jobs/src/nix-eval-jobs.cc b/subprojects/nix-eval-jobs/src/nix-eval-jobs.cc index 8705e3d06..cfb1d273d 100644 --- a/subprojects/nix-eval-jobs/src/nix-eval-jobs.cc +++ b/subprojects/nix-eval-jobs/src/nix-eval-jobs.cc @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -25,14 +27,11 @@ #include #include #include -#include #include -#include #include #include #include #include -#include #include #include #include @@ -75,8 +74,8 @@ class Collector { using Response = std::variant; RunningProgram child; - AutoCloseFD to; - std::optional from; + std::optional to; + std::optional from; Strings workerCmdline; @@ -105,11 +104,12 @@ class Collector { }; child = runProgram2(options); - to = std::move(toPipe.writeSide); - from.emplace(fromPipe.readSide.release()); + to.emplace(std::move(toPipe.writeSide)); + from.emplace(std::move(fromPipe.readSide)); } - void waitForWorkerReady() { + kj::Promise> waitForWorkerReady() + try { assert(child); std::visit(overloaded{ [](const Next &) {}, @@ -121,53 +121,66 @@ class Collector { to.reset(); from.reset(); }, - }, readResponse("checking worker process")); + }, LIX_TRY_AWAIT(readResponse("checking worker process"))); + co_return result::success(); + } catch (...) { + co_return result::current_exception(); } - void makeWorkerReady() { + kj::Promise> makeWorkerReady() + try { if (child) { - waitForWorkerReady(); + LIX_TRY_AWAIT(waitForWorkerReady()); } if (!child) { startWorker(); - waitForWorkerReady(); + LIX_TRY_AWAIT(waitForWorkerReady()); } if (!child) { throw Error("worker exited immediately"); } + co_return result::success(); + } catch (...) { + co_return result::current_exception(); } - Response readResponse(std::string_view msg) { + kj::Promise> readResponse(std::string_view msg) + try { assert(from); - auto line = from->readLine(); - if (line.empty()) { + auto line = LIX_TRY_AWAIT(from->readLine()); + if (!line) { handleBrokenPipe(msg); } else if (line == "next") { - return Next{}; + co_return Next{}; } else if (line == "restart") { - return Restart{}; + co_return Restart{}; } else { try { - return JsonResponse{JSON::parse(line)}; + co_return JsonResponse{JSON::parse(*line)}; } catch (const json::ParseError &e) { throw Error( "Received invalid JSON from worker: %s\n json: '%s'", - e.what(), line); + e.what(), *line); } } + } catch (...) { + co_return result::current_exception(); } - void writeRequest(Request request) { + kj::Promise> writeRequest(Request request) + try { assert(to); auto line = std::visit(overloaded{ [](const Do &request) { - return fmt("do %s", request.attrPath.dump()); + return fmt("do %s\n", request.attrPath.dump()); }, [](const Exit &) { - return std::string{"exit"}; + return std::string{"exit\n"}; }, }, request); - if (tryWriteLine(to.get(), line) < 0) { + try { + LIX_TRY_AWAIT(to->writeFull(line.data(), line.size())); + } catch (SysError &err) { auto msg = std::visit(overloaded{ [](const Do &request) { return fmt("sending attrPath '%s'", joinAttrPath(request.attrPath)); @@ -178,6 +191,9 @@ class Collector { }, request); handleBrokenPipe(msg); } + co_return result::success(); + } catch (...) { + co_return result::current_exception(); } [[noreturn]] void handleBrokenPipe(std::string_view msg) { @@ -240,10 +256,11 @@ public: } } - JSON evaluate(JSON attrPath) { - makeWorkerReady(); - writeRequest(Do{attrPath}); - return std::visit(overloaded{ + kj::Promise> evaluate(JSON attrPath) + try { + LIX_TRY_AWAIT(makeWorkerReady()); + LIX_TRY_AWAIT(writeRequest(Do{attrPath})); + co_return std::visit(overloaded{ [](const Next &) -> JSON { throw Error("unexpected response from worker: next"); }, @@ -253,93 +270,101 @@ public: [](const Restart &) -> JSON { throw Error("unexpected response from worker: restart"); }, - }, readResponse(fmt("reading result for attrPath '%s'", joinAttrPath(attrPath)))); - } - - void exit() { - if (child) { - waitForWorkerReady(); - } - if (child) { - writeRequest(Exit{}); - // The worker will print "restart" when exiting cleanly, even if due to an explicit exit request. - waitForWorkerReady(); - } - } -}; - -struct State { - std::set todo = JSON::array({JSON::array()}); - std::set active; - std::exception_ptr exc; - std::map jobs; -}; - -void collectorThread(MyArgs &myArgs, Sync &state_, std::condition_variable &wakeup) { - try { - Collector collector{myArgs.cmdline}; - - while (true) { - /* Wait for a job name to become available. */ - JSON attrPath; - while (true) { - checkInterrupt(); - auto state(state_.lock()); - if ((state->todo.empty() && state->active.empty()) || - state->exc) { - collector.exit(); - return; - } - if (!state->todo.empty()) { - attrPath = *state->todo.begin(); - state->todo.erase(state->todo.begin()); - state->active.insert(attrPath); - break; - } else - state.wait(wakeup); - } - - /* Tell the worker to evaluate it. */ - auto response = collector.evaluate(attrPath); - - /* Handle the response. */ - std::vector newAttrs; - if (response.find("attrs") != response.end()) { - for (auto &i : response["attrs"]) { - JSON newAttr = JSON(response["attrPath"]); - newAttr.emplace_back(i); - newAttrs.push_back(newAttr); - } - } else { - auto state(state_.lock()); - state->jobs.insert_or_assign(response["attr"], response); - if (nix::settings.readOnlyMode) { - response.erase("namedConstituents"); - response.erase("constituents"); - } - auto named = response.find("namedConstituents"); - if (named == response.end() || named->empty()) { - response.erase("namedConstituents"); - logger->writeToStdout(response.dump()); - } - } - - /* Add newly discovered job names to the queue. */ - { - auto state(state_.lock()); - state->active.erase(attrPath); - for (auto p : newAttrs) { - state->todo.insert(p); - } - wakeup.notify_all(); - } - } + }, LIX_TRY_AWAIT(readResponse(fmt("reading result for attrPath '%s'", joinAttrPath(attrPath))))); } catch (...) { - auto state(state_.lock()); - state->exc = std::current_exception(); - wakeup.notify_all(); + co_return result::current_exception(); } -} + + kj::Promise> exit() + try { + if (child) { + LIX_TRY_AWAIT(waitForWorkerReady()); + } + if (child) { + LIX_TRY_AWAIT(writeRequest(Exit{})); + // The worker will print "restart" when exiting cleanly, even if due to an explicit exit request. + LIX_TRY_AWAIT(waitForWorkerReady()); + } + co_return result::success(); + } catch (...) { + co_return result::current_exception(); + } +}; + +class Coordinator { + AsyncSemaphore semaphore; + kj::Array workers; + std::vector idleWorkers; + + kj::Promise> evaluateRecursively(JSON attrPath, std::map &jobs) + try { + JSON response; + { + auto _token = co_await semaphore.acquire(); + auto *worker = idleWorkers.back(); + idleWorkers.pop_back(); + Finally _returnWorker{[&]() { + idleWorkers.push_back(worker); + }}; + response = LIX_TRY_AWAIT(worker->evaluate(attrPath)); + } + + std::vector newAttrs; + if (response.find("attrs") != response.end()) { + for (auto &i : response["attrs"]) { + JSON newAttr = JSON(response["attrPath"]); + newAttr.emplace_back(i); + newAttrs.push_back(newAttr); + } + } else { + jobs.insert_or_assign(response["attr"], response); + if (nix::settings.readOnlyMode) { + response.erase("namedConstituents"); + response.erase("constituents"); + } + auto named = response.find("namedConstituents"); + if (named == response.end() || named->empty()) { + response.erase("namedConstituents"); + nix::logger->writeToStdout(response.dump()); + } + } + + LIX_TRY_AWAIT(asyncSpread(newAttrs, [&](const JSON &newAttr) { + return evaluateRecursively(newAttr, jobs); + })); + + co_return result::success(); + } catch (...) { + co_return result::current_exception(); + } + +public: + Coordinator(size_t nrWorkers, const Strings &cmdline) + : semaphore{static_cast(nrWorkers)} { + if (nrWorkers >= std::numeric_limits::max() - 1) { + throw Error("nix-eval-jobs cannot handle %d workers, please choose a reasonable number"); + } + + auto builder = kj::heapArrayBuilder(nrWorkers); + for (unsigned i = 0; i < nrWorkers; ++i) { + auto worker = &builder.add(cmdline); + idleWorkers.push_back(worker); + } + workers = builder.finish(); + } + + kj::Promise>> run() + try { + std::map jobs; + LIX_TRY_AWAIT(evaluateRecursively(JSON::array(), jobs)); + for (auto &worker : workers) { + LIX_TRY_AWAIT(worker.exit()); + } + co_return jobs; + } catch (...) { + co_return result::current_exception(); + } +}; int main(int argc, char **argv) { return handleExceptions(argv[0], [&]() { @@ -390,23 +415,8 @@ int main(int argc, char **argv) { return 0; } - Sync state_; - - /* Start a collector thread per worker process. */ - std::vector threads; - std::condition_variable wakeup; - for (size_t i = 0; i < myArgs.nrWorkers; i++) { - threads.emplace_back(std::bind(collectorThread, std::ref(myArgs), - std::ref(state_), std::ref(wakeup))); - } - - for (auto &thread : threads) - thread.join(); - - auto state(state_.lock()); - - if (state->exc) - std::rethrow_exception(state->exc); + Coordinator coordinator{myArgs.nrWorkers, myArgs.cmdline}; + auto jobs = aio.blockOn(coordinator.run()); if (myArgs.constituents) { auto store = aio.blockOn(myArgs.evalStoreUrl @@ -415,28 +425,28 @@ int main(int argc, char **argv) { std::visit( nix::overloaded{ [&](const std::vector &namedConstituents) { - rewriteAggregates(state->jobs, namedConstituents, store, + rewriteAggregates(jobs, namedConstituents, store, myArgs.gcRootsDir, aio); }, [&](const DependencyCycle &e) { printError( "Found dependency cycle between jobs '%s' and '%s'", e.a, e.b); - state->jobs[e.a]["error"] = e.message(); - state->jobs[e.b]["error"] = e.message(); + jobs[e.a]["error"] = e.message(); + jobs[e.b]["error"] = e.message(); - logger->writeToStdout(state->jobs[e.a].dump()); - logger->writeToStdout(state->jobs[e.b].dump()); + nix::logger->writeToStdout(jobs[e.a].dump()); + nix::logger->writeToStdout(jobs[e.b].dump()); for (const auto &jobName : e.remainingAggregates) { - state->jobs[jobName]["error"] = + jobs[jobName]["error"] = "Skipping aggregate because of a dependency " "cycle"; - logger->writeToStdout(state->jobs[jobName].dump()); + nix::logger->writeToStdout(jobs[jobName].dump()); } }, }, - resolveNamedConstituents(state->jobs)); + resolveNamedConstituents(jobs)); } return 0;