STOR-5552: Preserve in-flight WebSocket auto-responses (with fix)

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

@@ -258,7 +258,6 @@ IoOwn<LegacyWebSocketAdapter::Native> LegacyWebSocketAdapter::initNative(IoConte
   // We might have called `close()` when this WebSocket was previously active.
   // If so, we want to prevent any future calls to `send()`.
   nativeObj->closedOutgoing = closedOutgoingConn;
-  autoResponseStatus.isClosed = nativeObj->closedOutgoing;
   return ioContext.addObject(kj::mv(nativeObj));
 }
 
@@ -279,7 +278,11 @@ LegacyWebSocketAdapter::LegacyWebSocketAdapter(jsg::Lock& js,
           ws,
           kj::mv(KJ_REQUIRE_NONNULL(package.maybeTags)),
           package.closedOutgoingConnection)),
-      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())) {}
+      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())),
+      autoResponseStatusOwner(ioContext.addObject(kj::heap<AutoResponse>())),
+      autoResponseStatus(*autoResponseStatusOwner) {
+  autoResponseStatus.isClosed = farNative->closedOutgoing;
+}
 // This constructor is used when reinstantiating a websocket that had been hibernating, which is
 // why we can go straight to the Accepted state. However, note that we are actually in the
 // `Hibernatable` "sub-state"!
@@ -292,7 +295,9 @@ LegacyWebSocketAdapter::LegacyWebSocketAdapter(
                                                                         : BinaryType::ARRAYBUFFER),
       allowHalfOpen(!FeatureFlags::get(js).getWebSocketAutoReplyToClose()),
       farNative(nullptr),
