9 files changed, 28 insertions, 15 deletions
diff --git a/src/libstore/build/derivation-goal.cc b/src/libstore/build/derivation-goal.cc
index c6ac39ddf..edee09e13 100644
--- a/src/libstore/build/derivation-goal.cc
+++ b/src/libstore/build/derivation-goal.cc
@@ -125,7 +125,7 @@ Goal::WorkResult DerivationGoal::timedOut(Error && ex)
 }
 
 
-kj::Promise<Result<Goal::WorkResult>> DerivationGoal::work() noexcept
+kj::Promise<Result<Goal::WorkResult>> DerivationGoal::workImpl() noexcept
 {
     return useDerivation ? getDerivation() : haveDerivation();
 }
diff --git a/src/libstore/build/derivation-goal.hh b/src/libstore/build/derivation-goal.hh
index 0cacf75fd..b461b7d0d 100644
--- a/src/libstore/build/derivation-goal.hh
+++ b/src/libstore/build/derivation-goal.hh
@@ -250,7 +250,7 @@ struct DerivationGoal : public Goal
 
     WorkResult timedOut(Error && ex);
 
-    kj::Promise<Result<WorkResult>> work() noexcept override;
+    kj::Promise<Result<WorkResult>> workImpl() noexcept override;
 
     /**
      * Add wanted outputs to an already existing derivation goal.
diff --git a/src/libstore/build/drv-output-substitution-goal.cc b/src/libstore/build/drv-output-substitution-goal.cc
index 765908aeb..0d6f3fce0 100644
--- a/src/libstore/build/drv-output-substitution-goal.cc
+++ b/src/libstore/build/drv-output-substitution-goal.cc
@@ -24,7 +24,7 @@ DrvOutputSubstitutionGoal::DrvOutputSubstitutionGoal(
 }
 
 
-kj::Promise<Result<Goal::WorkResult>> DrvOutputSubstitutionGoal::work() noexcept
+kj::Promise<Result<Goal::WorkResult>> DrvOutputSubstitutionGoal::workImpl() noexcept
 try {
     trace("init");
 
diff --git a/src/libstore/build/drv-output-substitution-goal.hh b/src/libstore/build/drv-output-substitution-goal.hh
index 86020705e..f959e2a7b 100644
--- a/src/libstore/build/drv-output-substitution-goal.hh
+++ b/src/libstore/build/drv-output-substitution-goal.hh
@@ -70,7 +70,7 @@ public:
     kj::Promise<Result<WorkResult>> outPathValid() noexcept;
     kj::Promise<Result<WorkResult>> finished() noexcept;
 
-    kj::Promise<Result<WorkResult>> work() noexcept override;
+    kj::Promise<Result<WorkResult>> workImpl() noexcept override;
 
     JobCategory jobCategory() const override {
         return JobCategory::Substitution;
diff --git a/src/libstore/build/goal.cc b/src/libstore/build/goal.cc
index 5910c37b2..7fbf43045 100644
--- a/src/libstore/build/goal.cc
+++ b/src/libstore/build/goal.cc
@@ -1,6 +1,7 @@
 #include "goal.hh"
 #include "async-collect.hh"
 #include "worker.hh"
+#include <boost/outcome/try.hpp>
 #include <kj/time.h>
 
 namespace nix {
@@ -19,6 +20,23 @@ kj::Promise<void> Goal::waitForAWhile()
     return worker.aio.provider->getTimer().afterDelay(settings.pollInterval.get() * kj::SECONDS);
 }
 
+kj::Promise<Result<Goal::WorkResult>> Goal::work() noexcept
+try {
+    BOOST_OUTCOME_CO_TRY(auto result, co_await workImpl());
+
+    trace("done");
+    assert(!exitCode.has_value());
+    exitCode = result.exitCode;
+    ex = result.ex;
+
+    notify->fulfill();
+    cleanup();
+
+    co_return std::move(result);
+} catch (...) {
+    co_return result::failure(std::current_exception());
+}
+
 kj::Promise<Result<void>>
 Goal::waitForGoals(kj::Array<std::pair<GoalPtr, kj::Promise<void>>> dependencies) noexcept
 try {
diff --git a/src/libstore/build/goal.hh b/src/libstore/build/goal.hh
index 03c2cf309..10926fffc 100644
--- a/src/libstore/build/goal.hh
+++ b/src/libstore/build/goal.hh
@@ -92,7 +92,7 @@ struct Goal
      */
     BuildResult buildResult;
 
-    // for use by Worker only. will go away once work() is a promise.
+    // for use by Worker and Goal only. will go away once work() is a promise.
     kj::Own<kj::PromiseFulfiller<void>> notify;
 
 protected:
@@ -121,6 +121,8 @@ protected:
         return waitForGoals(kj::arrOf<std::pair<GoalPtr, kj::Promise<void>>>(std::move(goals)...));
     }
 
+    virtual kj::Promise<Result<WorkResult>> workImpl() noexcept = 0;
+
 public:
 
     /**
@@ -138,7 +140,7 @@ public:
         trace("goal destroyed");
     }
 
-    virtual kj::Promise<Result<WorkResult>> work() noexcept = 0;
+    kj::Promise<Result<WorkResult>> work() noexcept;
 
     virtual void waiteeDone(GoalPtr waitee) { }
 
diff --git a/src/libstore/build/substitution-goal.cc b/src/libstore/build/substitution-goal.cc
index c0dd95da5..206bc8649 100644
--- a/src/libstore/build/substitution-goal.cc
+++ b/src/libstore/build/substitution-goal.cc
@@ -46,7 +46,7 @@ Goal::WorkResult PathSubstitutionGoal::done(
 }
 
 
-kj::Promise<Result<Goal::WorkResult>> PathSubstitutionGoal::work() noexcept
+kj::Promise<Result<Goal::WorkResult>> PathSubstitutionGoal::workImpl() noexcept
 try {
     trace("init");
 
diff --git a/src/libstore/build/substitution-goal.hh b/src/libstore/build/substitution-goal.hh
index 41e665779..18b4262a4 100644
--- a/src/libstore/build/substitution-goal.hh
+++ b/src/libstore/build/substitution-goal.hh
@@ -87,7 +87,7 @@ public:
     );
     ~PathSubstitutionGoal();
 
-    kj::Promise<Result<WorkResult>> work() noexcept override;
+    kj::Promise<Result<WorkResult>> workImpl() noexcept override;
 
     /**
      * The states.
diff --git a/src/libstore/build/worker.cc b/src/libstore/build/worker.cc
index 0e8694a6d..e90f17678 100644
--- a/src/libstore/build/worker.cc
+++ b/src/libstore/build/worker.cc
@@ -195,19 +195,12 @@ static void removeGoal(std::shared_ptr<G> goal, auto & goalMap)
 
 void Worker::goalFinished(GoalPtr goal, Goal::WorkResult & f)
 {
-    goal->trace("done");
-    assert(!goal->exitCode.has_value());
-    goal->exitCode = f.exitCode;
-    goal->ex = f.ex;
-
     permanentFailure |= f.permanentFailure;
     timedOut |= f.timedOut;
     hashMismatch |= f.hashMismatch;
     checkMismatch |= f.checkMismatch;
 
     removeGoal(goal);
-    goal->notify->fulfill();
-    goal->cleanup();
 }
 
 void Worker::removeGoal(GoalPtr goal)