From a451899366d28e0f15c03fb35f029772e7297243 Mon Sep 17 00:00:00 2001 From: Ashley Peacock Date: Wed, 19 Aug 2026 12:32:50 +0100 Subject: [PATCH 1/2] STOR-5552: Preserve in-flight WebSocket auto-responses Track auto-response sends in the hibernation manager so revived WebSockets wait before writing. Queue completion alongside pending responses and retain the message buffer independently of the legacy adapter. Cover direct, queued, packaged, cancellation, and repeated-revival paths. --- src/workerd/api/web-socket.c++ | 40 +-- src/workerd/api/web-socket.h | 10 +- src/workerd/io/hibernation-manager-test.c++ | 255 ++++++++++++++---- src/workerd/io/legacy-hibernation-manager.c++ | 27 +- src/workerd/io/legacy-hibernation-manager.h | 6 +- 5 files changed, 247 insertions(+), 91 deletions(-) diff --git a/src/workerd/api/web-socket.c++ b/src/workerd/api/web-socket.c++ index 098fce96948..e179f1a84b3 100644 --- a/src/workerd/api/web-socket.c++ +++ b/src/workerd/api/web-socket.c++ @@ -940,6 +940,7 @@ void LegacyWebSocketAdapter::close( native.closedOutgoing = true; closedOutgoingForHib = true; + autoResponseStatus.isClosed = true; ensurePumping(js); } @@ -1132,18 +1133,15 @@ void LegacyWebSocketAdapter::ensurePumping(jsg::Lock& js) { } kj::Promise 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(); + 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 +1194,12 @@ kj::Promise 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 +1250,11 @@ kj::Promise 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 +1285,9 @@ kj::Promise 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 +1391,7 @@ kj::Promise> LegacyWebSocketAdapter::readLoop( native.closedOutgoing = true; closedOutgoingForHib = true; + autoResponseStatus.isClosed = true; ensurePumping(js); } shell.dispatchEventImpl( diff --git a/src/workerd/api/web-socket.h b/src/workerd/api/web-socket.h index 63ee58a1404..7c8f0732570 100644 --- a/src/workerd/api/web-socket.h +++ b/src/workerd/api/web-socket.h @@ -756,8 +756,14 @@ class LegacyWebSocketAdapter final: public WebSocketAdapter { struct AutoResponse { using OwnedAutoResponsePromise = kj::OneOf>, kj::Own>>; + 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> fulfiller; + }; kj::Maybe ongoingAutoResponse; - workerd::util::Queue pendingAutoResponseDeque; + workerd::util::Queue pendingAutoResponseDeque; size_t queuedAutoResponses = 0; bool isPumping = false; bool isClosed = false; @@ -765,7 +771,7 @@ class LegacyWebSocketAdapter final: public WebSocketAdapter { JSG_MEMORY_INFO(AutoResponse) { tracker.trackFieldWithSize("ongoingAutoResponse", sizeof(kj::Promise)); pendingAutoResponseDeque.forEach( - [&](const kj::String& message) { tracker.trackField(nullptr, message); }); + [&](const Pending& pending) { tracker.trackField(nullptr, pending.message); }); } }; diff --git a/src/workerd/io/hibernation-manager-test.c++ b/src/workerd/io/hibernation-manager-test.c++ index 70364bb1068..0a5f1c95d2c 100644 --- a/src/workerd/io/hibernation-manager-test.c++ +++ b/src/workerd/io/hibernation-manager-test.c++ @@ -12,10 +12,9 @@ // implementation as that work lands. The tests themselves are the source of // truth for the contract; comments are best-effort context. // -// A few tests use KJ_EXPECT_LOG to capture the production "another message -// send is already in progress" assertion as an expected ERROR log. They pass -// while the bug is present and fail loudly when the fix lands. Search the -// file for "regression test for EW-10817" to find them. +// One test uses KJ_EXPECT_LOG to capture the remaining "another message send +// is already in progress" assertion as an expected ERROR log. It passes while +// the general outgoing-queue bug is present and will fail when that fix lands. #include #include @@ -605,17 +604,9 @@ KJ_TEST("HibernationManager: in-flight DO close survives hibernation within one fixture.drainAndDestroy(kj::mv(request)); } -KJ_TEST("HibernationManager: in-flight auto-response orphans BlockedSend during hibernation") { - // Regression test for EW-10817. sendAutoResponse creates a BlockedSend on the pipe (held in - // a plain kj::Own outside any IoOwn — see web-socket.c++:874), then hibernation replaces - // activeOrPackage without carrying that state. The new api::WebSocket's pump skips the wait - // and trips on the orphaned BlockedSend. - // - // The KJ_EXPECT_LOG block below captures the bug's symptom (the assertion's ERROR log) so - // the test passes while EW-10817 is open. When the bug is fixed, the log won't fire and - // the KJ_EXPECT_LOG will fail — that's the signal to update this test (delete the - // EXPECT_LOG block and promote the receive() at the end to a positive assertion about the - // auto-response pong's content). +KJ_TEST("HibernationManager: in-flight auto-response survives repeated hibernation before close") { + // Each replacement api::WebSocket must wait for an auto-response started by the original + // instance before sending its close. DispatchStats stats; TestFixture fixture(stubLoopbackParams(stats, kj::str("ew-10817-autoresp"))); auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); @@ -626,35 +617,100 @@ KJ_TEST("HibernationManager: in-flight auto-response orphans BlockedSend during end1->send("ping"_kj).wait(fixture.getWaitScope()); fixture.pollEventLoop(); - // Hibernate. + // Hibernate, revive without consuming the pong, then hibernate again. The manager must retain + // its own branch of the pending send when it gives the first replacement a branch. + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto websockets = hm->getWebSockets(env.js, kj::none); + KJ_ASSERT(websockets.size() == 1); + }); fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); - // Unhibernate + close → hits orphaned BlockedSend. - { - KJ_EXPECT_LOG(ERROR, "another message send is already in progress"); - fixture.enterContext(*request, [&](const TestFixture::Environment& env) { - auto& js = env.js; - auto websockets = hm->getWebSockets(js, kj::none); - KJ_ASSERT(websockets.size() == 1); - websockets[0]->close(js, 1001, jsg::USVString(kj::str("stale"))); - }); + // Revive again and queue a close behind the in-flight auto-response. + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto& js = env.js; + auto websockets = hm->getWebSockets(js, kj::none); + KJ_ASSERT(websockets.size() == 1); + websockets[0]->close(js, 1001, jsg::USVString(kj::str("after-pong"))); + }); + fixture.pollEventLoop(); - fixture.pollEventLoop(); - } + auto pong = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(pong.is() && pong.get() == "pong"_kj); - // Receive the orphaned pong from the pipe (held outside any IoOwn — the very thing this - // test is documenting). This unblocks the stuck pump so drainAndDestroy() below can - // complete cleanly. Once EW-10817 is fixed, the orphan won't exist; this becomes a - // positive assertion about the pong's content. - end1->receive().wait(fixture.getWaitScope()); + auto closePromise = end1->receive(); + KJ_ASSERT(closePromise.poll(fixture.getWaitScope()), "close did not follow auto-response"); + auto closeMessage = closePromise.wait(fixture.getWaitScope()); + KJ_ASSERT(closeMessage.is()); + auto& close = closeMessage.get(); + KJ_ASSERT(close.code == 1001, close.code); + KJ_ASSERT(close.reason == "after-pong"_kj, close.reason); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST("HibernationManager: packaged in-flight auto-response finishes before close") { + DispatchStats stats; + TestFixture fixture(stubLoopbackParams(stats, kj::str("packaged-autoresp"))); + auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); + auto request = fixture.newIncomingRequest(); + auto end1 = acceptNewWebSocket(fixture, *request, *hm); + + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + + // The manager sends the pong directly while the api::WebSocket is packaged. + end1->send("ping"_kj).wait(fixture.getWaitScope()); + fixture.pollEventLoop(); + + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto websockets = hm->getWebSockets(env.js, kj::none); + KJ_ASSERT(websockets.size() == 1); + websockets[0]->close(env.js, 1001, jsg::USVString(kj::str("after-packaged-pong"))); + }); + fixture.pollEventLoop(); + + auto pong = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(pong.is() && pong.get() == "pong"_kj); + + auto closeMessage = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(closeMessage.is()); + auto& close = closeMessage.get(); + KJ_ASSERT(close.code == 1001, close.code); + KJ_ASSERT(close.reason == "after-packaged-pong"_kj, close.reason); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST("HibernationManager: rejected packaged auto-response removes revived WebSocket") { + DispatchStats stats; + TestFixture fixture(stubLoopbackParams(stats, kj::str("rejected-packaged-autoresp"))); + auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); + auto request = fixture.newIncomingRequest(); + auto end1 = acceptNewWebSocket(fixture, *request, *hm, "pending"_kj); + + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + end1->send("ping"_kj).wait(fixture.getWaitScope()); + fixture.pollEventLoop(); + + // Revival gives the replacement adapter a branch of the blocked send. Disconnecting the peer + // rejects both branches and must terminate the manager's socket exactly once. + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto websockets = hm->getWebSockets(env.js, kj::none); + KJ_ASSERT(websockets.size() == 1); + }); + end1 = nullptr; + fixture.pollEventLoop(); + + KJ_ASSERT(stats.customEventCalls == 1, stats.customEventCalls); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + KJ_ASSERT(hm->getWebSockets(env.js, kj::none).size() == 0); + KJ_ASSERT(hm->getWebSockets(env.js, "pending"_kj).size() == 0); + }); fixture.drainAndDestroy(kj::mv(request)); } KJ_TEST("HibernationManager: in-flight DO send orphans BlockedSend during hibernation") { // Regression test for EW-10817. Same shape as the auto-response variant above but driven by // a DO-side ws.send() — the pump creates a BlockedSend on the pipe (no BPT yet), hibernation - // orphans it, the next operation on the new api::WebSocket trips the assertion. See the - // auto-response variant above for the EXPECT_LOG / lifecycle details. + // orphans it, and the next operation on the new api::WebSocket trips the assertion. // DispatchStats stats; TestFixture fixture(stubLoopbackParams(stats, kj::str("ew-10817-dosend"))); @@ -817,13 +873,9 @@ KJ_TEST("HibernationManager: in-flight DO send lost across IoContext destruction "data frame was silently dropped across IoContext destruction; eyeball receives nothing"); } -KJ_TEST("HibernationManager: in-flight auto-response orphans BlockedSend across actor eviction") { - // Regression test for EW-10817 — the production failure mode. sendAutoResponse runs from - // the HM's readLoop (on the HM's TaskSet, NOT in an IoContext). It does a direct - // kj::WebSocket::send that creates a BlockedSend on the pipe. IoContext destruction cancels - // pump tasks but not sendAutoResponse, so the BlockedSend survives the IoContext's death. - // After actor eviction and revival, the new api::WebSocket's pump trips on the orphan. See - // the same-IoContext auto-response variant above for the EXPECT_LOG / lifecycle details. +KJ_TEST("HibernationManager: in-flight auto-response finishes before close across actor eviction") { + // sendAutoResponse runs from the HM's readLoop, outside the IoContext. Its in-flight send must + // remain visible after hibernation so a revived api::WebSocket waits before sending its close. DispatchStats stats; TestFixture fixture(stubLoopbackParams(stats, kj::str("ew-10817-cross-autoresp"))); auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); @@ -847,22 +899,26 @@ KJ_TEST("HibernationManager: in-flight auto-response orphans BlockedSend across request1 = nullptr; fixture.resetActor(); - // Phase 3: under a brand-new actor + IoContext, do something that starts a fresh pump. - // The new pump trips on the orphaned BlockedSend. + // Under a brand-new actor + IoContext, queue a close behind the in-flight auto-response. auto request2 = fixture.newIncomingRequest(); - { - KJ_EXPECT_LOG(ERROR, "another message send is already in progress"); - fixture.enterContext(*request2, [&](const TestFixture::Environment& env) { - auto& js = env.js; - auto websockets = hm->getWebSockets(js, kj::none); - KJ_ASSERT(websockets.size() == 1); - websockets[0]->close(js, 1001, jsg::USVString(kj::str("post-evict"))); - }); - fixture.pollEventLoop(); - } + fixture.enterContext(*request2, [&](const TestFixture::Environment& env) { + auto& js = env.js; + auto websockets = hm->getWebSockets(js, kj::none); + KJ_ASSERT(websockets.size() == 1); + websockets[0]->close(js, 1001, jsg::USVString(kj::str("post-evict"))); + }); + fixture.pollEventLoop(); - // Receive the orphaned pong (see same-IoContext variant above for why), then drain. - end1->receive().wait(fixture.getWaitScope()); + auto pong = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(pong.is() && pong.get() == "pong"_kj); + + auto closePromise = end1->receive(); + KJ_ASSERT(closePromise.poll(fixture.getWaitScope()), "close did not follow auto-response"); + auto closeMessage = closePromise.wait(fixture.getWaitScope()); + KJ_ASSERT(closeMessage.is()); + auto& close = closeMessage.get(); + KJ_ASSERT(close.code == 1001, close.code); + KJ_ASSERT(close.reason == "post-evict"_kj, close.reason); fixture.drainAndDestroy(kj::mv(request2)); } @@ -937,14 +993,46 @@ KJ_TEST("HibernationManager: DO close waits for the actor's output gate") { fixture.drainAndDestroy(kj::mv(request)); } -KJ_TEST("HibernationManager: auto-response (active) waits when pump is gate-blocked on a DO send") { +KJ_TEST("HibernationManager: auto-response is skipped after close is queued") { + DispatchStats stats; + TestFixture fixture(stubLoopbackParams(stats, kj::str("autoresp-after-close"))); + auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); + auto request = fixture.newIncomingRequest(); + auto end1 = acceptNewWebSocket(fixture, *request, *hm); + + auto paf = kj::newPromiseAndFulfiller(); + auto blocker = fixture.getActor().getOutputGate().lockWhile(kj::mv(paf.promise), nullptr); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto websockets = hm->getWebSockets(env.js, kj::none); + KJ_ASSERT(websockets.size() == 1); + websockets[0]->close(env.js, 1001, jsg::USVString(kj::str("already-closing"))); + }); + + // The pump is blocked before sending close. A later ping must not queue a pong which will be + // discarded when close is eventually sent. + end1->send("ping"_kj).wait(fixture.getWaitScope()); + fixture.pollEventLoop(); + + paf.fulfiller->fulfill(); + auto message = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(message.is()); + KJ_ASSERT(message.get().reason == "already-closing"_kj); + fixture.pollEventLoop(); + KJ_ASSERT(stats.customEventCalls == 0, stats.customEventCalls); + + blocker.wait(fixture.getWaitScope()); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST( + "HibernationManager: queued auto-response survives hibernation while pump is gate-blocked") { // When the pump is already running (isPumping == true) and stalled on the output gate for a // queued DO message, an arriving auto-response request causes sendAutoResponse to push the // pong onto pendingAutoResponseDeque. The pump only drains that deque after it finishes the // outer outgoingMessages loop, so the pong waits for the gate to release transitively. // - // Order at the eyeball: the gated DO message arrives first (after the gate releases), and - // the pong follows immediately after (line 998 in web-socket.c++). + // Order at the eyeball: the gated DO message arrives first after the gate releases, followed by + // the pong and any write queued by a replacement api::WebSocket. DispatchStats stats; TestFixture fixture(stubLoopbackParams(stats, kj::str("output-gate-autoresp-gated"))); auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); @@ -960,6 +1048,15 @@ KJ_TEST("HibernationManager: auto-response (active) waits when pump is gate-bloc // Eyeball sends ping. sendAutoResponse sees isPumping=true and queues "pong". end1->send("ping"_kj).wait(fixture.getWaitScope()); + // The replacement queues its close behind the pong's completion while the original pump remains + // blocked on the output gate. + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto websockets = hm->getWebSockets(env.js, kj::none); + KJ_ASSERT(websockets.size() == 1); + websockets[0]->close(env.js, 1001, jsg::USVString(kj::str("after-queued-pong"))); + }); + // Neither msg1 nor pong has arrived yet. auto receivePromise = end1->receive(); fixture.pollEventLoop(); @@ -972,11 +1069,53 @@ KJ_TEST("HibernationManager: auto-response (active) waits when pump is gate-bloc KJ_ASSERT(msg1.is() && msg1.get() == "msg1"_kj); auto msg2 = end1->receive().wait(fixture.getWaitScope()); KJ_ASSERT(msg2.is() && msg2.get() == "pong"_kj); + auto closeMessage = end1->receive().wait(fixture.getWaitScope()); + KJ_ASSERT(closeMessage.is()); + auto& close = closeMessage.get(); + KJ_ASSERT(close.code == 1001, close.code); + KJ_ASSERT(close.reason == "after-queued-pong"_kj, close.reason); blocker.wait(fixture.getWaitScope()); fixture.drainAndDestroy(kj::mv(request)); } +KJ_TEST("HibernationManager: canceled queued auto-response preserves revived WebSocket") { + DispatchStats stats; + TestFixture fixture(stubLoopbackParams(stats, kj::str("canceled-queued-autoresp"))); + auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); + auto request1 = fixture.newIncomingRequest(); + auto end1 = acceptNewWebSocket(fixture, *request1, *hm, "pending"_kj); + + auto paf = kj::newPromiseAndFulfiller(); + auto blocker = fixture.getActor().getOutputGate().lockWhile(kj::mv(paf.promise), nullptr); + sendFromDo(fixture, *request1, *hm, "blocked"_kj); + end1->send("ping"_kj).wait(fixture.getWaitScope()); + + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + fixture.enterContext(*request1, [&](const TestFixture::Environment& env) { + KJ_ASSERT(hm->getWebSockets(env.js, kj::none).size() == 1); + }); + + // Destroying the IoContext cancels the original pump. As before this fix, its queued pong is + // dropped without terminating the manager's WebSocket. + { + KJ_EXPECT_LOG(WARNING, "failed to invoke drain() on IncomingRequest before destroying it"); + request1 = nullptr; + } + paf.fulfiller->fulfill(); + blocker.wait(fixture.getWaitScope()); + fixture.pollEventLoop(); + + KJ_ASSERT(stats.customEventCalls == 0, stats.customEventCalls); + fixture.resetActor(); + auto request2 = fixture.newIncomingRequest(); + fixture.enterContext(*request2, [&](const TestFixture::Environment& env) { + KJ_ASSERT(hm->getWebSockets(env.js, kj::none).size() == 1); + KJ_ASSERT(hm->getWebSockets(env.js, "pending"_kj).size() == 1); + }); + fixture.drainAndDestroy(kj::mv(request2)); +} + KJ_TEST("HibernationManager: auto-response (active) bypasses the output gate") { // Documents CURRENT behavior: in active mode, sendAutoResponse uses a direct kj::WebSocket::send // that doesn't go through the pump, and therefore doesn't check waitForOutputLocksIfNecessary. diff --git a/src/workerd/io/legacy-hibernation-manager.c++ b/src/workerd/io/legacy-hibernation-manager.c++ index 0af14e77cba..877ab3cbc21 100644 --- a/src/workerd/io/legacy-hibernation-manager.c++ +++ b/src/workerd/io/legacy-hibernation-manager.c++ @@ -67,13 +67,16 @@ jsg::Ref LegacyHibernationManagerImpl::HibernatableWebSocket:: package.maybeTags = getTags(); // Now that we unhibernated the WebSocket, we can set the last received autoResponse timestamp - // that was stored in the corresponding HibernatableWebSocket. We also move autoResponsePromise - // from the hibernation manager to api::websocket to prevent possible ws.send races. + // that was stored in the corresponding HibernatableWebSocket. Share an autoResponsePromise + // branch with api::websocket while retaining the fork in case the socket hibernates again. + kj::Promise autoResponsePromise = kj::READY_NOW; + KJ_IF_SOME(promise, this->maybeAutoResponsePromise) { + autoResponsePromise = promise.addBranch(); + } activeOrPackage .init>( api::WebSocket::hibernatableFromNative(js, *KJ_REQUIRE_NONNULL(ws), kj::mv(package))) ->setAutoResponseStatus(autoResponseTimestamp, kj::mv(autoResponsePromise)); - autoResponsePromise = kj::READY_NOW; } return activeOrPackage.get>().addRef(); } @@ -346,13 +349,15 @@ kj::Promise LegacyHibernationManagerImpl::readLoop(HibernatableWebSocket& auto responseCopy = kj::str(KJ_REQUIRE_NONNULL(autoResponsePair->response)); KJ_SWITCH_ONEOF(hib.activeOrPackage) { KJ_CASE_ONEOF(apiWs, jsg::Ref) { - // If the actor is not hibernated/If the WebSocket is active, we need to update - // autoResponseTimestamp on the active websocket. - apiWs->setAutoResponseStatus(hib.autoResponseTimestamp, kj::READY_NOW); // Since we had a request set, we must have and response that's sent back using the // same websocket here. The sending of response is managed in web-socket to avoid // possible racing problems with regular websocket messages. - co_await apiWs->sendAutoResponse(kj::mv(responseCopy), ws); + hib.maybeAutoResponsePromise = + apiWs->sendAutoResponse(kj::mv(responseCopy), ws).fork(); + auto& promise = KJ_ASSERT_NONNULL(hib.maybeAutoResponsePromise); + apiWs->setAutoResponseStatus(hib.autoResponseTimestamp, promise.addBranch()); + KJ_DEFER(hib.maybeAutoResponsePromise = kj::none); + co_await promise; } KJ_CASE_ONEOF(package, api::WebSocket::HibernationPackage) { if (!package.closedOutgoingConnection) { @@ -360,10 +365,10 @@ kj::Promise LegacyHibernationManagerImpl::readLoop(HibernatableWebSocket& // If we do that, we have to provide it with the promise to avoid races. This can // happen if we have a websocket hibernating, that unhibernates and sends a // message while ws.send() for auto-response is also sending. - auto p = ws.send(responseCopy.asArray()).fork(); - hib.autoResponsePromise = p.addBranch(); - co_await p; - hib.autoResponsePromise = kj::READY_NOW; + hib.maybeAutoResponsePromise = + ws.send(responseCopy.asArray()).attach(kj::mv(responseCopy)).fork(); + KJ_DEFER(hib.maybeAutoResponsePromise = kj::none); + co_await KJ_ASSERT_NONNULL(hib.maybeAutoResponsePromise); } } } diff --git a/src/workerd/io/legacy-hibernation-manager.h b/src/workerd/io/legacy-hibernation-manager.h index 82e25ad54af..e08e5214350 100644 --- a/src/workerd/io/legacy-hibernation-manager.h +++ b/src/workerd/io/legacy-hibernation-manager.h @@ -139,9 +139,9 @@ class LegacyHibernationManagerImpl final: public Worker::Actor::HibernationManag // Stores the last received autoResponseRequest timestamp. kj::Maybe autoResponseTimestamp; - // Keeps track of the currently ongoing websocket auto-response send promise. This promise may - // be moved to api::websocket if an hibernating websocket unhibernates. - kj::Promise autoResponsePromise = kj::READY_NOW; + // Keeps track of the currently ongoing websocket auto-response send promise. A revived + // api::WebSocket receives a branch so the manager can retain this across repeated hibernation. + kj::Maybe> maybeAutoResponsePromise; friend LegacyHibernationManagerImpl; }; From 4ae16ee7be560db718569c05f35e223c410ec7ca Mon Sep 17 00:00:00 2001 From: Ashley Peacock Date: Fri, 21 Aug 2026 13:44:09 +0100 Subject: [PATCH 2/2] STOR-5552: Protect auto-response state from GC Move the WebSocket auto-response state under IoOwn so V8 garbage collection defers destruction of its KJ async objects to the owning IoContext. Retain direct access for hibernation-manager operations that run without a current IoContext.\n\nAdd coverage for collecting a hibernated WebSocket while its auto-response send is still in flight. --- src/workerd/api/web-socket.c++ | 15 ++++++--- src/workerd/api/web-socket.h | 5 ++- src/workerd/io/hibernation-manager-test.c++ | 35 +++++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/src/workerd/api/web-socket.c++ b/src/workerd/api/web-socket.c++ index e179f1a84b3..f8173230a8d 100644 --- a/src/workerd/api/web-socket.c++ +++ b/src/workerd/api/web-socket.c++ @@ -258,7 +258,6 @@ IoOwn 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())) {} + outgoingMessages(IoContext::current().addObject(kj::heap())), + autoResponseStatusOwner(ioContext.addObject(kj::heap())), + 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())) { + outgoingMessages(IoContext::current().addObject(kj::heap())), + autoResponseStatusOwner(IoContext::current().addObject(kj::heap())), + autoResponseStatus(*autoResponseStatusOwner) { auto nativeObj = kj::heap(); nativeObj->state.init(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())) { + outgoingMessages(IoContext::current().addObject(kj::heap())), + autoResponseStatusOwner(IoContext::current().addObject(kj::heap())), + autoResponseStatus(*autoResponseStatusOwner) { auto nativeObj = kj::heap(); nativeObj->state.init(); farNative = IoContext::current().addObject(kj::mv(nativeObj)); diff --git a/src/workerd/api/web-socket.h b/src/workerd/api/web-socket.h index 7c8f0732570..333e4e01c0c 100644 --- a/src/workerd/api/web-socket.h +++ b/src/workerd/api/web-socket.h @@ -886,7 +886,10 @@ class LegacyWebSocketAdapter final: public WebSocketAdapter { // the map without locking the isolate. IoOwn 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 autoResponseStatusOwner; + AutoResponse& autoResponseStatus; kj::Maybe> observer; diff --git a/src/workerd/io/hibernation-manager-test.c++ b/src/workerd/io/hibernation-manager-test.c++ index 0a5f1c95d2c..9051b36ef15 100644 --- a/src/workerd/io/hibernation-manager-test.c++ +++ b/src/workerd/io/hibernation-manager-test.c++ @@ -1222,5 +1222,40 @@ KJ_TEST("HibernationManager: hibernated auto-response copies buffer before suspe fixture.drainAndDestroy(kj::mv(request)); } +KJ_TEST("HibernationManager: GC collects WebSocket with in-flight auto-response") { + DispatchStats stats; + TestFixture fixture(stubLoopbackParams(stats, kj::str("gc-in-flight-autoresp"))); + auto hm = makeTestHm(fixture, "ping"_kj, "pong"_kj); + auto request = fixture.newIncomingRequest(); + + // The shared helper intentionally leaks a ref. This test instead creates the V8 wrapper that + // js.alloc() omits so GC owns the last reference after hibernation. + kj::Own end1; + jsg::WeakRef weakApiWs = nullptr; + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto pipe = kj::newWebSocketPipe(); + end1 = kj::mv(pipe.ends[0]); + auto apiWs = env.js.alloc(env.js, kj::mv(pipe.ends[1])); + weakApiWs = apiWs.getWeakRef(env.js); + auto& handler = KJ_ASSERT_NONNULL(env.js.tryGetTypeHandler>()); + auto wrapper KJ_UNUSED = handler.wrap(env.js, apiWs.addRef()); + hm->acceptWebSocket(kj::mv(apiWs), nullptr); + }); + + end1->send("ping"_kj).wait(fixture.getWaitScope()); + fixture.pollEventLoop(); + fixture.enterWorkerLock([&](Worker::Lock& lock) { hm->hibernateWebSockets(lock); }); + KJ_ASSERT(weakApiWs.isAlive()); + + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + env.isolate->LowMemoryNotification(); + KJ_ASSERT(!weakApiWs.isAlive()); + }); + + end1 = nullptr; + fixture.pollEventLoop(); + fixture.drainAndDestroy(kj::mv(request)); +} + } // namespace } // namespace workerd