-      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())) {
+      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())),
+      autoResponseStatusOwner(IoContext::current().addObject(kj::heap<AutoResponse>())),
+      autoResponseStatus(*autoResponseStatusOwner) {
   auto nativeObj = kj::heap<Native>();
   nativeObj->state.init<AwaitingAcceptanceOrCoupling>(kj::mv(native));
   farNative = IoContext::current().addObject(kj::mv(nativeObj));
@@ -305,7 +310,9 @@ LegacyWebSocketAdapter::LegacyWebSocketAdapter(jsg::Lock& js, WebSocket& shell,
                                                                         : BinaryType::ARRAYBUFFER),
       allowHalfOpen(!FeatureFlags::get(js).getWebSocketAutoReplyToClose()),
       farNative(nullptr),
-      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())) {
+      outgoingMessages(IoContext::current().addObject(kj::heap<OutgoingMessagesMap>())),
+      autoResponseStatusOwner(IoContext::current().addObject(kj::heap<AutoResponse>())),
+      autoResponseStatus(*autoResponseStatusOwner) {
   auto nativeObj = kj::heap<Native>();
   nativeObj->state.init<AwaitingConnection>();
   farNative = IoContext::current().addObject(kj::mv(nativeObj));
@@ -940,6 +947,7 @@ void LegacyWebSocketAdapter::close(
 
   native.closedOutgoing = true;
   closedOutgoingForHib = true;
+  autoResponseStatus.isClosed = true;
   ensurePumping(js);
 }
 
@@ -1132,18 +1140,15 @@ void LegacyWebSocketAdapter::ensurePumping(jsg::Lock& js) {
 }
 
 kj::Promise<void> LegacyWebSocketAdapter::sendAutoResponse(kj::String message, kj::WebSocket& ws) {
-  if (autoResponseStatus.isPumping) {
-    autoResponseStatus.pendingAutoResponseDeque.push(kj::mv(message));
-  } else if (!autoResponseStatus.isClosed) {
-    auto p = ws.send(message).fork();
-    KJ_IF_SOME(context, IoContext::tryCurrent()) {
-      autoResponseStatus.ongoingAutoResponse.emplace(context.addObject(kj::heap(p.addBranch())));
-    } else {
-      // Called outside an IoContext (e.g. from the hibernation manager's readLoop).
-      autoResponseStatus.ongoingAutoResponse.emplace(kj::heap(p.addBranch()));
-    }
-    co_await p;
-    autoResponseStatus.ongoingAutoResponse = kj::none;
+  if (autoResponseStatus.isClosed) {
+    return kj::READY_NOW;
+  } else if (autoResponseStatus.isPumping) {
+    auto completion = kj::newPromiseAndFulfiller<void>();
+    autoResponseStatus.pendingAutoResponseDeque.push(
+        AutoResponse::Pending{kj::mv(message), kj::mv(completion.fulfiller)});
+    return kj::mv(completion.promise);
+  } else {
+    return ws.send(message).attach(kj::mv(message));
   }
 }
 
@@ -1196,7 +1201,12 @@ kj::Promise<void> LegacyWebSocketAdapter::pump(IoContext& context,
 
     autoResponse.isPumping = false;
 
-    autoResponse.pendingAutoResponseDeque.clear();
+    // Preserve the existing behavior of silently dropping queued auto-responses when the pump
+    // exits before sending them.
+    while (!autoResponse.pendingAutoResponseDeque.empty()) {
+      auto pending = KJ_ASSERT_NONNULL(autoResponse.pendingAutoResponseDeque.pop());
+      pending.fulfiller->fulfill();
+    }
 
     if (!completed) {
       // We didn't make it to `completed = true` at the end of this function, so either an
@@ -1247,10 +1257,11 @@ kj::Promise<void> LegacyWebSocketAdapter::pump(IoContext& context,
       auto size = countBytesFromMessage(gatedMessage.message);
 
       while (gatedMessage.pendingAutoResponses > 0) {
-        auto message = KJ_ASSERT_NONNULL(autoResponse.pendingAutoResponseDeque.pop());
+        auto pending = KJ_ASSERT_NONNULL(autoResponse.pendingAutoResponseDeque.pop());
+        KJ_DEFER(pending.fulfiller->fulfill());
         gatedMessage.pendingAutoResponses--;
         autoResponse.queuedAutoResponses--;
-        co_await ws.send(message);
+        co_await ws.send(pending.message);
       }
 
       KJ_SWITCH_ONEOF(gatedMessage.message) {
@@ -1281,8 +1292,9 @@ kj::Promise<void> LegacyWebSocketAdapter::pump(IoContext& context,
     // If there are any auto-responses left to process, we should do it now.
     // We should also check if the last sent message was a close. Shouldn't happen.
     while (!autoResponse.pendingAutoResponseDeque.empty() && !autoResponse.isClosed) {
-      auto message = KJ_ASSERT_NONNULL(autoResponse.pendingAutoResponseDeque.pop());
-      co_await ws.send(message);
+      auto pending = KJ_ASSERT_NONNULL(autoResponse.pendingAutoResponseDeque.pop());
+      KJ_DEFER(pending.fulfiller->fulfill());
+      co_await ws.send(pending.message);
     }
 
     // While we were `co_await`ing the auto-response send, more messages could have been queued
@@ -1386,6 +1398,7 @@ kj::Promise<kj::Maybe<kj::Exception>> LegacyWebSocketAdapter::readLoop(
 
               native.closedOutgoing = true;
               closedOutgoingForHib = true;
+              autoResponseStatus.isClosed = true;
               ensurePumping(js);
             }
             shell.dispatchEventImpl(

src/workerd/api/web-socket.h

@@ -756,16 +756,22 @@ class LegacyWebSocketAdapter final: public WebSocketAdapter {
   struct AutoResponse {
     using OwnedAutoResponsePromise =
         kj::OneOf<IoOwn<kj::Promise<void>>, kj::Own<kj::Promise<void>>>;
+    struct Pending {
+      kj::String message;
+      // Queued auto-responses are historically fire-and-forget. Completion means the pump no
+      // longer owns this item, not necessarily that the send succeeded.
+      kj::Own<kj::PromiseFulfiller<void>> fulfiller;
+    };
     kj::Maybe<OwnedAutoResponsePromise> ongoingAutoResponse;
-    workerd::util::Queue<kj::String> pendingAutoResponseDeque;
+    workerd::util::Queue<Pending> pendingAutoResponseDeque;
     size_t queuedAutoResponses = 0;
     bool isPumping = false;
     bool isClosed = false;
 
     JSG_MEMORY_INFO(AutoResponse) {
       tracker.trackFieldWithSize("ongoingAutoResponse", sizeof(kj::Promise<void>));
       pendingAutoResponseDeque.forEach(
-          [&](const kj::String& message) { tracker.trackField(nullptr, message); });
+          [&](const Pending& pending) { tracker.trackField(nullptr, pending.message); });
     }
   };
 
@@ -880,7 +886,10 @@ class LegacyWebSocketAdapter final: public WebSocketAdapter {
   // the map without locking the isolate.
   IoOwn<OutgoingMessagesMap> outgoingMessages;
 
-  AutoResponse autoResponseStatus;
+  // Auto-responses can run without a current IoContext, so they access the state directly while
+  // the IoOwn ensures it is destroyed by the owning IoContext.
+  IoOwn<AutoResponse> autoResponseStatusOwner;
+  AutoResponse& autoResponseStatus;
 
   kj::Maybe<kj::Own<WebSocketObserver>> observer;