Bun.serve websocket: deliver queued publish() and corked send() before a forced close or process exit

Closed
RO
1mo ago
+342-8

Problem

  • server.publish() of a message under 16 KB returns the byte count, but no subscriber receives it when server.stop(true), ws.terminate(), server.unref() or process.exit() follows in the same tick. ws.send() inside a handler is lost the same way.
  • us_socket_close closes the fd before WebSocketContext::onClose runs (packages/bun-uws/src/WebSocketContext.h:269). onClose frees the subscriber without a drain and discards the cork buffer.
  • At process exit the loop does not tick again, so nothing drains the TopicTree (App.h:533).

Fix

  • WebSocket::flushBeforeForcedClose() drains the socket's subscriber, then uncorks. WebSocket::close() (ws.terminate()) and the parser's forceClose call it before us_socket_close. An open handler that throws uses close(false), so the corked 101 stays unsent.
  • App::close() (server.stop(true)) drains the TopicTree and uncorks its corked websocket before it closes the groups.
  • New Loop::flushPendingWrites() runs the uWS pre handler. VirtualMachine::on_exit calls it after the exit listeners.
  • Each flush writes only what publish() and send() already reported as sent. #32852 drains the same way. Verified: six new tests in test/js/bun/websocket/websocket-server.test.ts. Five fail on bun 1.4.2. All pass on a debug build. Self-reviewed: 3 concerns raised, 1 addressed (Notes).

Background

  • The TopicTree is the uWS pub/sub index. A small publish() is stored once, with an index per subscriber. A loop pre or post handler, or the socket's next send(), writes the batch.
  • A cork is a 16 KB per-loop buffer. uWS corks a socket while its data handler runs, so send() frames go out in one write when the handler returns.
  • ws.terminate(), server.stop(true) and parser errors close the TCP socket at once, with no Close frame.
Notes

Repro for the first case (from the report):

const rc = server.publish("all", "bye"); // rc = 3
server.stop(true);                       // subscribers receive nothing

Controls that always delivered: ws.send() from a timer, a publish of 16 KB or more (it bypasses the batch), ws.close() after the publish, setTimeout(() => server.stop(true), 0). The batch also writes itself out at the 32nd pending message, so 40 small publishes then stop(true) delivered exactly the first 32.

Test cases added:

  • send() and publish() then server.stop(true) inside a message handler. The handler's socket is corked, so its publishes drain into the cork buffer. It must receive sent, bye0, bye1, bye2. The other subscriber must receive bye0, bye1, bye2.
  • The same with ws.terminate() on every socket.
  • A child process that publishes from a timer, then calls server.unref() or process.exit(0). Skipped on Windows: the exit resets the connection there and the reset can discard what the client has not read, the same reason serve.test.ts skips its process-exit test.
  • A raw TCP client that writes a legal text frame and an orphan continuation frame in one write. The echo of the first frame must arrive before the 1006 close.
  • An open handler that sends a frame and then throws, after a synchronous server.upgrade(). The raw client must receive no response at all. This passes on bun 1.4.2. It fails on the first version of this PR, where close() flushed the corked 101 and the client saw open, the frame, then a 1006 close.

on_exit runs for a natural exit, process.exit(), and a worker that exits on its own. A worker that its parent terminated runs no script, and the flush skips it. on_exit uses uws::Loop::get() and not vm.uws_loop(): Bun.spawnSync swaps event_loop_handle to an isolated loop that has no uWS LoopData. The drain callback only corks, sends and uncorks, so the flush runs no JavaScript.

Self-review concerns:

  • Addressed: the two process-exit tests were racy on Windows. They now skip there.
  • Not changed: on wss://, stop(true) ends with a reset after the TLS close_notify (packages/bun-usockets/src/context.c:127). That is the existing behavior for every byte written before stop(true), ws.send() included.
  • Not changed: a publish() made from a close handler while stop(true) walks the sockets is still dropped for the sockets that close later in the same walk. A flush there needs a hook before each us_socket_close inside us_socket_group_close_all.

Also not changed: onEnd (the peer sent FIN) still frees the subscriber without a drain. On Linux the FIN is dispatched in a later loop iteration than the data before it (usockets does not register EPOLLRDHUP), so the loop's own drain runs first and a test for it passes with or without a flush there. kqueue can report EOF on the data's event (packages/bun-usockets/src/loop.c, the eof drain), so a half-closing client on macOS can still lose a same-tick publish. I could not produce a failing case on this machine, so that path is left for a separate change.

#43013 edits the same lines of WebSocket::close() and forceClose for a different bug (a TLS forced close that waits for the peer). The two changes compose: the flush goes before whichever us_socket_close call runs.

Suites run on the debug build: test/js/bun/websocket/, test/js/bun/http/serve.test.ts, bun-server.test.ts, serve-listen.test.ts, six files of test/js/web/websocket/, test/js/web/workers/worker.test.ts, test/js/node/process/process.test.js. cargo check --target x86_64-pc-windows-msvc -p bun_uws_sys -p bun_jsc passes.

Local failures that this diff does not cause:

  • websocket-server.test.ts: eight send()/sendText()/sendBinary() tests time out at 10 s on the debug build while the concurrent (benchmark) test runs for 29 s. Each passes alone.
  • test/js/first_party/ws/: debug panic assertion failed: !self.body_read_ref.get().has, the subject of #43427.
  • serve.test.ts: two tests need a non-root user and open egress. process.test.js: one test needs USER in the environment.

no test proof · iteration 1 · platform-specific test(s) that do not run on this machine, deferring to CI, which covers all platforms: test/js/bun/websocket/websocket-server.test.ts