Implement setInactivityTimeout in ContainerClient

setInactivityTimeout allows a container to outlive its associated actor.
The current implementation tightly couples the lifetime of the container
client to the lifetime of the actor, so that the ContainerClient is
always destroyed whenever the actor is destroyed.

To fix this we make ContainerClient reference counted and provide a
reference to the actor. Whenever setInactivityTimeout is called we add a
reference to the client. If the actor shuts down, the remaining
reference keeps the container alive until the timeout expires. If
setInactivityTimeout was never called then there is no remaining
reference, and the container shuts down when the actor does.

We also implement the ability for an actor to reconnect to an already
running container when it restarts. We do this by maintaining a map of
containerId->ContainerClient in the ActorNamespace. This map does not
hold a reference to the ContainerClient: only actors and the inactivity
timeout promise hold references to ensure we do not leak resources by
holding a reference indefinitely. Whenever an actor starts it consults
the map to see if a ContainerClient already exists. If it doesn't we
create one and add it to the map.

In order to clear entries from the map when a ContainerClient is
destroyed we provide a callback function (`cleanupCallback`) to the
ContainerClient which is run when ContainerClient is destroyed.

Author: gpanders

src/workerd/server/container-client.c++

@@ -107,16 +107,22 @@ ContainerClient::ContainerClient(capnp::ByteStreamFactory& byteStreamFactory,
     kj::String dockerPath,
     kj::String containerName,
     kj::String imageName,
-    kj::TaskSet& waitUntilTasks)
+    kj::TaskSet& waitUntilTasks,
+    kj::Function<void()> cleanupCallback)
     : byteStreamFactory(byteStreamFactory),
       timer(timer),
       network(network),
       dockerPath(kj::mv(dockerPath)),
-      containerName(kj::encodeUriComponent(kj::mv(containerName))),
+      containerName(kj::encodeUriComponent(containerName)),
       imageName(kj::mv(imageName)),
-      waitUntilTasks(waitUntilTasks) {}
+      waitUntilTasks(waitUntilTasks),
+      cleanupCallback(kj::mv(cleanupCallback)) {}
 
 ContainerClient::~ContainerClient() noexcept(false) {
+  // Call the cleanup callback to remove this client from the ActorNamespace map
+  cleanupCallback();
+
+  // Destroy the Docker container
   waitUntilTasks.add(dockerApiRequest(network, kj::str(dockerPath), kj::HttpMethod::DELETE,
       kj::str("/containers/", containerName, "?force=true"))
                          .ignoreResult());
@@ -427,7 +433,19 @@ kj::Promise<void> ContainerClient::signal(SignalContext context) {
 }
 
 kj::Promise<void> ContainerClient::setInactivityTimeout(SetInactivityTimeoutContext context) {
-  // empty implementation on purpose
+  auto params = context.getParams();
+  auto durationMs = params.getDurationMs();
+
+  JSG_REQUIRE(
+      durationMs > 0, Error, "setInactivityTimeout() requires durationMs > 0, got ", durationMs);
+
+  auto timeout = durationMs * kj::MILLISECONDS;
+
+  // Add a timer task that holds a reference to this ContainerClient.
+  waitUntilTasks.add(timer.afterDelay(timeout).then([self = kj::addRef(*this)]() {
+    // This callback does nothing but drop the reference
+  }));
+
   co_return;
 }
 
@@ -444,4 +462,8 @@ kj::Promise<void> ContainerClient::listenTcp(ListenTcpContext context) {
   KJ_UNIMPLEMENTED("listenTcp not implemented for Docker containers - use port mapping instead");
 }
 
+kj::Own<ContainerClient> ContainerClient::addRef() {
+  return kj::addRef(*this);
+}
+
 }  // namespace workerd::server

src/workerd/server/container-client.h

@@ -11,6 +11,7 @@
 #include <kj/async.h>
 #include <kj/compat/http.h>
 #include <kj/map.h>
