From 48168af261f47e249a2d66e52d7a930ad493e933 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch <maximilian@mbosch.me>
Date: Fri, 28 Feb 2025 19:08:54 +0100
Subject: [PATCH] Update flake and fix build with mostly async liblixstore

The vast majority of changes was pretty straight-forward: we need an
AsyncIoRoot per thread to resolve promises we get from Lix. The `aio`
object is created per thread and passed as lvalue ref instead of making
it a field of the State struct since that struct is shared between the
threads.

The biggest change affects the NAR extractor: parseAndCopyDump has been
removed, instead the `nar::parse` and `nar::dump` parts are public now.
This won't work with the original parse visitor anymore, so instead the
NAR is parsed once and dumped again to copy it to the binary cache. In
between, the generator is walked through to search for relevant files
(such as build products) and after that the generators are recreated for
`AsyncGeneratorInputStream` to read again.
---
 flake.lock                                   |  14 +-
 src/hydra-queue-runner/build-remote.cc       |  40 +++--
 src/hydra-queue-runner/build-result.cc       |   3 +-
 src/hydra-queue-runner/builder.cc            |  22 +--
 src/hydra-queue-runner/dispatcher.cc         |  11 +-
 src/hydra-queue-runner/hydra-build-result.hh |   1 +
 src/hydra-queue-runner/hydra-queue-runner.cc |  26 ++-
 src/hydra-queue-runner/nar-extractor.cc      | 157 ++++++++++++-------
 src/hydra-queue-runner/queue-monitor.cc      |  22 +--
 src/hydra-queue-runner/state.hh              |  21 +--
 10 files changed, 188 insertions(+), 129 deletions(-)

diff --git a/flake.lock b/flake.lock
index 4fb3b4dba..a98d73a53 100644
--- a/flake.lock
+++ b/flake.lock
@@ -27,11 +27,11 @@
         "pre-commit-hooks": "pre-commit-hooks"
       },
       "locked": {
-        "lastModified": 1739016888,
-        "narHash": "sha256-JUOvTAYx+/bAec9H7C/nnpT1NRm/6tdy5EZ/XuUGHlA=",
+        "lastModified": 1741082941,
+        "narHash": "sha256-mxMbmNSXLZ0G+4uPEXCodjRJffqh/Jq4X5pgFuQFZB0=",
         "ref": "refs/heads/main",
-        "rev": "72326c404487bbf8dc6ee069930c6c2a0319857e",
-        "revCount": 17371,
+        "rev": "ca89e431a31527a014bfd0d529da2a8099027a5f",
+        "revCount": 17577,
         "type": "git",
         "url": "https://git.lix.systems/lix-project/lix"
       },
