ws: close server sockets through the drain path and keep them CLOSING until 'close' is emitted

Closed
RO
1mo ago
+468-36

Rewritten after the first review: no immediates, the close goes through the drain path.

Problem

  • A server socket of the built-in ws shim is CLOSED when close() returns, and a 'close' listener added after it never fires. npm ws is CLOSING and emits 'close' later, with a Buffer reason. Cause: uws runs the close callback inside ServerWebSocket.close() (packages/bun-uws/src/WebSocket.h:270) and #close (ws.js:1056 on main) emitted right there.
  • close() with messages still queued for drain lost them. send() after close dropped its callback, bufferedAmount ignored getBufferedAmount(), and #drain resent messages uws had accepted.
  • uws never sent the FIN for a close() issued inside the drain handler. The connection only went away through the end() timer, 16 s by default. Plain Bun.serve reproduces it.

Fix

  • close() records the pending close and runs the drain step: queued messages go out first, the native close follows once getBufferedAmount() is 0. The close callback stays in CLOSING, fails the callbacks of undelivered messages, and emits 'close' on the next tick with a Buffer reason.
  • send() on a closing or closed socket calls back with the "not open" error and counts the bytes. bufferedAmount adds what uws has buffered. A native 0 is a drop only for a non-empty payload.
  • onWritable shuts the socket down after the drain handler when end() was called in it, as onData already does (WebSocketContext.h:339). Bun runs the handler corked, so end() leaves the FIN to the caller.
  • Verified: test/js/first_party/ws/ws.test.ts (block "server socket close lifecycle") and test/js/bun/websocket/websocket-server-backpressure-buffer.test.ts. All 9 new tests fail on the released build. More in the notes.

Background

  • "ws" always resolves to the shim. Its server socket is a JS state machine over the Bun.serve websocket callbacks installed by node:http (_http_server.ts:648).
  • Native send() returns the byte count, -1 if uws buffered the message, 0 if it dropped it (after 16 MiB). The shim queues dropped messages and retries them from drain.
  • After sending a Close frame uws ignores incoming data, so the peer's reply code is not observable. 'close' carries the code passed to close(), which a compliant peer echoes anyway.
Notes

Reference behavior was taken from node v26 with ws 8.18.3 (test/node_modules/ws) using a raw RFC 6455 client, so both implementations saw identical bytes. The shim now matches on: readyState after close() and terminate(), the send and ping callbacks after close, 'close' args (4000, Buffer("bye")) for a local close with an echoing peer, (1006, Buffer(0)) after terminate(), (4444, Buffer("other")) for a peer close, bufferedAmount === 4 after a late send("late"), and send("", cb) calling back with one empty frame on the wire.

Known remaining differences: the close code after a local close() is the local one (native limitation above). The binary frame shape for binaryType = "arraybuffer" was fixed on main by #39808 (now in the base of this branch), the text frame shape is #36061. The deferred 'close' relies on request handlers running with the event loop entered (#39826, in the base of this branch): before that, a nextTick queued inside close() ran before close() returned when called from the 'connection' handler. The test "close() inside the 'connection' handler" pins that case.

Why the close waits for getBufferedAmount() === 0 and not only for the shim queue: uws end() sends the Close frame with send(), and when the socket is over its backpressure limit that returns DROPPED, which end() treats as sent and shuts the socket down at once. The released build delivers 2.5 MiB of 26 MiB in the new drain test for that reason. That is a separate uws issue, reported separately. Waiting for an empty buffer also matches when npm ws emits 'close'.

Probe for the FIN bug with plain Bun.serve: pause a raw client, ws.send() 1 MiB until it returns -1, in drain call ws.close() once getBufferedAmount() is 0, resume the client. Released build: the client gets the Close frame and no FIN for 4 to 16 s (the end() timer). This branch: FIN right after the Close frame. The new test relies on that timer being 16 s at the default idleTimeout, so it times out on the released build instead of passing late.

The backpressure tests are skipped on Windows like the ones already in that file: Winsock's loopback takes the whole payload into the kernel. On Linux and macOS the kernel takes under 8 MiB from a paused reader (the existing 8 MiB test depends on the same), so the queue test sends 26 x 1 MiB to get past the 16 MiB limit. The two blocks take about 1.5 s on a debug ASAN build and were stable over 6 runs.

#drain callbacks moved from queueMicrotask to process.nextTick, which is what send() uses for the direct path.

From the second review round (1607583): payloadByteLength sizes a payload that is neither a string nor buffer-like as the text frame of its string form, which is what the native send() sends, so a drop of such a message is not mistaken for a delivered empty frame. The close callback subtracts the entries it discards instead of zeroing the counter, so bytes counted by a send() during CLOSING stay in bufferedAmount as in npm ws. The test "terminate() with messages still queued fails their callbacks" pins that and the error callbacks of undelivered messages.

Suites run on the debug build: test/js/first_party/ws/ (4 files), node-http-with-ws.test.ts, the ws regression tests (012040, 03844, 26358, 29684, 32734, 3613), test/js/bun/websocket/ (the 8 failures in websocket-server.test.ts are the debug-build fixture timeouts that #39370 addresses, the file passes on the released build here too), test/js/web/websocket/ (2 failures need the public internet). ws-proxy.test.ts needs HTTP_PROXY unset in this container, on the released build as well.
Overlap with #39843: that PR fixes both uws findings from this one in end() itself (Close frame past the backpressure limit, FIN sent by end() after uncork), which covers the drain handler case too. The onWritable hunk here then becomes redundant and can be dropped when #39843 lands first, or #39843 can remove it on rebase. Both can coexist at runtime: shutdown() on an already shut down socket is a no-op in usockets (raw and SSL). The two PRs also both add tests to websocket-server-backpressure-buffer.test.ts, so whichever lands second needs a small rebase.

Rebased onto main after #39642 and #39808 landed in ws.js. The only conflicts were adjacent additions (the controlPayload helper next to payloadByteLength/wasDropped, and the EventEmitter/isWindows imports in the test file), kept both sides.

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