+#include <kj/refcount.h>
 #include <kj/string.h>
 
 namespace workerd::server {
@@ -19,15 +20,20 @@ namespace workerd::server {
 // so it can be used as a rpc::Container::Client via kj::heap<ContainerClient>().
 // This allows the Container JSG class to use Docker directly without knowing
 // it's talking to Docker instead of a real RPC service.
-class ContainerClient final: public rpc::Container::Server {
+//
+// ContainerClient is reference-counted to support actor reconnection with inactivity timeouts.
+// When setInactivityTimeout() is called, a timer holds a reference to prevent premature
+// destruction. The ContainerClient can be shared across multiple actor lifetimes.
+class ContainerClient final: public rpc::Container::Server, public kj::Refcounted {
  public:
   ContainerClient(capnp::ByteStreamFactory& byteStreamFactory,
       kj::Timer& timer,
       kj::Network& network,
       kj::String dockerPath,
       kj::String containerName,
       kj::String imageName,
-      kj::TaskSet& waitUntilTasks);
+      kj::TaskSet& waitUntilTasks,
+      kj::Function<void()> cleanupCallback);
 
   ~ContainerClient() noexcept(false);
 
@@ -41,6 +47,9 @@ class ContainerClient final: public rpc::Container::Server {
   kj::Promise<void> listenTcp(ListenTcpContext context) override;
   kj::Promise<void> setInactivityTimeout(SetInactivityTimeoutContext context) override;
 
+  // Reference counting support
+  kj::Own<ContainerClient> addRef();
+
  private:
   capnp::ByteStreamFactory& byteStreamFactory;
   kj::Timer& timer;
@@ -82,6 +91,9 @@ class ContainerClient final: public rpc::Container::Server {
   kj::Promise<void> stopContainer();
   kj::Promise<void> killContainer(uint32_t signal);
   kj::Promise<void> destroyContainer();
+
+  // Cleanup callback to remove from ActorNamespace map when destroyed
+  kj::Function<void()> cleanupCallback;
 };
 
 }  // namespace workerd::server

src/workerd/server/server.c++

@@ -2347,6 +2347,11 @@ class Server::WorkerService final: public Service,
           auto reason = 0;
           a->shutdown(reason);
         }
+
+        // Drop the container client reference
+        // If setInactivityTimeout() was called, there's still a timer holding a reference
+        // If not, this may be the last reference and the ContainerClient destructor will run
+        containerClient = kj::none;
       }
 
       void active() override {
@@ -2475,6 +2480,7 @@ class Server::WorkerService final: public Service,
         manager = kj::none;
         tracker->shutdown();
         actor = kj::none;
+        containerClient = kj::none;
 
         KJ_IF_SOME(r, reason) {
           brokenReason = kj::cp(r);
@@ -2546,6 +2552,9 @@ class Server::WorkerService final: public Service,
       kj::Maybe<kj::Promise<void>> onBrokenTask;
       kj::Maybe<kj::Exception> brokenReason;
 
+      // Reference to the ContainerClient (if container is enabled for this actor)
+      kj::Maybe<kj::Own<ContainerClient>> containerClient;
+
       // If this is a `ForkedPromise<void>`, await the promise. When it has resolved, then
       // `classAndId` will have been replaced with the resolved `ClassAndId` value.
       kj::OneOf<ClassAndId, kj::ForkedPromise<void>> classAndId;
@@ -2717,6 +2726,11 @@ class Server::WorkerService final: public Service,
         }
         // Destroy the last strong Worker::Actor reference.
         actor = kj::none;
+
+        // Drop our reference to the ContainerClient
+        // If setInactivityTimeout() was called, the timer still holds a reference
+        // so the container stays alive until the timeout expires
+        containerClient = kj::none;
       }
 
       void start(kj::Own<ActorClass>& actorClass, Worker::Actor::Id& id) {
@@ -2801,30 +2815,35 @@ class Server::WorkerService final: public Service,
 
         auto loopback = kj::refcounted<Loopback>(*this);
 
-        kj::Maybe<rpc::Container::Client> containerClient = kj::none;
+        kj::Maybe<rpc::Container::Client> containerClientCapability = kj::none;
         KJ_IF_SOME(config, containerOptions) {
-          auto& dockerPathRef = KJ_ASSERT_NONNULL(ns.dockerPath,
-              "dockerPath needs to be defined in order enable containers on this durable object.");
           KJ_ASSERT(config.hasImageName(), "Image name is required");
           auto imageName = config.getImageName();
-          kj::String containerId;
+          kj::String containerIdStr;
           KJ_SWITCH_ONEOF(id) {
             KJ_CASE_ONEOF(globalId, kj::Own<ActorIdFactory::ActorId>) {
-              containerId = globalId->toString();
+              containerIdStr = globalId->toString();
             }
             KJ_CASE_ONEOF(existingId, kj::String) {
-              containerId = kj::str(existingId);
+              containerIdStr = kj::str(existingId);
             }
           }
-          containerClient = kj::heap<ContainerClient>(ns.byteStreamFactory, timer, ns.dockerNetwork,
-              kj::str(dockerPathRef),
-              kj::str("workerd-", KJ_ASSERT_NONNULL(uniqueKey), "-", containerId),
-              kj::str(imageName), ns.waitUntilTasks);
+
+          // Create the full container name
+          auto fullContainerId =
+              kj::str("workerd-", KJ_ASSERT_NONNULL(uniqueKey), "-", containerIdStr);
+
+          // Get or create the ContainerClient (may reuse existing one from previous actor instance)
+          containerClient = ns.getOrCreateContainerClient(fullContainerId, imageName);
+
+          // Get a capability to pass to the Actor
+          // This creates an additional reference that the Actor's IoContext holds
+          containerClientCapability = KJ_ASSERT_NONNULL(containerClient)->addRef();
         }
 
         auto actor = actorClass->newActor(getTracker(), Worker::Actor::cloneId(id),
             kj::mv(makeActorCache), kj::mv(makeStorage), kj::mv(loopback), tryGetManagerRef(),
-            kj::mv(containerClient), *this);
+            kj::mv(containerClientCapability), *this);
         onBrokenTask = monitorOnBroken(*actor);
         this->actor = kj::mv(actor);
       }
