fix(websocket): implement receive backpressure for WebSocketStream

lib/web/websocket/stream/websocketstream.js

@@ -33,6 +33,7 @@ class WebSocketStream {
   #readableStream
   /** @type {ReadableStreamDefaultController} */
   #readableStreamController
+  #parserBackpressured = false
 
   // Each WebSocketStream object has an associated writable stream , which is a WritableStream .
   /** @type {WritableStream} */
@@ -47,9 +48,13 @@ class WebSocketStream {
     onConnectionEstablished: (response, extensions) => this.#onConnectionEstablished(response, extensions),
     onMessage: (opcode, data) => this.#onMessage(opcode, data),
     onParserError: (err) => failWebsocketConnection(this.#handler, null, err.message),
-    onParserDrain: () => this.#handler.socket.resume(),
+    onParserDrain: () => {
+      this.#parserBackpressured = false
+      this.#resumeSocketIfNeeded()
+    },
     onSocketData: (chunk) => {
       if (!this.#parser.write(chunk)) {
+        this.#parserBackpressured = true
         this.#handler.socket.pause()
       }
     },
@@ -117,7 +122,9 @@ class WebSocketStream {
     this.#closedPromise = Promise.withResolvers()
 
     // 7. Apply backpressure to the WebSocket.
-    // TODO
+    // The socket and readable stream are not available until the connection is
+    // established; backpressure is applied from the readable pull algorithm and
+    // after messages are enqueued.
 
     // 8.  If options [" signal "] exists ,
     if (options.signal != null) {
@@ -193,6 +200,11 @@ class WebSocketStream {
     // 2. Let reason be closeInfo [" reason "].
     const reason = closeInfo.reason
 
+    // The closing handshake needs to read the peer's Close frame off the
+    // socket to complete; a socket paused for receive backpressure (#5503)
+    // would otherwise never see it and the closed promise would hang forever.
+    this.#resumeSocketForClose()
+
     // 3. Close the WebSocket with this , code , and reason .
     closeWebSocketConnection(this.#handler, code, reason, true)
   }
@@ -291,6 +303,7 @@ class WebSocketStream {
       start: (controller) => {
         this.#readableStreamController = controller
       },
+      pull: () => this.#resumeSocketIfNeeded(),
       cancel: (reason) => this.#cancel(reason)
     })
 
@@ -350,6 +363,35 @@ class WebSocketStream {
     this.#readableStreamController.enqueue(chunk)
 
     // 4. Apply backpressure to the WebSocket.
+    this.#applyBackpressureToWebSocket()
+  }
+
+  #applyBackpressureToWebSocket () {
+    if (this.#readableStreamController.desiredSize <= 0) {
+      this.#handler.socket.pause()
+    }
+  }
+
+  #resumeSocketIfNeeded () {
+    if (
+      !this.#parserBackpressured &&
+      this.#handler.socket != null &&
+      !this.#handler.socket.destroyed &&
+      this.#readableStreamController != null &&
+      this.#readableStreamController.desiredSize > 0
+    ) {
+      this.#handler.socket.resume()
+    }
+  }
+
+  // Once closing starts there is no more reader to protect from backpressure,
+  // and the closing handshake needs to keep reading off the socket to see the
+  // peer's Close frame - so unconditionally resume, ignoring both backpressure
+  // sources above.
+  #resumeSocketForClose () {
+    if (this.#handler.socket != null && !this.#handler.socket.destroyed) {
+      this.#handler.socket.resume()
+    }
   }
 
   /** @type {import('../websocket').Handler['onSocketClose']} */
@@ -438,6 +480,9 @@ class WebSocketStream {
       reasonString = reason.reason
     }
 
+    // See close() for why this needs to happen before closeWebSocketConnection.
+    this.#resumeSocketForClose()
+
     // 4. Close the WebSocket with stream , code , and reasonString . If this throws an exception,
     //    discard code and reasonString and close the WebSocket with stream .
     closeWebSocketConnection(this.#handler, code, reasonString)

test/websocket/stream/issue-5503.js

@@ -0,0 +1,171 @@
+'use strict'
+
+const { test } = require('node:test')
+const net = require('node:net')
+const { WebSocketServer } = require('ws')
+
+const { WebSocketStream } = require('../../..')
+
+function waitFor (predicate) {
+  return new Promise((resolve, reject) => {
+    const started = Date.now()
+
+    const check = () => {
+      if (predicate()) {
+        resolve()
+      } else if (Date.now() - started > 1000) {
+        reject(new Error('timed out waiting for condition'))
+      } else {
+        setTimeout(check, 10)
+      }
+    }
+
+    check()
+  })
+}
+
+// These tests assert that `closed` does NOT hang - if it regresses, the
+// point of the test is defeated by also hanging (up to the runner's own
+// timeout) instead of failing fast with a clear message.
+function withTimeout (promise, ms, message) {
+  let timer
+  const timeout = new Promise((resolve, reject) => {
+    timer = setTimeout(() => reject(new Error(message)), ms)
+  })
+
+  return Promise.race([promise, timeout]).finally(() => clearTimeout(timer))
+}
+
+test('WebSocketStream applies receive backpressure to the socket', async (t) => {
+  const server = new WebSocketServer({ port: 0 })
+
+  server.on('connection', (socket) => {
+    socket.send('one')
+  })
+
+  const { port } = server.address()
+  const originalPause = net.Socket.prototype.pause
+  const originalResume = net.Socket.prototype.resume
+  let pauseCount = 0
+  let resumeCount = 0
+
+  net.Socket.prototype.pause = function (...args) {
+    if (this.remotePort === port) {
+      pauseCount++
+    }
+
+    return originalPause.apply(this, args)
+  }
+
+  net.Socket.prototype.resume = function (...args) {
+    if (this.remotePort === port) {
+      resumeCount++
+    }
+
+    return originalResume.apply(this, args)
+  }
+
+  t.after(() => {
+    net.Socket.prototype.pause = originalPause
+    net.Socket.prototype.resume = originalResume
+
+    for (const client of server.clients) {
+      client.terminate()
+    }
+
+    server.close()
+  })
+
+  const stream = new WebSocketStream(`ws://127.0.0.1:${port}`)
+  const { readable } = await stream.opened
+
+  await waitFor(() => pauseCount > 0)
+
+  const resumesBeforeRead = resumeCount
+  const reader = readable.getReader()
+  const first = await reader.read()
+
+  t.assert.deepStrictEqual(first, { done: false, value: 'one' })
+
+  await waitFor(() => resumeCount > resumesBeforeRead)
+
+  await reader.cancel()
+  await stream.closed.catch(() => {})
+})
+
+test('WebSocketStream#close resolves closed even while the socket is paused for receive backpressure', async (t) => {
+  const server = new WebSocketServer({ port: 0 })
+
+  server.on('connection', (socket) => {
+    socket.send('one')
+    socket.send('two')
+  })
+
+  const { port } = server.address()
+  const originalPause = net.Socket.prototype.pause
+  let pauseCount = 0
+
+  net.Socket.prototype.pause = function (...args) {
+    if (this.remotePort === port) {
+      pauseCount++
+    }
+
+    return originalPause.apply(this, args)
+  }
+
+  t.after(() => {
+    net.Socket.prototype.pause = originalPause
+
+    for (const client of server.clients) {
+      client.terminate()
+    }
+
+    server.close()
+  })
+
+  const stream = new WebSocketStream(`ws://127.0.0.1:${port}`)
+  await stream.opened
+
+  // Never read from `readable` - wait until the unread backlog has actually
+  // paused the socket (desiredSize <= 0) before closing, otherwise closing
+  // before the messages even arrive wouldn't exercise the paused-socket
+  // path at all. Closing must still complete: it needs to read the peer's
+  // own Close frame off that same paused socket.
+  await waitFor(() => pauseCount > 0)
+  stream.close()
+
+  const result = await withTimeout(stream.closed, 2000, 'stream.closed hung after close() with an unread backlog')
+  t.assert.strictEqual(result.closeCode, 1005)
+})
+
+test('WebSocketStream reader#cancel resolves closed even while the socket is paused for receive backpressure', async (t) => {
+  const server = new WebSocketServer({ port: 0 })
+
+  server.on('connection', (socket) => {
+    socket.send('one')
+    socket.send('two')
+  })
+
+  const { port } = server.address()
+
+  t.after(() => {
+    for (const client of server.clients) {
+      client.terminate()
+    }
+
+    server.close()
+  })
+
+  const stream = new WebSocketStream(`ws://127.0.0.1:${port}`)
+  const { readable } = await stream.opened
+  const reader = readable.getReader()
+
+  // Read the first message so the reader is locked and has seen data, then
+  // leave the second one unread - the socket stays paused. Cancelling must
+  // still complete rather than hang waiting for the peer's Close frame.
+  await reader.read()
+  await reader.cancel()
+
+  const result = await withTimeout(stream.closed, 2000, 'stream.closed hung after reader.cancel() with an unread backlog')
+  t.assert.strictEqual(result.closeCode, 1005)
+})