Skip to content

Commit 1089808

Browse files
authored
fix(websocket): support process.unref (#5578)
1 parent 05ec71b commit 1089808

2 files changed

Lines changed: 82 additions & 0 deletions

File tree

‎lib/web/websocket/websocket.js‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ const { SendQueue } = require('./sender')
2525
const { WebsocketFrameSend } = require('./frame')
2626
const { channels } = require('../../core/diagnostics')
2727

28+
const kRef = Symbol.for('nodejs.ref')
29+
const kUnref = Symbol.for('nodejs.unref')
30+
2831
function getSocketAddress (socket) {
2932
if (typeof socket?.address === 'function') {
3033
return socket.address()
@@ -68,6 +71,7 @@ class WebSocket extends EventTarget {
6871
#bufferedAmount = 0
6972
#protocol = ''
7073
#extensions = ''
74+
#refed = true
7175

7276
/** @type {SendQueue} */
7377
#sendQueue
@@ -194,6 +198,20 @@ class WebSocket extends EventTarget {
194198
this.#binaryType = 'blob'
195199
}
196200

201+
[kRef] () {
202+
webidl.brandCheck(this, WebSocket)
203+
204+
this.#refed = true
205+
this.#handler.socket?.ref?.()
206+
}
207+
208+
[kUnref] () {
209+
webidl.brandCheck(this, WebSocket)
210+
211+
this.#refed = false
212+
this.#handler.socket?.unref?.()
213+
}
214+
197215
/**
198216
* @see https://websockets.spec.whatwg.org/#dom-websocket-close
199217
* @param {number|undefined} code
@@ -468,6 +486,10 @@ class WebSocket extends EventTarget {
468486
// once this happens, the connection is open
469487
this.#handler.socket = response.socket
470488

489+
if (!this.#refed) {
490+
this.#handler.socket.unref?.()
491+
}
492+
471493
// Get options from dispatcher options
472494
const maxFragments = this.#handler.controller.dispatcher?.webSocketOptions?.maxFragments
473495
const maxPayloadSize = this.#handler.controller.dispatcher?.webSocketOptions?.maxPayloadSize

‎test/websocket/process-ref.js‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
'use strict'
2+
3+
const { test } = require('node:test')
4+
const { spawn } = require('node:child_process')
5+
const { once } = require('node:events')
6+
const { join } = require('node:path')
7+
const { WebSocketServer } = require('ws')
8+
9+
function waitForExit (child, timeout) {
10+
return new Promise((resolve, reject) => {
11+
const timer = setTimeout(() => {
12+
reject(new Error('Child process did not exit after WebSocket was unrefed'))
13+
}, timeout)
14+
15+
child.once('error', (error) => {
16+
clearTimeout(timer)
17+
reject(error)
18+
})
19+
20+
child.once('exit', (code, signal) => {
21+
clearTimeout(timer)
22+
resolve({ code, signal })
23+
})
24+
})
25+
}
26+
27+
test('process.unref allows the process to exit with an open WebSocket', async (t) => {
28+
const server = new WebSocketServer({ port: 0 })
29+
let connected = false
30+
server.once('connection', () => {
31+
connected = true
32+
})
33+
await once(server, 'listening')
34+
35+
t.after(() => server.close())
36+
37+
const url = `ws://127.0.0.1:${server.address().port}`
38+
const undici = join(__dirname, '../..')
39+
const child = spawn(process.execPath, ['-e', `
40+
const { WebSocket } = require(${JSON.stringify(undici)})
41+
const ws = new WebSocket(${JSON.stringify(url)})
42+
if (typeof ws[Symbol.for('nodejs.ref')] !== 'function' ||
43+
typeof ws[Symbol.for('nodejs.unref')] !== 'function') {
44+
throw new Error('WebSocket does not implement the Refable protocol')
45+
}
46+
ws.addEventListener('open', () => {
47+
process.unref(ws)
48+
process.ref(ws)
49+
process.unref(ws)
50+
})
51+
`], { stdio: 'ignore' })
52+
53+
t.after(() => child.kill())
54+
55+
const { code, signal } = await waitForExit(child, 5000)
56+
57+
t.assert.strictEqual(connected, true)
58+
t.assert.strictEqual(code, 0)
59+
t.assert.strictEqual(signal, null)
60+
})

0 commit comments

Comments
 (0)