@@ -2862,11 +2881,46 @@ class Server::WorkerService final: public Service,
       })->addRef();
     }
 
+    kj::Own<ContainerClient> getOrCreateContainerClient(
+        kj::StringPtr containerId, kj::StringPtr imageName) {
+      // Check if ContainerClient already exists in map
+      KJ_IF_SOME(existingClient, containerClients.find(containerId)) {
+        // Found existing client, return a new reference to it
+        return existingClient->addRef();
+      }
+
+      // Not found, create new ContainerClient
+      auto& dockerPathRef = KJ_ASSERT_NONNULL(
+          dockerPath, "dockerPath must be defined to enable containers on this Durable Object.");
+
+      // Create cleanup callback to remove from map when destroyed
+      kj::Function<void()> cleanupCallback = [this, containerId]() {
+        // Remove raw pointer from map when destructor runs
+        containerClients.erase(containerId);
+      };
+
+      // Create new ContainerClient
+      auto client = kj::refcounted<ContainerClient>(byteStreamFactory, timer, dockerNetwork,
+          kj::str(dockerPathRef), kj::str(containerId), kj::str(imageName), waitUntilTasks,
+          kj::mv(cleanupCallback));
+
+      // Store raw pointer in map (does not own)
+      containerClients.insert(kj::str(containerId), client.get());
+
+      // Return owned reference
+      return kj::mv(client);
+    }
+
     void abortAll(kj::Maybe<const kj::Exception&> reason) {
       for (auto& actor: actors) {
         actor.value->abort(reason);
       }
       actors.clear();
+
+      // Clean up all container clients
+      // This drops the ActorNamespace's reference to each ContainerClient
+      // If any have active timers, they'll still be kept alive until timer expires
+      containerClients.clear();
     }
 
    private:
@@ -2890,6 +2944,12 @@ class Server::WorkerService final: public Service,
     // inactivity, we keep the ActorContainer in the map but drop the Own<Worker::Actor>. When a new
     // request comes in, we recreate the Own<Worker::Actor>.
     ActorMap actors;
+
+    // Map of container IDs to ContainerClients (for reconnection support with inactivity timeouts).
+    // The map holds raw pointers (not ownership) - ContainerClients are owned by actors and timers.
+    // When the last reference is dropped, the destructor removes the entry from this map.
+    kj::HashMap<kj::StringPtr, ContainerClient*> containerClients;
+
     kj::Maybe<kj::Promise<void>> cleanupTask;
     kj::Timer& timer;
     capnp::ByteStreamFactory& byteStreamFactory;