From 649d8cd08fa16c71e3580f16d77c2122540f3195 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Sat, 5 Oct 2024 00:38:35 +0200
Subject: [PATCH] libstore: remove Worker::removeGoal

we can use our newfound powers of Goal::work Is A Real Promise to remove
completed goals from continuation promises. apart from being much easier
to follow it's also a lot more efficient because we have the iterator to
the item we are trying to remove, skipping a linear search of the cache.

Change-Id: Ie0190d051c5f4b81304d98db478348b20c209df5
---
 src/libstore/build/goal.hh   |  5 ----
 src/libstore/build/worker.cc | 56 ++++++++++++------------------------
 src/libstore/build/worker.hh |  7 +----
 3 files changed, 20 insertions(+), 48 deletions(-)

diff --git a/src/libstore/build/goal.hh b/src/libstore/build/goal.hh
index b524d3118..de1c92c85 100644
--- a/src/libstore/build/goal.hh
+++ b/src/libstore/build/goal.hh
@@ -82,11 +82,6 @@ struct Goal
      */
     std::string name;
 
-    struct WorkResult;
-
-    // for use by Worker and Goal only. will go away once work() is a promise.
-    kj::Own<kj::PromiseFulfiller<Result<WorkResult>>> notify;
-
 protected:
     AsyncSemaphore::Token slotToken;
 
diff --git a/src/libstore/build/worker.cc b/src/libstore/build/worker.cc
index 4e8fa38db..82acbdb3d 100644
--- a/src/libstore/build/worker.cc
+++ b/src/libstore/build/worker.cc
@@ -71,25 +71,37 @@ std::pair<std::shared_ptr<G>, kj::Promise<Result<Goal::WorkResult>>> Worker::mak
     // and then we only want to recreate the goal *once*. concurrent accesses
     // to the worker are not sound, we want to catch them if at all possible.
     for ([[maybe_unused]] auto _attempt : {1, 2}) {
-        auto & goal_weak = it->second;
-        auto goal = goal_weak.goal.lock();
+        auto & cachedGoal = it->second;
+        auto & goal = cachedGoal.goal;
         if (!goal) {
             goal = create();
-            goal_weak.goal = goal;
             // do not start working immediately. if we are not yet running we
             // may create dependencies as though they were toplevel goals, in
             // which case the dependencies will not report build errors. when
             // we are running we may be called for this same goal more times,
             // and then we want to modify rather than recreate when possible.
-            goal_weak.promise = kj::evalLater([goal] { return goal->work(); }).fork();
-            childStarted(goal, goal_weak.promise.addBranch());
+            auto removeWhenDone = [goal, &map, it] {
+                // c++ lambda coroutine capture semantics are *so* fucked up.
+                return [](auto goal, auto & map, auto it) -> kj::Promise<Result<Goal::WorkResult>> {
+                    auto result = co_await goal->work();
+                    // a concurrent call to makeGoalCommon may have reset our
+                    // cached goal and replaced it with a new instance. don't
+                    // remove the goal in this case, otherwise we will crash.
+                    if (goal == it->second.goal) {
+                        map.erase(it);
+                    }
+                    co_return result;
+                }(goal, map, it);
+            };
+            cachedGoal.promise = kj::evalLater(std::move(removeWhenDone)).fork();
+            childStarted(goal, cachedGoal.promise.addBranch());
         } else {
             if (!modify(*goal)) {
-                goal_weak = {};
+                cachedGoal = {};
                 continue;
             }
         }
-        return {goal, goal_weak.promise.addBranch()};
+        return {goal, cachedGoal.promise.addBranch()};
     }
     assert(false && "could not make a goal. possible concurrent worker access");
 }
@@ -184,44 +196,14 @@ std::pair<GoalPtr, kj::Promise<Result<Goal::WorkResult>>> Worker::makeGoal(const
 }
 
 
-template<typename G>
-static void removeGoal(std::shared_ptr<G> goal, auto & goalMap)
-{
-    /* !!! inefficient */
-    for (auto i = goalMap.begin();
-         i != goalMap.end(); )
-        if (i->second.goal.lock() == goal) {
-            auto j = i; ++j;
-            goalMap.erase(i);
-            i = j;
-        }
-        else ++i;
-}
-
-
 void Worker::goalFinished(GoalPtr goal, Goal::WorkResult & f)
 {
     permanentFailure |= f.permanentFailure;
     timedOut |= f.timedOut;
     hashMismatch |= f.hashMismatch;
     checkMismatch |= f.checkMismatch;
-
-    removeGoal(goal);
 }
 
-void Worker::removeGoal(GoalPtr goal)
-{
-    if (auto drvGoal = std::dynamic_pointer_cast<DerivationGoal>(goal))
-        nix::removeGoal(drvGoal, derivationGoals);
-    else if (auto subGoal = std::dynamic_pointer_cast<PathSubstitutionGoal>(goal))
-        nix::removeGoal(subGoal, substitutionGoals);
-    else if (auto subGoal = std::dynamic_pointer_cast<DrvOutputSubstitutionGoal>(goal))
-        nix::removeGoal(subGoal, drvOutputSubstitutionGoals);
-    else
-        assert(false);
-}
-
-
 void Worker::childStarted(GoalPtr goal, kj::Promise<Result<Goal::WorkResult>> promise)
 {
     children.add(promise
diff --git a/src/libstore/build/worker.hh b/src/libstore/build/worker.hh
index 923092b51..fcf0ad8c7 100644
--- a/src/libstore/build/worker.hh
+++ b/src/libstore/build/worker.hh
@@ -95,7 +95,7 @@ private:
     template<typename G>
     struct CachedGoal
     {
-        std::weak_ptr<G> goal;
+        std::shared_ptr<G> goal;
         kj::ForkedPromise<Result<Goal::WorkResult>> promise{nullptr};
     };
     /**
@@ -134,11 +134,6 @@ private:
 
     void goalFinished(GoalPtr goal, Goal::WorkResult & f);
 
-    /**
-     * Remove a dead goal.
-     */
-    void removeGoal(GoalPtr goal);
-
     /**
      * Registers a running child process.
      */