Merge pull request #664 from cloudflare/milan/hibernation-lifetime

Move owned WebSocket back to api::WebSocket on final hibernation event

Author: MellowYarker

src/workerd/api/global-scope.c++

@@ -527,8 +527,12 @@ void ServiceWorkerGlobalScope::sendHibernatableWebSocketClose(
 
   KJ_IF_MAYBE(h, exportedHandler) {
     KJ_IF_MAYBE(handler, h->webSocketClose) {
-      auto websocket = event->getWebSocket(lock);
-      websocket->initiateHibernatableRelease(lock, api::WebSocket::HibernatableReleaseState::CLOSE);
+      // We won't be dispatching any further events because we've received a close, so we return the
+      // owned websocket back to the api::WebSocket.
+      auto releasePackage = event->prepareForRelease(lock);
+      auto& websocket = releasePackage.webSocketRef;
+      websocket->initiateHibernatableRelease(lock,
+          kj::mv(releasePackage.ownedWebSocket), api::WebSocket::HibernatableReleaseState::CLOSE);
       auto promise = (*handler)(lock, kj::mv(websocket), close.code, kj::mv(close.reason),
                                 close.wasClean);
       event->waitUntil(kj::mv(promise));
@@ -545,8 +549,12 @@ void ServiceWorkerGlobalScope::sendHibernatableWebSocketError(
 
   KJ_IF_MAYBE(h, exportedHandler) {
     KJ_IF_MAYBE(handler, h->webSocketError) {
-      auto websocket = event->getWebSocket(lock);
-      websocket->initiateHibernatableRelease(lock, WebSocket::HibernatableReleaseState::ERROR);
+      // We won't be dispatching any further events because we've encountered an error, so we return
+      // the owned websocket back to the api::WebSocket.
+      auto releasePackage = event->prepareForRelease(lock);
+      auto& websocket = releasePackage.webSocketRef;
+      websocket->initiateHibernatableRelease(lock,
+          kj::mv(releasePackage.ownedWebSocket), WebSocket::HibernatableReleaseState::ERROR);
       jsg::Lock& js(lock);
       auto promise = (*handler)(js, kj::mv(websocket), js.exceptionToJs(kj::mv(e)));
       event->waitUntil(kj::mv(promise));

src/workerd/api/hibernatable-web-socket.c++

@@ -12,10 +12,22 @@ namespace workerd::api {
 HibernatableWebSocketEvent::HibernatableWebSocketEvent()
     : ExtendableEvent("webSocketMessage") {};
 
+HibernatableWebSocketEvent::ItemsForRelease HibernatableWebSocketEvent::prepareForRelease(
+    jsg::Lock &lock) {
+  auto& actor = KJ_REQUIRE_NONNULL(IoContext::current().getActor());
+  auto& manager = kj::downcast<HibernationManagerImpl>(
+      KJ_REQUIRE_NONNULL(actor.getHibernationManager()));
+  auto& hibernatableWebSocket = KJ_REQUIRE_NONNULL(manager.webSocketForEventHandler);
+
+  // `getWebSocket()` requires `HibernatableWebSocket::ws` be non-null, so it must be called first.
+  // The explicit ctor makes it less likely we accidentally move the owned websocket first.
+  return ItemsForRelease(getWebSocket(lock), kj::mv(KJ_REQUIRE_NONNULL(hibernatableWebSocket.ws)));
+}
+
 jsg::Ref<WebSocket> HibernatableWebSocketEvent::getWebSocket(jsg::Lock& lock) {
-  auto& manager = static_cast<HibernationManagerImpl&>(
-      KJ_REQUIRE_NONNULL(
-          KJ_REQUIRE_NONNULL(IoContext::current().getActor()).getHibernationManager()));
+  auto& actor = KJ_REQUIRE_NONNULL(IoContext::current().getActor());
+  auto& manager = kj::downcast<HibernationManagerImpl>(
+      KJ_REQUIRE_NONNULL(actor.getHibernationManager()));
   auto& hibernatableWebSocket = KJ_REQUIRE_NONNULL(manager.webSocketForEventHandler);
   return hibernatableWebSocket.getActiveOrUnhibernate(lock);
 }

src/workerd/api/hibernatable-web-socket.h

@@ -22,6 +22,23 @@ class HibernatableWebSocketEvent final: public ExtendableEvent {
 
   static jsg::Ref<HibernatableWebSocketEvent> constructor(kj::String type) = delete;
 
+  struct ItemsForRelease {
+    // When we call a close or error event, we need to move the owned websocket back into the
+    // api::WebSocket to extend its lifetime. The way we obtain the websocket from the
+    // HibernationManager is somewhat fragile, so it's better if we group the reference and owned
+    // websocket together.
+    jsg::Ref<WebSocket> webSocketRef;
+    kj::Own<kj::WebSocket> ownedWebSocket;
+
+    explicit ItemsForRelease(jsg::Ref<WebSocket> ref, kj::Own<kj::WebSocket> owned)
+        : webSocketRef(kj::mv(ref)), ownedWebSocket(kj::mv(owned)) {}
+  };
+
+  ItemsForRelease prepareForRelease(jsg::Lock& lock);
+  // Only call this once (when transferring ownership of the websocket back to the api::WebSocket).
+  // Gets a reference to the api::WebSocket, and moves the owned kj::WebSocket out of the
+  // HibernatableWebSocket whose event we are currently delivering.
+
   jsg::Ref<WebSocket> getWebSocket(jsg::Lock& lock);
 
   JSG_RESOURCE_TYPE(HibernatableWebSocketEvent) {

src/workerd/api/web-socket.c++

@@ -373,7 +373,7 @@ WebSocket::Accepted::Accepted(kj::Own<kj::WebSocket> wsParam, Native& native, Io
 }
 
 WebSocket::Accepted::Accepted(Hibernatable wsParam, Native& native, IoContext& context)
-    : ws(wsParam),
+    : ws(kj::mv(wsParam)),
       whenAbortedTask(createAbortTask(native, context)) {
   KJ_IF_MAYBE(a, context.getActor()) {
     auto& metrics = a->getMetrics();
@@ -699,7 +699,7 @@ void WebSocket::ensurePumping(jsg::Lock& js) {
           KJ_FAIL_ASSERT("Unexpected native web socket state", native.state);
         }
       }
-    }, [this](jsg::Lock& js, jsg::Value&& exception) mutable {
+    }, [this, thisHandle = JSG_THIS](jsg::Lock& js, jsg::Value&& exception) mutable {
       if (awaitingHibernatableRelease()) {
         // We have a hibernatable websocket -- we don't want to dispatch a regular error event.
         tryReleaseNative(js);

src/workerd/api/web-socket.h

@@ -290,7 +290,9 @@ class WebSocket: public EventTarget {
     ERROR
   };
 
-  void initiateHibernatableRelease(jsg::Lock& js, HibernatableReleaseState releaseState) {
+  void initiateHibernatableRelease(jsg::Lock& js,
+      kj::Own<kj::WebSocket> ws,
+      HibernatableReleaseState releaseState) {
     // Called when a Hibernatable WebSocket wants to dispatch a close/error event, this modifies
     // our `Accepted` state to prepare the state to transition to `Released`.
     //
@@ -300,7 +302,7 @@ class WebSocket: public EventTarget {
     KJ_IF_MAYBE(state, farNative->state.tryGet<Accepted>()) {
       KJ_REQUIRE(state->isHibernatable(),
           "tried to initiate hibernatable release but websocket wasn't hibernatable");
-      state->ws.initiateHibernatableRelease(js, releaseState);
+      state->ws.initiateHibernatableRelease(js, kj::mv(ws), releaseState);
       farNative->closedIncoming = true;
     } else {
       KJ_LOG(WARNING, "Unexpected Hibernatable WebSocket state on release", farNative->state);
@@ -456,6 +458,11 @@ class WebSocket: public EventTarget {
       // A `Hibernatable` WebSocket shares a sub-set of behavior that's already implemented for an
       // `Accepted` WebSocket, so we can think of it a sub-state.
       kj::WebSocket& ws;
+      kj::Maybe<kj::Own<void>> attachedForClose;
+      // If we have initiated a hibernatable error/close event, we need to take back ownership of
+      // the kj::WebSocket so any final queued messages will deliver. We store this owned websocket
+      // in `attachedForClose`. Since the `ws` reference is still valid, we prevent usage of
+      // `attachedForClose` directly in favor of using continuing to use `ws` directly.
       HibernatableReleaseState releaseState = HibernatableReleaseState::NONE;
       // We can't move the state to Released after the Hibernatable Close/Error event runs, since
       // we don't have a request on the thread by the time the event completes.
@@ -514,10 +521,15 @@ class WebSocket: public EventTarget {
         return inner.tryGet<Hibernatable>();
       }
 
-      void initiateHibernatableRelease(jsg::Lock& js, HibernatableReleaseState state) {
+      void initiateHibernatableRelease(jsg::Lock& js,
+          kj::Own<kj::WebSocket> ws,
+          HibernatableReleaseState state) {
         // Transitions our Hibernatable websocket to a "Releasing" state.
         // The websocket will transition to `Released` when convenient.
-        KJ_REQUIRE_NONNULL(getIfHibernatable()).releaseState = state;
+        auto& hibernatable = KJ_REQUIRE_NONNULL(getIfHibernatable());
+        hibernatable.releaseState = state;
+        // Note that we move the owned kj::WebSocket here.
+        hibernatable.attachedForClose = kj::mv(ws);
       }
 
       bool isAwaitingRelease() {
@@ -544,18 +556,18 @@ class WebSocket: public EventTarget {
       return ws.getIfNotHibernatable() == nullptr;
     }
 
+    kj::Promise<void> createAbortTask(Native& native, IoContext& context);
+    kj::Promise<void> whenAbortedTask = nullptr;
+    // Listens for ws->whenAborted() and possibly triggers a proactive shutdown.
+
+    kj::Maybe<kj::Own<ActorObserver>> actorMetrics;
+
     kj::Canceler canceler;
     // This canceler wraps the pump loop as a precaution to make sure we can't exit the Accepted
     // state with a pump task still happening asychronously. In practice the canceler should usually
     // be empty when destroyed because we do not leave the Accepted state if we're still pumping.
     // Even in the case of IoContext premature cancellation, the pump task should be canceled
     // by the IoContext before the Canceler is destroyed.
-
-    kj::Promise<void> createAbortTask(Native& native, IoContext& context);
-    kj::Promise<void> whenAbortedTask = nullptr;
-    // Listens for ws->whenAborted() and possibly triggers a proactive shutdown.
-
-    kj::Maybe<kj::Own<ActorObserver>> actorMetrics;
   };
 
   struct Released {};

src/workerd/io/hibernation-manager.c++

@@ -7,7 +7,7 @@
 
 namespace workerd {
 
-HibernationManagerImpl::~HibernationManagerImpl() {
+HibernationManagerImpl::~HibernationManagerImpl() noexcept(false) {
   // Note that the HibernatableWebSocket destructor handles removing any references to itself in
   // `tagToWs`, and even removes the hashmap entry if there are no more entries in the bucket.
   allWs.clear();
@@ -156,7 +156,7 @@ kj::Promise<void> HibernationManagerImpl::handleSocketTermination(
 
 kj::Promise<void> HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) {
   // Like the api::WebSocket readLoop(), but we dispatch different types of events.
-  auto& ws = *hib.ws;
+  auto& ws = *KJ_REQUIRE_NONNULL(hib.ws);
   while (true) {
     kj::WebSocket::Message message = co_await ws.receive();
     // Note that errors are handled by the callee of `readLoop`, since we throw from `receive()`.

src/workerd/io/hibernation-manager.h

@@ -22,7 +22,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager {
         hibernationEventType(hibernationEventType),
         onDisconnect(DisconnectHandler{}),
         readLoopTasks(onDisconnect) {}
-  ~HibernationManagerImpl();
+  ~HibernationManagerImpl() noexcept(false);
 
   void acceptWebSocket(jsg::Ref<api::WebSocket> ws, kj::ArrayPtr<kj::String> tags) override;
   // Tells the HibernationManager to create a new HibernatableWebSocket with the associated tags
@@ -39,7 +39,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager {
   // Hibernates all the websockets held by the HibernationManager.
   // This converts our activeOrPackage from an api::WebSocket to a HibernationPackage.
 
-  friend jsg::Ref<api::WebSocket> api::HibernatableWebSocketEvent::getWebSocket(jsg::Lock& lock);
+  friend class api::HibernatableWebSocketEvent;
 
 private:
   class HibernatableWebSocket;
@@ -73,7 +73,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager {
           ws(activeOrPackage.get<jsg::Ref<api::WebSocket>>()->acceptAsHibernatable()),
           manager(manager) {}
 
-    ~HibernatableWebSocket() {
+    ~HibernatableWebSocket() noexcept(false) {
       // We expect this dtor to be called when we're removing a HibernatableWebSocket
       // from our `allWs` collection in the HibernationManager.
 
@@ -100,7 +100,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager {
       // to the api::WebSocket.
       KJ_IF_MAYBE(package, activeOrPackage.tryGet<api::WebSocket::HibernationPackage>()) {
         activeOrPackage.init<jsg::Ref<api::WebSocket>>(
-            api::WebSocket::hibernatableFromNative(js, *ws, kj::mv(*package)));
+            api::WebSocket::hibernatableFromNative(js, *KJ_REQUIRE_NONNULL(ws), kj::mv(*package)));
       }
       return activeOrPackage.get<jsg::Ref<api::WebSocket>>().addRef();
     }
@@ -116,7 +116,11 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager {
     kj::OneOf<jsg::Ref<api::WebSocket>, api::WebSocket::HibernationPackage> activeOrPackage;
     // If active, we have an api::WebSocket reference, otherwise, we're hibernating, so we retain
     // the websocket's properties in a HibernationPackage until it's time to wake up.
-    kj::Own<kj::WebSocket> ws;
+    kj::Maybe<kj::Own<kj::WebSocket>> ws;
+    // This is an owned websocket that we extract from the api::WebSocket after accepting as
+    // hibernatable. It becomes null once we dispatch a close or error event because we want its
+    // lifetime to be managed by IoContext's DeleteQueue. This helps prevent a situation where the
+    // HibernationManager drops the websocket before all queued messages have sent.
     HibernationManagerImpl& manager;
     // TODO(someday): We (currently) only use the HibernationManagerImpl reference to refer to
     // `tagToWs` when running the dtor for `HibernatableWebSocket`. This feels a bit excessive,