@@ -58,11 +58,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1738843498,
-        "narHash": "sha256-7x+Q4xgFj9UxZZO9aUDCR8h4vyYut4zPUvfj3i+jBHE=",
+        "lastModified": 1741048562,
+        "narHash": "sha256-W4YZ3fvWZiFYYyd900kh8P8wU6DHSiwaH0j4+fai1Sk=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "f5a32fa27df91dfc4b762671a0e0a859a8a0058f",
+        "rev": "6af28b834daca767a7ef99f8a7defa957d0ade6f",
         "type": "github"
       },
       "original": {
diff --git a/src/hydra-queue-runner/build-remote.cc b/src/hydra-queue-runner/build-remote.cc
index f1a5ce5e5..4fe4e2baa 100644
--- a/src/hydra-queue-runner/build-remote.cc
+++ b/src/hydra-queue-runner/build-remote.cc
@@ -11,6 +11,7 @@
 #include "lix/libstore/serve-protocol-impl.hh"
 #include "lix/libstore/serve-protocol.hh"
 #include "lix/libstore/ssh.hh"
+#include "lix/libutil/async.hh"
 #include "lix/libutil/current-process.hh"
 #include "lix/libutil/finally.hh"
 #include "lix/libutil/url.hh"
@@ -113,10 +114,11 @@ static void copyClosureTo(
     ::Machine::Connection & conn,
     Store & destStore,
     const StorePathSet & paths,
+    AsyncIoRoot & aio,
     SubstituteFlag useSubstitutes = NoSubstitute)
 {
     StorePathSet closure;
-    destStore.computeFSClosure(paths, closure);
+    aio.blockOn(destStore.computeFSClosure(paths, closure));
 
     /* Send the "query valid paths" command with the "lock" option
        enabled. This prevents a race where the remote host
@@ -133,7 +135,7 @@ static void copyClosureTo(
 
     if (present.size() == closure.size()) return;
 
-    auto sorted = destStore.topoSortPaths(closure);
+    auto sorted = aio.blockOn(destStore.topoSortPaths(closure));
 
     StorePathSet missing;
     for (auto & i : std::ranges::reverse_view(sorted))
@@ -145,7 +147,7 @@ static void copyClosureTo(
         std::chrono::seconds(600));
 
     conn.to << ServeProto::Command::ImportPaths;
-    destStore.exportPaths(missing, conn.to);
+    aio.blockOn(destStore.exportPaths(missing, conn.to));
     conn.to.flush();
 
     if (readInt(conn.from) != 1)
@@ -218,6 +220,7 @@ static void handshake(::Machine::Connection & conn, unsigned int repeats)
 }
 
 static BasicDerivation sendInputs(
+    AsyncIoRoot & aio,
     State & state,
     Step & step,
     Store & localStore,
@@ -236,7 +239,7 @@ static BasicDerivation sendInputs(
        to do that, however, but we would not use it here.)
      */
     BasicDerivation basicDrv = ({
-        auto maybeBasicDrv = step.drv->tryResolve(destStore, &localStore);
+        auto maybeBasicDrv = aio.blockOn(step.drv->tryResolve(destStore, &localStore));
         if (!maybeBasicDrv)
             throw Error(
                 "the derivation '%s' can’t be resolved. It’s probably "
@@ -250,9 +253,9 @@ static BasicDerivation sendInputs(
        this will copy the inputs to the binary cache from the local
        store. */
     if (&localStore != &destStore) {
-        copyClosure(localStore, destStore,
+        aio.blockOn(copyClosure(localStore, destStore,
             step.drv->inputSrcs,
-            NoRepair, NoCheckSigs, NoSubstitute);
+            NoRepair, NoCheckSigs, NoSubstitute));
     }
 
     {
@@ -268,10 +271,10 @@ static BasicDerivation sendInputs(
         /* Copy the input closure. */
         if (conn.machine->isLocalhost()) {
             StorePathSet closure;
-            destStore.computeFSClosure(basicDrv.inputSrcs, closure);
-            copyPaths(destStore, localStore, closure, NoRepair, NoCheckSigs, NoSubstitute);
+            aio.blockOn(destStore.computeFSClosure(basicDrv.inputSrcs, closure));
+            aio.blockOn(copyPaths(destStore, localStore, closure, NoRepair, NoCheckSigs, NoSubstitute));
         } else {
-            copyClosureTo(conn, destStore, basicDrv.inputSrcs, Substitute);
+            copyClosureTo(conn, destStore, basicDrv.inputSrcs, aio, Substitute);
         }
 
         auto now2 = std::chrono::steady_clock::now();
@@ -394,7 +397,8 @@ static void copyPathFromRemote(
     NarMemberDatas & narMembers,
     Store & localStore,
     Store & destStore,
-    const ValidPathInfo & info
+    const ValidPathInfo & info,
+    AsyncIoRoot & aio
 )
 {
       /* Receive the NAR from the remote and add it to the
@@ -413,9 +417,9 @@ static void copyPathFromRemote(
 
           co_yield extractNarDataFilter(conn.from, localStore.printStorePath(info.path), narMembers);
       };
-      GeneratorSource source2{coro()};
+      AsyncGeneratorInputStream inputStream{coro()};
 
-      destStore.addToStore(info, source2, NoRepair, NoCheckSigs);
+      aio.blockOn(destStore.addToStore(info, inputStream, NoRepair, NoCheckSigs));
 }
 
 static void copyPathsFromRemote(
@@ -423,14 +427,15 @@ static void copyPathsFromRemote(
     NarMemberDatas & narMembers,
     Store & localStore,
     Store & destStore,
-    const std::map<StorePath, ValidPathInfo> & infos
+    const std::map<StorePath, ValidPathInfo> & infos,
+    AsyncIoRoot & aio
 )
 {
       auto pathsSorted = reverseTopoSortPaths(infos);
 
       for (auto & path : pathsSorted) {
           auto & info = infos.find(path)->second;
-          copyPathFromRemote(conn, narMembers, localStore, destStore, info);
+          copyPathFromRemote(conn, narMembers, localStore, destStore, info, aio);
       }
 
 }
@@ -505,7 +510,8 @@ private:
     T* sem;
 };
 
-void State::buildRemote(ref<Store> destStore,
+void State::buildRemote(AsyncIoRoot & aio,
+    ref<Store> destStore,
     MachineReservation::ptr & reservation,
     ::Machine::ptr machine, Step::ptr step,
     const BuildOptions & buildOptions,
@@ -578,7 +584,7 @@ void State::buildRemote(ref<Store> destStore,
            copy the immediate sources of the derivation and the required
            outputs of the input derivations. */
         updateStep(ssSendingInputs);
-        BasicDerivation resolvedDrv = build_remote::sendInputs(*this, *step, *localStore, *destStore, conn, result.overhead, nrStepsWaiting, nrStepsCopyingTo);
+        BasicDerivation resolvedDrv = build_remote::sendInputs(aio, *this, *step, *localStore, *destStore, conn, result.overhead, nrStepsWaiting, nrStepsCopyingTo);
 
         logFileDel.cancel();
 
@@ -665,7 +671,7 @@ void State::buildRemote(ref<Store> destStore,
             printMsg(lvlDebug, "copying outputs of ‘%s’ from ‘%s’ (%d bytes)",
                 localStore->printStorePath(step->drvPath), machine->sshName, totalNarSize);
 
-            build_remote::copyPathsFromRemote(conn, narMembers, *localStore, *destStore, infos);
+            build_remote::copyPathsFromRemote(conn, narMembers, *localStore, *destStore, infos, aio);
             auto now2 = std::chrono::steady_clock::now();
 
             result.overhead += std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();
diff --git a/src/hydra-queue-runner/build-result.cc b/src/hydra-queue-runner/build-result.cc
index 2672c4df3..051784471 100644
--- a/src/hydra-queue-runner/build-result.cc
+++ b/src/hydra-queue-runner/build-result.cc
@@ -9,6 +9,7 @@ using namespace nix;
 
 
 BuildOutput getBuildOutput(
+    AsyncIoRoot & aio,
     nix::ref<Store> store,
     NarMemberDatas & narMembers,
     const OutputPathMap derivationOutputs)
@@ -19,7 +20,7 @@ BuildOutput getBuildOutput(
     StorePathSet outputs;
     StorePathSet closure;
     for (auto& [outputName, outputPath] : derivationOutputs) {
-        store->computeFSClosure(outputPath, closure);
+        aio.blockOn(store->computeFSClosure(outputPath, closure));
         outputs.insert(outputPath);
         res.outputs.insert({outputName, outputPath});
     }
diff --git a/src/hydra-queue-runner/builder.cc b/src/hydra-queue-runner/builder.cc
index f3bbba4d8..a4a814e8f 100644
--- a/src/hydra-queue-runner/builder.cc
+++ b/src/hydra-queue-runner/builder.cc
@@ -19,6 +19,8 @@ void setThreadName(const std::string & name)
 
 void State::builder(MachineReservation::ptr reservation)
 {
+    AsyncIoRoot aio;
+
     setThreadName("bld~" + std::string(reservation->step->drvPath.to_string()));
 
     StepResult res = sRetry;
@@ -41,7 +43,7 @@ void State::builder(MachineReservation::ptr reservation)
         try {
             auto destStore = getDestStore();
             // Might release the reservation.
-            res = doBuildStep(destStore, reservation, *conn, activeStep);
+            res = doBuildStep(aio, destStore, reservation, *conn, activeStep);
         } catch (pqxx::broken_connection & e) {
             printMsg(lvlError, "db lost while building ‘%s’ on ‘%s’: %s (retriable)",
                 localStore->printStorePath(activeStep->step->drvPath),
@@ -83,7 +85,8 @@ void State::builder(MachineReservation::ptr reservation)
 }
 
 
-State::StepResult State::doBuildStep(nix::ref<Store> destStore,
+State::StepResult State::doBuildStep(AsyncIoRoot & aio,
+    nix::ref<Store> destStore,
     MachineReservation::ptr & reservation,
     Connection & conn,
     std::shared_ptr<ActiveStep> activeStep)
@@ -204,7 +207,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
         {
             auto mc = startDbUpdate();
             pqxx::work txn(conn);
-            stepNr = createBuildStep(txn, result.startTime, buildId, step, machine->sshName, bsBusy);
+            stepNr = createBuildStep(aio, txn, result.startTime, buildId, step, machine->sshName, bsBusy);
             txn.commit();
         }
 
@@ -219,7 +222,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
 
         try {
             /* FIXME: referring builds may have conflicting timeouts. */
-            buildRemote(destStore, reservation, machine, step, buildOptions, result, activeStep, updateStep, narMembers);
+            buildRemote(aio, destStore, reservation, machine, step, buildOptions, result, activeStep, updateStep, narMembers);
         } catch (Error & e) {
             if (activeStep->state_.lock()->cancelled) {
                 printInfo("marking step %d of build %d as cancelled", stepNr, buildId);
@@ -234,7 +237,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
 
         if (result.stepStatus == bsSuccess) {
             updateStep(ssPostProcessing);
-            res = getBuildOutput(destStore, narMembers, destStore->queryDerivationOutputMap(step->drvPath, &*localStore));
+            res = getBuildOutput(aio, destStore, narMembers, aio.blockOn(destStore->queryDerivationOutputMap(step->drvPath, &*localStore)));
         }
     }
 
@@ -261,7 +264,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
     /* Finish the step in the database. */
     if (stepNr) {
         pqxx::work txn(conn);
-        finishBuildStep(txn, result, buildId, stepNr, machine->sshName);
+        finishBuildStep(aio, txn, result, buildId, stepNr, machine->sshName);
         txn.commit();
     }
 
@@ -288,7 +291,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
 
         assert(stepNr);
 
-        for (auto & [outputName, optOutputPath] : destStore->queryPartialDerivationOutputMap(step->drvPath, &*localStore)) {
+        for (auto & [outputName, optOutputPath] : aio.blockOn(destStore->queryPartialDerivationOutputMap(step->drvPath, &*localStore))) {
             if (!optOutputPath)
                 throw Error(
                     "Missing output %s for derivation %d which was supposed to have succeeded",
@@ -393,7 +396,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
         }
 
     } else
-        failStep(conn, step, buildId, result, machine, stepFinished);
+        failStep(aio, conn, step, buildId, result, machine, stepFinished);
 
     // FIXME: keep stats about aborted steps?
     nrStepsDone++;
@@ -410,6 +413,7 @@ State::StepResult State::doBuildStep(nix::ref<Store> destStore,
 
 
 void State::failStep(
+    AsyncIoRoot & aio,
     Connection & conn,
     Step::ptr step,
     BuildID buildId,
@@ -459,7 +463,7 @@ void State::failStep(
                     ((result.stepStatus != bsCachedFailure && result.stepStatus != bsUnsupported) && buildId == build->id) ||
                     build->finishedInDB)
                     continue;
-                createBuildStep(txn,
+                createBuildStep(aio, txn,
                     0, build->id, step, machine ? machine->sshName : "",
                     result.stepStatus, result.errorMsg, buildId == build->id ? 0 : buildId);
             }
diff --git a/src/hydra-queue-runner/dispatcher.cc b/src/hydra-queue-runner/dispatcher.cc
index 050e20ad3..f744d3ba2 100644
--- a/src/hydra-queue-runner/dispatcher.cc
+++ b/src/hydra-queue-runner/dispatcher.cc
@@ -34,6 +34,8 @@ void State::dispatcher()
     printMsg(lvlDebug, "Waiting for the machines parsing to have completed at least once");
     machinesReadyLock.lock();
 
+    AsyncIoRoot aio;
+
     while (true) {
         try {
             printMsg(lvlDebug, "dispatcher woken up");
@@ -41,7 +43,7 @@ void State::dispatcher()
 
             auto t_before_work = std::chrono::steady_clock::now();
 
-            auto sleepUntil = doDispatch();
+            auto sleepUntil = doDispatch(aio);
 
             auto t_after_work = std::chrono::steady_clock::now();
 
@@ -76,7 +78,7 @@ void State::dispatcher()
 }
 
 
-system_time State::doDispatch()
+system_time State::doDispatch(AsyncIoRoot & aio)
 {
     /* Prune old historical build step info from the jobsets. */
     {
@@ -321,7 +323,7 @@ system_time State::doDispatch()
 
     } while (keepGoing);
 
-    abortUnsupported();
+    abortUnsupported(aio);
 
     return sleepUntil;
 }
@@ -337,7 +339,7 @@ void State::wakeDispatcher()
 }
 
 
-void State::abortUnsupported()
+void State::abortUnsupported(AsyncIoRoot & aio)
 {
     /* Make a copy of 'runnable' and 'machines' so we don't block them
        very long. */
@@ -397,6 +399,7 @@ void State::abortUnsupported()
             bool stepFinished = false;
 
             failStep(
+                aio,
                 *conn, step, build->id,
                 RemoteResult {
                     .stepStatus = bsUnsupported,
diff --git a/src/hydra-queue-runner/hydra-build-result.hh b/src/hydra-queue-runner/hydra-build-result.hh
index 9e92bdf8d..cb7079f9c 100644
--- a/src/hydra-queue-runner/hydra-build-result.hh
+++ b/src/hydra-queue-runner/hydra-build-result.hh
@@ -42,6 +42,7 @@ struct BuildOutput
 };
 
 BuildOutput getBuildOutput(
+    nix::AsyncIoRoot & aio,
     nix::ref<nix::Store> store,
     NarMemberDatas & narMembers,
     const nix::OutputPathMap derivationOutputs);
diff --git a/src/hydra-queue-runner/hydra-queue-runner.cc b/src/hydra-queue-runner/hydra-queue-runner.cc
index 7d788a3b1..c59b0216d 100644
--- a/src/hydra-queue-runner/hydra-queue-runner.cc
+++ b/src/hydra-queue-runner/hydra-queue-runner.cc
@@ -1,4 +1,5 @@
 #include <iostream>
+#include <lix/libstore/pathlocks.hh>
 #include <thread>
 #include <optional>
 #include <type_traits>
@@ -20,7 +21,6 @@
 #include "lix/libmain/shared.hh"
 #include "lix/libstore/globals.hh"
 #include "lix/libstore/s3-binary-cache-store.hh"
-#include "lix/libutil/args.hh"
 #include "lix/libutil/async.hh"
 
 using namespace nix;
@@ -333,7 +333,7 @@ unsigned int State::allocBuildStep(pqxx::work & txn, BuildID buildId)
 }
 
 
-unsigned int State::createBuildStep(pqxx::work & txn, time_t startTime, BuildID buildId, Step::ptr step,
+unsigned int State::createBuildStep(AsyncIoRoot & aio, pqxx::work & txn, time_t startTime, BuildID buildId, Step::ptr step,
     const std::string & machine, BuildStatus status, const std::string & errorMsg, BuildID propagatedFrom)
 {
  restart:
@@ -356,7 +356,7 @@ unsigned int State::createBuildStep(pqxx::work & txn, time_t startTime, BuildID
 
     if (r.affected_rows() == 0) goto restart;
 
-    for (auto & [name, output] : getDestStore()->queryPartialDerivationOutputMap(step->drvPath, &*localStore))
+    for (auto & [name, output] : aio.blockOn(getDestStore()->queryPartialDerivationOutputMap(step->drvPath, &*localStore)))
         txn.exec_params0
             ("insert into BuildStepOutputs (build, stepnr, name, path) values ($1, $2, $3, $4)",
             buildId, stepNr, name,
@@ -382,7 +382,7 @@ void State::updateBuildStep(pqxx::work & txn, BuildID buildId, unsigned int step
 }
 
 
-void State::finishBuildStep(pqxx::work & txn, const RemoteResult & result,
+void State::finishBuildStep(AsyncIoRoot & aio, pqxx::work & txn, const RemoteResult & result,
     BuildID buildId, unsigned int stepNr, const std::string & machine)
 {
     assert(result.startTime);
@@ -406,7 +406,7 @@ void State::finishBuildStep(pqxx::work & txn, const RemoteResult & result,
         assert(res.size());
         StorePath drvPath = localStore->parseStorePath(res[0].as<std::string>());
         // If we've finished building, all the paths should be known
-        for (auto & [name, output] : getDestStore()->queryDerivationOutputMap(drvPath, &*localStore))
+        for (auto & [name, output] : aio.blockOn(getDestStore()->queryDerivationOutputMap(drvPath, &*localStore)))
             txn.exec_params0
                 ("update BuildStepOutputs set path = $4 where build = $1 and stepnr = $2 and name = $3",
                   buildId, stepNr, name, localStore->printStorePath(output));
@@ -589,16 +589,13 @@ void State::notifyBuildFinished(pqxx::work & txn, BuildID buildId,
 }
 
 
-std::shared_ptr<PathLocks> State::acquireGlobalLock()
+std::optional<PathLock> State::acquireGlobalLock()
 {
     Path lockPath = hydraData + "/queue-runner/lock";
 
     createDirs(dirOf(lockPath));
 
-    auto lock = std::make_shared<PathLocks>();
-    if (!lock->tryLockPaths(PathSet({lockPath}))) return nullptr;
-
-    return lock;
+    return nix::tryLockPath(lockPath);
 }
 
 
@@ -846,7 +843,7 @@ void State::unlock()
 }
 
 
-void State::run(BuildID buildOne)
+void State::run(AsyncIoRoot & aio, BuildID buildOne)
 {
     /* Can't be bothered to shut down cleanly. Goodbye! */
     auto callback = createInterruptCallback([&]() { std::_Exit(0); });
@@ -873,10 +870,10 @@ void State::run(BuildID buildOne)
     StoreConfig::Params localParams;
     localParams["max-connections"] = "16";
     localParams["max-connection-age"] = "600";
-    localStore = openStore(getEnv("NIX_REMOTE").value_or(""), localParams);
+    localStore = aio.blockOn(openStore(getEnv("NIX_REMOTE").value_or(""), localParams));
 
     auto storeUri = config->getStrOption("store_uri");
-    _destStore = storeUri == "" ? localStore : openStore(storeUri);
+    _destStore = storeUri == "" ? localStore : aio.blockOn(openStore(storeUri));
 
     useSubstitutes = config->getBoolOption("use-substitutes", false);
 
@@ -975,6 +972,7 @@ int main(int argc, char * * argv)
         BuildID buildOne = 0;
         std::optional<std::string> metricsAddrOpt = std::nullopt;
 
+        // By definition the `aio` for the main thread.
         nix::AsyncIoRoot aio;
         LegacyArgs(aio, argv[0], [&](Strings::iterator & arg, const Strings::iterator & end) {
             if (*arg == "--unlock")
@@ -1001,6 +999,6 @@ int main(int argc, char * * argv)
         else if (unlock)
             state.unlock();
         else
-            state.run(buildOne);
+            state.run(aio, buildOne);
     });
 }
diff --git a/src/hydra-queue-runner/nar-extractor.cc b/src/hydra-queue-runner/nar-extractor.cc
index 0378a9b81..329a63435 100644
--- a/src/hydra-queue-runner/nar-extractor.cc
+++ b/src/hydra-queue-runner/nar-extractor.cc
@@ -2,77 +2,112 @@
 
 #include "lix/libutil/archive.hh"
 
+#include <lix/libutil/box_ptr.hh>
+#include <lix/libutil/serialise.hh>
+#include <lix/libutil/types.hh>
 #include <unordered_set>
 #include <utility>
 
 using namespace nix;
 
-struct Extractor : NARParseVisitor
+static Generator<Bytes> readHydraData(
+    NarMemberData &narData,
+    Generator<Bytes> contents,
+    uint64_t size)
 {
-    class MyFileHandle : public FileHandle
-    {
-        NarMemberData & memberData;
-        uint64_t expectedSize;
-        std::unique_ptr<HashSink> hashSink;
+    HashSink hashSink(HashType::SHA256);
+    while (auto block = contents.next()) {
+        std::string_view data = {block->data(), block->size()};
 
-        public:
-            MyFileHandle(NarMemberData & memberData, uint64_t size) : memberData(memberData), expectedSize(size)
-            {
-                hashSink = std::make_unique<HashSink>(HashType::SHA256);
-            }
+        *narData.fileSize += block->size();
+        hashSink(data);
+        if (narData.contents) {
+            narData.contents->append(data);
+        }
+        assert(narData.fileSize <= size);
+        if (narData.fileSize == size) {
+            auto [hash, len] = hashSink.finish();
+            assert(narData.fileSize == len);
+            narData.sha256 = hash;
+        }
 
-            void receiveContents(std::string_view data) override
-            {
-                *memberData.fileSize += data.size();
-                (*hashSink)(data);
-                if (memberData.contents) {
-                    memberData.contents->append(data);
-                }
-                assert(memberData.fileSize <= expectedSize);
-                if (memberData.fileSize == expectedSize) {
-                    auto [hash, len] = hashSink->finish();
-                    assert(memberData.fileSize == len);
-                    memberData.sha256 = hash;
-                    hashSink.reset();
-                }
-            }
-    };
+        co_yield *block;
+    }
+}
 
-    std::unordered_set<Path> filesToKeep {
+static Generator<nar::Entry> discoverHydraDataInNAR(
+    NarMemberDatas &members,
+    const Path &prefix,
+    const Path pathSuffix,
+    nar::Entry root)
+{
+    static const std::unordered_set<Path> filesToKeep{
         "/nix-support/hydra-build-products",
         "/nix-support/hydra-release-name",
         "/nix-support/hydra-metrics",
     };
 
-    NarMemberDatas & members;
-    Path prefix;
+    nix::overloaded handlers {
+        [&](nar::File f) -> nar::Entry {
+            return nar::File{
+                f.executable,
+                f.size,
+                [](
+                    nar::File f,
+                    NarMemberDatas &members,
+                    const Path &pathSuffix,
+                    const Path &prefix
+                ) -> Generator<Bytes> {
+                    auto pref = prefix + pathSuffix;
+                    auto memberData = NarMemberData{
+                        .type = FSAccessor::Type::tRegular,
+                        .fileSize = 0,
+                        .contents = filesToKeep.count(pathSuffix)
+                            ? std::optional("")
+                            : std::nullopt,
+                    };
+                    co_yield readHydraData(memberData, std::move(f.contents), f.size);
+                    members.insert_or_assign(pref, memberData);
+                }(std::move(f), members, pathSuffix, prefix),
+            };
+        },
+        [&](nar::Symlink sl) -> nar::Entry {
+            members.insert_or_assign(
+                prefix + pathSuffix,
+                NarMemberData{.type = FSAccessor::Type::tSymlink});
+            return sl;
+        },
+        [&](nar::Directory d) -> nar::Entry {
+            auto pref = prefix + pathSuffix;
+            members.insert_or_assign(
+                pref, NarMemberData{.type = FSAccessor::Type::tDirectory});
 
-    Extractor(NarMemberDatas & members, Path  prefix)
-        : members(members), prefix(std::move(prefix))
-    { }
-
-    void createDirectory(const Path & path) override
-    {
-        members.insert_or_assign(prefix + path, NarMemberData { .type = FSAccessor::Type::tDirectory });
-    }
-
-    std::unique_ptr<FileHandle> createRegularFile(const Path & path, uint64_t size, bool executable) override
-    {
-        auto memberData = &members.insert_or_assign(prefix + path, NarMemberData {
-            .type = FSAccessor::Type::tRegular,
-            .fileSize = 0,
-            .contents = filesToKeep.count(path) ? std::optional("") : std::nullopt,
-        }).first->second;
-
-        return std::make_unique<MyFileHandle>(*memberData, size);
-    }
-
-    void createSymlink(const Path & path, const std::string & target) override
-    {
-        members.insert_or_assign(prefix + path, NarMemberData { .type = FSAccessor::Type::tSymlink });
-    }
-};
+            return nar::Directory{
+                ([](
+                    NarMemberDatas &members,
+                    nar::Directory d,
+                    const Path &prefix,
+                    const Path &pathSuffix
+                ) -> Generator<std::pair<const std::string&, nar::Entry>> {
+                    while (auto e = d.contents.next()) {
+                        auto iter = discoverHydraDataInNAR(
+                            members,
+                            prefix,
+                            pathSuffix + "/" + e->first,
+                            std::move(e->second)
+                        );
+                        co_yield std::pair{
+                            std::cref(e->first),
+                            *iter.next(),
+                        };
+                        assert(!iter.next().has_value());
+                    }
+                }(members, std::move(d), prefix, pathSuffix))};
+        }
+    };
 
+    co_yield std::visit(handlers, std::move(root));
+}
 
 void extractNarData(
     box_ptr<Source> & source,
@@ -91,7 +126,15 @@ nix::WireFormatGenerator extractNarDataFilter(
     NarMemberDatas & members)
 {
     return [](Source & source, const Path & prefix, NarMemberDatas & members) -> WireFormatGenerator {
-        Extractor extractor(members, prefix);
-        co_yield parseAndCopyDump(extractor, source);
+        auto items = nar::parse(source);
+        auto inner = discoverHydraDataInNAR(
+            members,
+            prefix,
+            "",
+            std::move(*items.next())
+        );
+        co_yield nar::dump(*inner.next());
+        assert(!items.next().has_value());
+        assert(!inner.next().has_value());
     }(source, prefix, members);
 }
diff --git a/src/hydra-queue-runner/queue-monitor.cc b/src/hydra-queue-runner/queue-monitor.cc
index 6231449bc..04900a2db 100644
--- a/src/hydra-queue-runner/queue-monitor.cc
+++ b/src/hydra-queue-runner/queue-monitor.cc
@@ -202,7 +202,7 @@ bool State::getQueuedBuilds(AsyncIoRoot & aio, Connection & conn,
                 if (!res[0].is_null()) propagatedFrom = res[0].as<BuildID>();
 
                 if (!propagatedFrom) {
-                    for (auto & [outputName, optOutputPath] : destStore->queryPartialDerivationOutputMap(ex.step->drvPath, &*localStore)) {
+                    for (auto & [outputName, optOutputPath] : aio.blockOn(destStore->queryPartialDerivationOutputMap(ex.step->drvPath, &*localStore))) {
                         constexpr std::string_view common = "select max(s.build) from BuildSteps s join BuildStepOutputs o on s.build = o.build where startTime != 0 and stopTime != 0 and status = 1";
                         auto res = optOutputPath
                             ? txn.exec_params(
@@ -219,7 +219,7 @@ bool State::getQueuedBuilds(AsyncIoRoot & aio, Connection & conn,
                     }
                 }
 
-                createBuildStep(txn, 0, build->id, ex.step, "", bsCachedFailure, "", propagatedFrom);
+                createBuildStep(aio, txn, 0, build->id, ex.step, "", bsCachedFailure, "", propagatedFrom);
                 txn.exec_params
                     ("update Builds set finished = 1, buildStatus = $2, startTime = $3, stopTime = $3, isCachedBuild = 1, notificationPendingSince = $3 "
                      "where id = $1 and finished = 0",
@@ -250,9 +250,9 @@ bool State::getQueuedBuilds(AsyncIoRoot & aio, Connection & conn,
         /* If we didn't get a step, it means the step's outputs are
            all valid. So we mark this as a finished, cached build. */
         if (!step) {
-            BuildOutput res = getBuildOutputCached(conn, destStore, build->drvPath);
+            BuildOutput res = getBuildOutputCached(aio, conn, destStore, build->drvPath);
 
-            for (auto & i : destStore->queryDerivationOutputMap(build->drvPath, &*localStore))
+            for (auto & i : aio.blockOn(destStore->queryDerivationOutputMap(build->drvPath, &*localStore)))
                 addRoot(i.second);
 
             {
@@ -518,7 +518,7 @@ Step::ptr State::createStep(AsyncIoRoot & aio, ref<Store> destStore,
     /* Are all outputs valid? */
     auto outputHashes = staticOutputHashes(*localStore, *(step->drv));
     std::map<DrvOutput, std::optional<StorePath>> paths;
-    for (auto & [outputName, maybeOutputPath] : destStore->queryPartialDerivationOutputMap(drvPath, &*localStore)) {
+    for (auto & [outputName, maybeOutputPath] : aio.blockOn(destStore->queryPartialDerivationOutputMap(drvPath, &*localStore))) {
         auto outputHash = outputHashes.at(outputName);
         paths.insert({{.drvHash=outputHash, .outputName=outputName}, maybeOutputPath});
     }
@@ -549,7 +549,7 @@ Step::ptr State::createStep(AsyncIoRoot & aio, ref<Store> destStore,
                 avail++;
             else if (useSubstitutes) {
                 SubstitutablePathInfos infos;
-                localStore->querySubstitutablePathInfos({{path, {}}}, infos);
+                aio.blockOn(localStore->querySubstitutablePathInfos({{path, {}}}, infos));
                 if (infos.size() == 1)
                     avail++;
             }
@@ -578,9 +578,9 @@ Step::ptr State::createStep(AsyncIoRoot & aio, ref<Store> destStore,
                         // FIXME: should copy directly from substituter to destStore.
                     }
 
-                    copyClosure(*localStore, *destStore,
+                    aio.blockOn(copyClosure(*localStore, *destStore,
                         StorePathSet { path },
-                        NoRepair, CheckSigs, NoSubstitute);
+                        NoRepair, CheckSigs, NoSubstitute));
 
                     time_t stopTime = time(nullptr);
 
@@ -692,9 +692,9 @@ void State::processJobsetSharesChange(Connection & conn)
 }
 
 
-BuildOutput State::getBuildOutputCached(Connection & conn, nix::ref<nix::Store> destStore, const nix::StorePath & drvPath)
+BuildOutput State::getBuildOutputCached(nix::AsyncIoRoot & aio, Connection & conn, nix::ref<nix::Store> destStore, const nix::StorePath & drvPath)
 {
-    auto derivationOutputs = destStore->queryDerivationOutputMap(drvPath, &*localStore);
+    auto derivationOutputs = aio.blockOn(destStore->queryDerivationOutputMap(drvPath, &*localStore));
 
     {
     pqxx::work txn(conn);
@@ -758,5 +758,5 @@ BuildOutput State::getBuildOutputCached(Connection & conn, nix::ref<nix::Store>
     }
 
     NarMemberDatas narMembers;
-    return getBuildOutput(destStore, narMembers, derivationOutputs);
+    return getBuildOutput(aio, destStore, narMembers, derivationOutputs);
 }
diff --git a/src/hydra-queue-runner/state.hh b/src/hydra-queue-runner/state.hh
index 9e5f21b6f..16f2f2ca6 100644
--- a/src/hydra-queue-runner/state.hh
+++ b/src/hydra-queue-runner/state.hh
@@ -523,13 +523,13 @@ private:
 
     unsigned int allocBuildStep(pqxx::work & txn, BuildID buildId);
 
-    unsigned int createBuildStep(pqxx::work & txn, time_t startTime, BuildID buildId, Step::ptr step,
+    unsigned int createBuildStep(nix::AsyncIoRoot & aio, pqxx::work & txn, time_t startTime, BuildID buildId, Step::ptr step,
         const std::string & machine, BuildStatus status, const std::string & errorMsg = "",
         BuildID propagatedFrom = 0);
 
     void updateBuildStep(pqxx::work & txn, BuildID buildId, unsigned int stepNr, StepState stepState);
 
-    void finishBuildStep(pqxx::work & txn, const RemoteResult & result, BuildID buildId, unsigned int stepNr,
+    void finishBuildStep(nix::AsyncIoRoot & aio, pqxx::work & txn, const RemoteResult & result, BuildID buildId, unsigned int stepNr,
         const std::string & machine);
 
     unsigned int createSubstitutionStep(pqxx::work & txn, time_t startTime, time_t stopTime,
@@ -547,7 +547,7 @@ private:
     /* Handle cancellation, deletion and priority bumps. */
     void processQueueChange(Connection & conn);
 
-    BuildOutput getBuildOutputCached(Connection & conn, nix::ref<nix::Store> destStore,
+    BuildOutput getBuildOutputCached(nix::AsyncIoRoot & aio, Connection & conn, nix::ref<nix::Store> destStore,
         const nix::StorePath & drvPath);
 
     /* Returns paths missing from the remote store. Paths are processed in
@@ -562,6 +562,7 @@ private:
         std::set<Step::ptr> & newSteps, std::set<Step::ptr> & newRunnable);
 
     void failStep(
+        nix::AsyncIoRoot & aio,
         Connection & conn,
         Step::ptr step,
         BuildID buildId,
@@ -579,23 +580,25 @@ private:
     /* The thread that selects and starts runnable builds. */
     void dispatcher();
 
-    system_time doDispatch();
+    system_time doDispatch(nix::AsyncIoRoot & aio);
 
     void wakeDispatcher();
 
-    void abortUnsupported();
+    void abortUnsupported(nix::AsyncIoRoot & aio);
 
     void builder(MachineReservation::ptr reservation);
 
     /* Perform the given build step. Return true if the step is to be
        retried. */
     enum StepResult { sDone, sRetry, sMaybeCancelled };
-    StepResult doBuildStep(nix::ref<nix::Store> destStore,
+    StepResult doBuildStep(nix::AsyncIoRoot & aio,
+        nix::ref<nix::Store> destStore,
         MachineReservation::ptr & reservation,
         Connection & conn,
         std::shared_ptr<ActiveStep> activeStep);
 
-    void buildRemote(nix::ref<nix::Store> destStore,
+    void buildRemote(nix::AsyncIoRoot & aio,
+        nix::ref<nix::Store> destStore,
         MachineReservation::ptr & reservation,
         Machine::ptr machine, Step::ptr step,
         const BuildOptions & buildOptions,
@@ -615,7 +618,7 @@ private:
 
     /* Acquire the global queue runner lock, or null if somebody else
        has it. */
-    std::shared_ptr<nix::PathLocks> acquireGlobalLock();
+    std::optional<nix::PathLock> acquireGlobalLock();
 
     void dumpStatus(Connection & conn);
 
@@ -627,5 +630,5 @@ public:
 
     void unlock();
 
-    void run(BuildID buildOne = 0);
+    void run(nix::AsyncIoRoot & aio, BuildID buildOne = 0);
 };