From 53df814534c727c69d769e9fca2dba252d1590a5 Mon Sep 17 00:00:00 2001
From: Milan Miladinovic <milan@cloudflare.com>
Date: Wed, 11 Oct 2023 15:16:13 -0400
Subject: [PATCH] fixup! Evict actors after 10 sec inactivity in Workerd

---
 src/workerd/server/server.c++ | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++
index 7486a3a2211..0922d3041e2 100644
--- a/src/workerd/server/server.c++
+++ b/src/workerd/server/server.c++
@@ -1415,8 +1415,11 @@ public:
           ref.container = kj::none;
         }
 
+        // Don't erase the onBrokenTask if it is the reason we are being destroyed.
+        if (!onBrokenTriggered) {
+          parent.onBrokenTasks.erase(key);
+        }
         // We need to make sure we're removed from the actors map.
-        parent.onBrokenTasks.erase(key);
         parent.actors.erase(key);
       }
 
@@ -1506,6 +1509,9 @@ public:
       bool hasClients() { return containerRef != kj::none; }
       kj::Maybe<ActorContainerRef&> getContainerRef() { return containerRef; }
 
+      // `onBrokenTriggered` indicates the actor has been broken.
+      void setOnBroken() { onBrokenTriggered = true; }
+
       // The actor is constructed after the ActorContainer so it starts off empty.
       kj::Maybe<kj::Own<Worker::Actor>> actor;
     private:
@@ -1516,6 +1522,7 @@ public:
       kj::TimePoint lastAccess;
       kj::Maybe<kj::Own<Worker::Actor::HibernationManager>> manager;
       kj::Maybe<kj::Promise<void>> shutdownTask;
+      bool onBrokenTriggered = false;
 
       // Non-empty if at least one client has a reference to this actor.
       // If no clients are connected, we may be evicted by `cleanupLoop`.
@@ -1779,7 +1786,10 @@ public:
       }
       // Note that we remove the entire ActorContainer from the map -- this drops the
       // HibernationManager so any connected hibernatable websockets will be disconnected.
-      actors.erase(entryRef.getKey());
+      entryRef.setOnBroken();
+      auto key = kj::str(entryRef.getKey());
+      actors.erase(key.asPtr());
+      onBrokenTasks.erase(key);
     }
   };