7 files changed, 133 insertions, 106 deletions

diff --git a/src/libstore/build/derivation-goal.cc b/src/libstore/build/derivation-goal.cc
index 96140e10b..765df5f5a 100644
--- a/src/libstore/build/derivation-goal.cc
+++ b/src/libstore/build/derivation-goal.cc
@@ -127,7 +127,15 @@ Goal::WorkResult DerivationGoal::timedOut(Error && ex)
 
 kj::Promise<Result<Goal::WorkResult>> DerivationGoal::workImpl() noexcept
 {
-    return useDerivation ? getDerivation() : haveDerivation();
+    KJ_DEFER({
+        act.reset();
+        actLock.reset();
+        builderActivities.clear();
+    });
+
+    BOOST_OUTCOME_CO_TRY(auto result, co_await (useDerivation ? getDerivation() : haveDerivation()));
+    result.storePath = drvPath;
+    co_return result;
 }
 
 bool DerivationGoal::addWantedOutputs(const OutputsSpec & outputs)
diff --git a/src/libstore/build/entry-points.cc b/src/libstore/build/entry-points.cc
index 808179a4d..36ad35be0 100644
--- a/src/libstore/build/entry-points.cc
+++ b/src/libstore/build/entry-points.cc
@@ -6,26 +6,20 @@
 
 namespace nix {
 
-static auto runWorker(Worker & worker, auto mkGoals)
-{
-    return worker.run(mkGoals);
-}
-
 void Store::buildPaths(const std::vector<DerivedPath> & reqs, BuildMode buildMode, std::shared_ptr<Store> evalStore)
 {
     auto aio = kj::setupAsyncIo();
-    Worker worker(*this, evalStore ? *evalStore : *this, aio);
 
-    auto goals = runWorker(worker, [&](GoalFactory & gf) {
+    auto results = processGoals(*this, evalStore ? *evalStore : *this, aio, [&](GoalFactory & gf) {
         Worker::Targets goals;
         for (auto & br : reqs)
-            goals.emplace(gf.makeGoal(br, buildMode));
+            goals.emplace_back(gf.makeGoal(br, buildMode));
         return goals;
-    });
+    }).wait(aio.waitScope).value();
 
     StringSet failed;
     std::shared_ptr<Error> ex;
-    for (auto & [i, result] : goals) {
+    for (auto & [i, result] : results.goals) {
         if (result.ex) {
             if (ex)
                 logError(result.ex->info());
@@ -33,19 +27,17 @@ void Store::buildPaths(const std::vector<DerivedPath> & reqs, BuildMode buildMod
                 ex = result.ex;
         }
         if (result.exitCode != Goal::ecSuccess) {
-            if (auto i2 = dynamic_cast<DerivationGoal *>(i.get()))
-                failed.insert(printStorePath(i2->drvPath));
-            else if (auto i2 = dynamic_cast<PathSubstitutionGoal *>(i.get()))
-                failed.insert(printStorePath(i2->storePath));
+            if (result.storePath)
+                failed.insert(printStorePath(*result.storePath));
         }
     }
 
     if (failed.size() == 1 && ex) {
-        ex->withExitStatus(worker.failingExitStatus());
+        ex->withExitStatus(results.failingExitStatus);
         throw std::move(*ex);
     } else if (!failed.empty()) {
         if (ex) logError(ex->info());
-        throw Error(worker.failingExitStatus(), "build of %s failed", concatStringsSep(", ", quoteStrings(failed)));
+        throw Error(results.failingExitStatus, "build of %s failed", concatStringsSep(", ", quoteStrings(failed)));
     }
 }
 
@@ -55,24 +47,19 @@ std::vector<KeyedBuildResult> Store::buildPathsWithResults(
     std::shared_ptr<Store> evalStore)
 {
     auto aio = kj::setupAsyncIo();
-    Worker worker(*this, evalStore ? *evalStore : *this, aio);
 
-    std::vector<std::pair<const DerivedPath &, GoalPtr>> state;
-
-    auto goals = runWorker(worker, [&](GoalFactory & gf) {
+    auto goals = processGoals(*this, evalStore ? *evalStore : *this, aio, [&](GoalFactory & gf) {
         Worker::Targets goals;
         for (const auto & req : reqs) {
-            auto goal = gf.makeGoal(req, buildMode);
-            state.push_back({req, goal.first});
-            goals.emplace(std::move(goal));
+            goals.emplace_back(gf.makeGoal(req, buildMode));
         }
         return goals;
-    });
+    }).wait(aio.waitScope).value().goals;
 
     std::vector<KeyedBuildResult> results;
 
-    for (auto & [req, goalPtr] : state)
-        results.emplace_back(goals[goalPtr].result.restrictTo(req));
+    for (auto && [goalIdx, req] : enumerate(reqs))
+        results.emplace_back(goals[goalIdx].result.restrictTo(req));
 
     return results;
 }
@@ -81,15 +68,14 @@ BuildResult Store::buildDerivation(const StorePath & drvPath, const BasicDerivat
     BuildMode buildMode)
 {
     auto aio = kj::setupAsyncIo();
-    Worker worker(*this, *this, aio);
 
     try {
-        auto goals = runWorker(worker, [&](GoalFactory & gf) {
+        auto results = processGoals(*this, *this, aio, [&](GoalFactory & gf) {
             Worker::Targets goals;
-            goals.emplace(gf.makeBasicDerivationGoal(drvPath, drv, OutputsSpec::All{}, buildMode));
+            goals.emplace_back(gf.makeBasicDerivationGoal(drvPath, drv, OutputsSpec::All{}, buildMode));
             return goals;
-        });
-        auto [goal, result] = *goals.begin();
+        }).wait(aio.waitScope).value();
+        auto & result = results.goals.begin()->second;
         return result.result.restrictTo(DerivedPath::Built {
             .drvPath = makeConstantStorePathRef(drvPath),
             .outputs = OutputsSpec::All {},
@@ -109,21 +95,20 @@ void Store::ensurePath(const StorePath & path)
     if (isValidPath(path)) return;
 
     auto aio = kj::setupAsyncIo();
-    Worker worker(*this, *this, aio);
 
-    auto goals = runWorker(worker, [&](GoalFactory & gf) {
+    auto results = processGoals(*this, *this, aio, [&](GoalFactory & gf) {
         Worker::Targets goals;
-        goals.emplace(gf.makePathSubstitutionGoal(path));
+        goals.emplace_back(gf.makePathSubstitutionGoal(path));
         return goals;
-    });
-    auto [goal, result] = *goals.begin();
+    }).wait(aio.waitScope).value();
+    auto & result = results.goals.begin()->second;
 
     if (result.exitCode != Goal::ecSuccess) {
         if (result.ex) {
-            result.ex->withExitStatus(worker.failingExitStatus());
+            result.ex->withExitStatus(results.failingExitStatus);
             throw std::move(*result.ex);
         } else
-            throw Error(worker.failingExitStatus(), "path '%s' does not exist and cannot be created", printStorePath(path));
+            throw Error(results.failingExitStatus, "path '%s' does not exist and cannot be created", printStorePath(path));
     }
 }
 
@@ -131,23 +116,22 @@ void Store::ensurePath(const StorePath & path)
 void Store::repairPath(const StorePath & path)
 {
     auto aio = kj::setupAsyncIo();
-    Worker worker(*this, *this, aio);
 
-    auto goals = runWorker(worker, [&](GoalFactory & gf) {
+    auto results = processGoals(*this, *this, aio, [&](GoalFactory & gf) {
         Worker::Targets goals;
-        goals.emplace(gf.makePathSubstitutionGoal(path, Repair));
+        goals.emplace_back(gf.makePathSubstitutionGoal(path, Repair));
         return goals;
-    });
-    auto [goal, result] = *goals.begin();
+    }).wait(aio.waitScope).value();
+    auto & result = results.goals.begin()->second;
 
     if (result.exitCode != Goal::ecSuccess) {
         /* Since substituting the path didn't work, if we have a valid
            deriver, then rebuild the deriver. */
         auto info = queryPathInfo(path);
         if (info->deriver && isValidPath(*info->deriver)) {
-            worker.run([&](GoalFactory & gf) {
+            processGoals(*this, *this, aio, [&](GoalFactory & gf) {
                 Worker::Targets goals;
-                goals.emplace(gf.makeGoal(
+                goals.emplace_back(gf.makeGoal(
                     DerivedPath::Built{
                         .drvPath = makeConstantStorePathRef(*info->deriver),
                         // FIXME: Should just build the specific output we need.
@@ -156,9 +140,9 @@ void Store::repairPath(const StorePath & path)
                     bmRepair
                 ));
                 return goals;
-            });
+            }).wait(aio.waitScope).value();
         } else
-            throw Error(worker.failingExitStatus(), "cannot repair path '%s'", printStorePath(path));
+            throw Error(results.failingExitStatus, "cannot repair path '%s'", printStorePath(path));
     }
 }
 
diff --git a/src/libstore/build/goal.cc b/src/libstore/build/goal.cc
index 02b22b8ad..db16e2cf8 100644
--- a/src/libstore/build/goal.cc
+++ b/src/libstore/build/goal.cc
@@ -22,6 +22,12 @@ kj::Promise<void> Goal::waitForAWhile()
 
 kj::Promise<Result<Goal::WorkResult>> Goal::work() noexcept
 try {
+    // always clear the slot token, no matter what happens. not doing this
+    // can cause builds to get stuck on exceptions (or other early exist).
+    // ideally we'd use scoped slot tokens instead of keeping them in some
+    // goal member variable, but we cannot do this yet for legacy reasons.
+    KJ_DEFER({ slotToken = {}; });
+
     BOOST_OUTCOME_CO_TRY(auto result, co_await workImpl());
 
     trace("done");
diff --git a/src/libstore/build/goal.hh b/src/libstore/build/goal.hh
index 29540dcd3..5b66a2c08 100644
--- a/src/libstore/build/goal.hh
+++ b/src/libstore/build/goal.hh
@@ -94,6 +94,9 @@ public:
         bool timedOut = false;
         bool hashMismatch = false;
         bool checkMismatch = false;
+        /// Store path this goal relates to. Will be set to drvPath for
+        /// derivations, or the substituted store path for substitions.
+        std::optional<StorePath> storePath = {};
     };
 
 protected:
diff --git a/src/libstore/build/substitution-goal.cc b/src/libstore/build/substitution-goal.cc
index e0ca23a86..2d0c594ac 100644
--- a/src/libstore/build/substitution-goal.cc
+++ b/src/libstore/build/substitution-goal.cc
@@ -3,6 +3,7 @@
 #include "nar-info.hh"
 #include "signals.hh"
 #include "finally.hh"
+#include <boost/outcome/try.hpp>
 #include <kj/array.h>
 #include <kj/vector.h>
 
@@ -54,7 +55,7 @@ try {
 
     /* If the path already exists we're done. */
     if (!repair && worker.store.isValidPath(storePath)) {
-        return {done(ecSuccess, BuildResult::AlreadyValid)};
+        co_return done(ecSuccess, BuildResult::AlreadyValid);
     }
 
     if (settings.readOnlyMode)
@@ -62,9 +63,11 @@ try {
 
     subs = settings.useSubstitutes ? getDefaultSubstituters() : std::list<ref<Store>>();
 
-    return tryNext();
+    BOOST_OUTCOME_CO_TRY(auto result, co_await tryNext());
+    result.storePath = storePath;
+    co_return result;
 } catch (...) {
-    return {std::current_exception()};
+    co_return result::failure(std::current_exception());
 }
 
 
diff --git a/src/libstore/build/worker.cc b/src/libstore/build/worker.cc
index 10f58f5d3..2a764b193 100644
--- a/src/libstore/build/worker.cc
+++ b/src/libstore/build/worker.cc
@@ -229,8 +229,8 @@ try {
     co_return result::failure(std::current_exception());
 }
 
-Worker::Results Worker::run(std::function<Targets (GoalFactory &)> req)
-{
+kj::Promise<Result<Worker::Results>> Worker::run(std::function<Targets (GoalFactory &)> req)
+try {
     auto topGoals = req(goalFactory());
 
     assert(!running);
@@ -252,16 +252,18 @@ Worker::Results Worker::run(std::function<Targets (GoalFactory &)> req)
         promise = promise.exclusiveJoin(boopGC(*localStore));
     }
 
-    return promise.wait(aio.waitScope).value();
+    co_return co_await promise;
+} catch (...) {
+    co_return result::failure(std::current_exception());
 }
 
 kj::Promise<Result<Worker::Results>> Worker::runImpl(Targets topGoals)
 try {
     debug("entered goal loop");
 
-    kj::Vector<Targets::value_type> promises(topGoals.size());
-    for (auto & gp : topGoals) {
-        promises.add(std::move(gp));
+    kj::Vector<std::pair<size_t, kj::Promise<Result<Goal::WorkResult>>>> promises(topGoals.size());
+    for (auto && [idx, gp] : enumerate(topGoals)) {
+        promises.add(idx, std::move(gp.second));
     }
 
     Results results;
@@ -270,7 +272,7 @@ try {
     while (auto done = co_await collect.next()) {
         // propagate goal exceptions outward
         BOOST_OUTCOME_CO_TRY(auto result, done->second);
-        results.emplace(done->first, result);
+        results.goals.emplace(done->first, result);
 
         /* If a top-level goal failed, then kill all other goals
            (unless keepGoing was set). */
@@ -285,6 +287,25 @@ try {
        --keep-going *is* set, then they must all be finished now. */
     assert(!settings.keepGoing || children.isEmpty());
 
+    results.failingExitStatus = [&] {
+        // See API docs in header for explanation
+        unsigned int mask = 0;
+        bool buildFailure = permanentFailure || timedOut || hashMismatch;
+        if (buildFailure)
+            mask |= 0x04;  // 100
+        if (timedOut)
+            mask |= 0x01;  // 101
+        if (hashMismatch)
+            mask |= 0x02;  // 102
+        if (checkMismatch) {
+            mask |= 0x08;  // 104
+        }
+
+        if (mask)
+            mask |= 0x60;
+        return mask ? mask : 1;
+    }();
+
     co_return std::move(results);
 } catch (...) {
     co_return result::failure(std::current_exception());
@@ -301,27 +322,6 @@ try {
 }
 
 
-unsigned int Worker::failingExitStatus()
-{
-    // See API docs in header for explanation
-    unsigned int mask = 0;
-    bool buildFailure = permanentFailure || timedOut || hashMismatch;
-    if (buildFailure)
-        mask |= 0x04;  // 100
-    if (timedOut)
-        mask |= 0x01;  // 101
-    if (hashMismatch)
-        mask |= 0x02;  // 102
-    if (checkMismatch) {
-        mask |= 0x08;  // 104
-    }
-
-    if (mask)
-        mask |= 0x60;
-    return mask ? mask : 1;
-}
-
-
 bool Worker::pathContentsGood(const StorePath & path)
 {
     auto i = pathContentsGoodCache.find(path);
diff --git a/src/libstore/build/worker.hh b/src/libstore/build/worker.hh
index 1a913ca16..369e58b41 100644
--- a/src/libstore/build/worker.hh
+++ b/src/libstore/build/worker.hh
@@ -85,8 +85,40 @@ protected:
 class Worker : public WorkerBase
 {
 public:
-    using Targets = std::map<GoalPtr, kj::Promise<Result<Goal::WorkResult>>>;
-    using Results = std::map<GoalPtr, Goal::WorkResult>;
+    using Targets = std::vector<std::pair<GoalPtr, kj::Promise<Result<Goal::WorkResult>>>>;
+    struct Results {
+        /** Results of individual goals, if available. Goal results will be
+          * added to this map with the index they had in the `Targets` list
+          * returned by the goal factory function passed to `work`. If some
+          * goals did not complete processing, e.g. due to an early exit on
+          * goal failures, not all indices will be set. This may be used to
+          * detect which of the goals were cancelled before they completed.
+          */
+        std::map<size_t, Goal::WorkResult> goals;
+
+        /**
+         * The exit status in case of failure.
+         *
+         * In the case of a build failure, returned value follows this
+         * bitmask:
+         *
+         * ```
+         * 0b1100100
+         *      ^^^^
+         *      |||`- timeout
+         *      ||`-- output hash mismatch
+         *      |`--- build failure
+         *      `---- not deterministic
+         * ```
+         *
+         * In other words, the failure code is at least 100 (0b1100100), but
+         * might also be greater.
+         *
+         * Otherwise (no build failure, but some other sort of failure by
+         * assumption), this returned value is 1.
+         */
+        unsigned int failingExitStatus;
+    };
 
 private:
 
@@ -192,6 +224,7 @@ public:
     NotifyingCounter<uint64_t> expectedNarSize{[this] { updateStatisticsLater(); }};
     NotifyingCounter<uint64_t> doneNarSize{[this] { updateStatisticsLater(); }};
 
+private:
     Worker(Store & store, Store & evalStore, kj::AsyncIoContext & aio);
     ~Worker();
 
@@ -202,7 +235,6 @@ public:
     /**
      * @ref DerivationGoal "derivation goal"
      */
-private:
     template<typename ID, std::derived_from<Goal> G>
     std::pair<std::shared_ptr<G>, kj::Promise<Result<Goal::WorkResult>>> makeGoalCommon(
         std::map<ID, CachedGoal<G>> & map,
@@ -246,30 +278,7 @@ public:
     /**
      * Loop until the specified top-level goals have finished.
      */
-    Results run(std::function<Targets (GoalFactory &)> req);
-
-    /***
-     * The exit status in case of failure.
-     *
-     * In the case of a build failure, returned value follows this
-     * bitmask:
-     *
-     * ```
-     * 0b1100100
-     *      ^^^^
-     *      |||`- timeout
-     *      ||`-- output hash mismatch
-     *      |`--- build failure
-     *      `---- not deterministic
-     * ```
-     *
-     * In other words, the failure code is at least 100 (0b1100100), but
-     * might also be greater.
-     *
-     * Otherwise (no build failure, but some other sort of failure by
-     * assumption), this returned value is 1.
-     */
-    unsigned int failingExitStatus();
+    kj::Promise<Result<Results>> run(std::function<Targets (GoalFactory &)> req);
 
     /**
      * Check whether the given valid path exists and has the right
@@ -278,6 +287,20 @@ public:
     bool pathContentsGood(const StorePath & path);
 
     void markContentsGood(const StorePath & path);
+
+    template<typename MkGoals>
+    friend kj::Promise<Result<Results>> processGoals(
+        Store & store, Store & evalStore, kj::AsyncIoContext & aio, MkGoals && mkGoals
+    ) noexcept;
 };
 
+template<typename MkGoals>
+kj::Promise<Result<Worker::Results>> processGoals(
+    Store & store, Store & evalStore, kj::AsyncIoContext & aio, MkGoals && mkGoals
+) noexcept
+try {
+    co_return co_await Worker(store, evalStore, aio).run(std::forward<MkGoals>(mkGoals));
+} catch (...) {
+    co_return result::failure(std::current_exception());
+}
 }