diff --git a/src/workerd/api/web-socket.c++ b/src/workerd/api/web-socket.c++ index 098fce96948..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)); @@ -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 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 +1201,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 +1257,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 +1292,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 +1398,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..333e4e01c0c 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); }); } }; @@ -880,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 70364bb1068..9051b36ef15 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. @@ -1083,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 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; };