Refactor reference counting with IPC (#1089)

lib/ipc/buffer-messages.js

@@ -25,6 +25,7 @@ export const waitForIpcOutput = async ({
 
 	for await (const message of loopOnMessages({
 		anyProcess: subprocess,
+		channel: subprocess.channel,
 		isSubprocess: false,
 		ipc,
 		shouldAwait: false,

lib/ipc/forward.js

@@ -5,7 +5,7 @@ import {undoAddedReferences} from './reference.js';
 // Forward the `message` and `disconnect` events from the process and subprocess to a proxy emitter.
 // This prevents the `error` event from stopping IPC.
 // This also allows debouncing the `message` event.
-export const getIpcEmitter = (anyProcess, isSubprocess) => {
+export const getIpcEmitter = (anyProcess, channel, isSubprocess) => {
 	if (IPC_EMITTERS.has(anyProcess)) {
 		return IPC_EMITTERS.get(anyProcess);
 	}
@@ -15,7 +15,12 @@ export const getIpcEmitter = (anyProcess, isSubprocess) => {
 	const ipcEmitter = new EventEmitter();
 	ipcEmitter.connected = true;
 	IPC_EMITTERS.set(anyProcess, ipcEmitter);
-	forwardEvents(ipcEmitter, anyProcess, isSubprocess);
+	forwardEvents({
+		ipcEmitter,
+		anyProcess,
+		channel,
+		isSubprocess,
+	});
 	return ipcEmitter;
 };
 
@@ -24,11 +29,16 @@ const IPC_EMITTERS = new WeakMap();
 // The `message` and `disconnect` events are buffered in the subprocess until the first listener is setup.
 // However, unbuffering happens after one tick, so this give enough time for the caller to setup the listener on the proxy emitter first.
 // See https://github.com/nodejs/node/blob/2aaeaa863c35befa2ebaa98fb7737ec84df4d8e9/lib/internal/child_process.js#L721
-const forwardEvents = (ipcEmitter, anyProcess, isSubprocess) => {
+const forwardEvents = ({ipcEmitter, anyProcess, channel, isSubprocess}) => {
 	const boundOnMessage = onMessage.bind(undefined, anyProcess, ipcEmitter);
 	anyProcess.on('message', boundOnMessage);
-	anyProcess.once('disconnect', onDisconnect.bind(undefined, {anyProcess, ipcEmitter, boundOnMessage}));
-	undoAddedReferences(anyProcess, isSubprocess);
+	anyProcess.once('disconnect', onDisconnect.bind(undefined, {
+		anyProcess,
+		channel,
+		ipcEmitter,
+		boundOnMessage,
+	}));
+	undoAddedReferences(channel, isSubprocess);
 };
 
 // Check whether there might still be some `message` events to receive

lib/ipc/get-each.js

@@ -4,28 +4,30 @@ import {getIpcEmitter, isConnected} from './forward.js';
 import {addReference, removeReference} from './reference.js';
 
 // Like `[sub]process.on('message')` but promise-based
-export const getEachMessage = ({anyProcess, isSubprocess, ipc}) => loopOnMessages({
+export const getEachMessage = ({anyProcess, channel, isSubprocess, ipc}) => loopOnMessages({
 	anyProcess,
+	channel,
 	isSubprocess,
 	ipc,
 	shouldAwait: !isSubprocess,
 });
 
 // Same but used internally
-export const loopOnMessages = ({anyProcess, isSubprocess, ipc, shouldAwait}) => {
+export const loopOnMessages = ({anyProcess, channel, isSubprocess, ipc, shouldAwait}) => {
 	validateIpcMethod({
 		methodName: 'getEachMessage',
 		isSubprocess,
 		ipc,
 		isConnected: isConnected(anyProcess),
 	});
 
-	addReference(anyProcess);
-	const ipcEmitter = getIpcEmitter(anyProcess, isSubprocess);
+	addReference(channel);
+	const ipcEmitter = getIpcEmitter(anyProcess, channel, isSubprocess);
 	const controller = new AbortController();
 	stopOnDisconnect(anyProcess, ipcEmitter, controller);
 	return iterateOnMessages({
 		anyProcess,
+		channel,
 		ipcEmitter,
 		isSubprocess,
 		shouldAwait,
@@ -40,14 +42,14 @@ const stopOnDisconnect = async (anyProcess, ipcEmitter, controller) => {
 	} catch {}
 };
 
-const iterateOnMessages = async function * ({anyProcess, ipcEmitter, isSubprocess, shouldAwait, controller}) {
+const iterateOnMessages = async function * ({anyProcess, channel, ipcEmitter, isSubprocess, shouldAwait, controller}) {
 	try {
 		for await (const [message] of on(ipcEmitter, 'message', {signal: controller.signal})) {
 			yield message;
 		}
 	} catch {} finally {
 		controller.abort();
-		removeReference(anyProcess);
+		removeReference(channel);
 
 		if (!isSubprocess) {
 			disconnect(anyProcess);

lib/ipc/get-one.js

@@ -4,20 +4,25 @@ import {getIpcEmitter, isConnected} from './forward.js';
 import {addReference, removeReference} from './reference.js';
 
 // Like `[sub]process.once('message')` but promise-based
-export const getOneMessage = ({anyProcess, isSubprocess, ipc}, {filter} = {}) => {
+export const getOneMessage = ({anyProcess, channel, isSubprocess, ipc}, {filter} = {}) => {
 	validateIpcMethod({
 		methodName: 'getOneMessage',
 		isSubprocess,
 		ipc,
 		isConnected: isConnected(anyProcess),
 	});
 
-	return getOneMessageAsync(anyProcess, isSubprocess, filter);
+	return getOneMessageAsync({
+		anyProcess,
+		channel,
+		isSubprocess,
+		filter,
+	});
 };
 
-const getOneMessageAsync = async (anyProcess, isSubprocess, filter) => {
-	addReference(anyProcess);
-	const ipcEmitter = getIpcEmitter(anyProcess, isSubprocess);
+const getOneMessageAsync = async ({anyProcess, channel, isSubprocess, filter}) => {
+	addReference(channel);
+	const ipcEmitter = getIpcEmitter(anyProcess, channel, isSubprocess);
 	const controller = new AbortController();
 	try {
 		return await Promise.race([
@@ -26,7 +31,7 @@ const getOneMessageAsync = async (anyProcess, isSubprocess, filter) => {
 		]);
 	} finally {
 		controller.abort();
-		removeReference(anyProcess);
+		removeReference(channel);
 	}
 };
 

lib/ipc/incoming.js

@@ -1,6 +1,7 @@
 import {once} from 'node:events';
 import {scheduler} from 'node:timers/promises';
 import {waitForOutgoingMessages} from './outgoing.js';
+import {redoAddedReferences} from './reference.js';
 
 // Debounce the `message` event so it is emitted at most once per macrotask.
 // This allows users to call `await getOneMessage()`/`getEachMessage()` multiple times in a row.
@@ -28,14 +29,15 @@ export const onMessage = async (anyProcess, ipcEmitter, message) => {
 const INCOMING_MESSAGES = new WeakMap();
 
 // If the `message` event is currently debounced, the `disconnect` event must wait for it
-export const onDisconnect = async ({anyProcess, ipcEmitter, boundOnMessage}) => {
+export const onDisconnect = async ({anyProcess, channel, ipcEmitter, boundOnMessage}) => {
 	const incomingMessages = INCOMING_MESSAGES.get(anyProcess);
 	while (incomingMessages?.length > 0) {
 		// eslint-disable-next-line no-await-in-loop
 		await once(ipcEmitter, 'message');
 	}
 
 	anyProcess.removeListener('message', boundOnMessage);
+	redoAddedReferences(channel);
 	ipcEmitter.connected = false;
 	ipcEmitter.emit('disconnect');
 };

lib/ipc/methods.js

@@ -17,14 +17,25 @@ const getIpcMethods = (anyProcess, isSubprocess, ipc) => {
 	const anyProcessSend = anyProcess.send === undefined
 		? undefined
 		: promisify(anyProcess.send.bind(anyProcess));
+	const {channel} = anyProcess;
 	return {
 		sendMessage: sendMessage.bind(undefined, {
 			anyProcess,
 			anyProcessSend,
 			isSubprocess,
 			ipc,
 		}),
-		getOneMessage: getOneMessage.bind(undefined, {anyProcess, isSubprocess, ipc}),
-		getEachMessage: getEachMessage.bind(undefined, {anyProcess, isSubprocess, ipc}),
+		getOneMessage: getOneMessage.bind(undefined, {
+			anyProcess,
+			channel,
+			isSubprocess,
+			ipc,
+		}),
+		getEachMessage: getEachMessage.bind(undefined, {
+			anyProcess,
+			channel,
+			isSubprocess,
+			ipc,
+		}),
 	};
 };

lib/ipc/reference.js

@@ -1,21 +1,32 @@
 // By default, Node.js keeps the subprocess alive while it has a `message` or `disconnect` listener.
 // We replicate the same logic for the events that we proxy.
-// This ensures the subprocess is kept alive while `sendMessage()`, `getOneMessage()` and `getEachMessage()` are ongoing.
+// This ensures the subprocess is kept alive while `getOneMessage()` and `getEachMessage()` are ongoing.
+// This is not a problem with `sendMessage()` since Node.js handles that method automatically.
+// We do not use `anyProcess.channel.ref()` since this would prevent the automatic `.channel.refCounted()` Node.js is doing.
+// We keep a reference to `anyProcess.channel` since it might be `null` while `getOneMessage()` or `getEachMessage()` is still processing debounced messages.
 // See https://github.com/nodejs/node/blob/2aaeaa863c35befa2ebaa98fb7737ec84df4d8e9/lib/internal/child_process.js#L547
-export const addReference = anyProcess => {
-	anyProcess.channel?.refCounted();
+export const addReference = channel => {
+	channel.refCounted();
 };
 
-export const removeReference = anyProcess => {
-	anyProcess.channel?.unrefCounted();
+export const removeReference = channel => {
+	channel.unrefCounted();
 };
 
 // To proxy events, we setup some global listeners on the `message` and `disconnect` events.
 // Those should not keep the subprocess alive, so we remove the automatic counting that Node.js is doing.
 // See https://github.com/nodejs/node/blob/1b965270a9c273d4cf70e8808e9d28b9ada7844f/lib/child_process.js#L180
-export const undoAddedReferences = (anyProcess, isSubprocess) => {
+export const undoAddedReferences = (channel, isSubprocess) => {
 	if (isSubprocess) {
-		removeReference(anyProcess);
-		removeReference(anyProcess);
+		removeReference(channel);
+		removeReference(channel);
+	}
+};
+
+// Reverse it during `disconnect`
+export const redoAddedReferences = (channel, isSubprocess) => {
+	if (isSubprocess) {
+		addReference(channel);
+		addReference(channel);
 	}
 };

lib/ipc/send.js

@@ -5,7 +5,6 @@ import {
 	disconnect,
 } from './validation.js';
 import {startSendMessage, endSendMessage} from './outgoing.js';
-import {addReference, removeReference} from './reference.js';
 
 // Like `[sub]process.send()` but promise-based.
 // We do not `await subprocess` during `.sendMessage()` nor `.getOneMessage()` since those methods are transient.
@@ -28,7 +27,6 @@ export const sendMessage = ({anyProcess, anyProcessSend, isSubprocess, ipc}, mes
 };
 
 const sendMessageAsync = async ({anyProcess, anyProcessSend, isSubprocess, message}) => {
-	addReference(anyProcess);
 	const outgoingMessagesState = startSendMessage(anyProcess);
 	try {
 		await anyProcessSend(message);
@@ -39,6 +37,5 @@ const sendMessageAsync = async ({anyProcess, anyProcessSend, isSubprocess, messa
 		throw error;
 	} finally {
 		endSendMessage(outgoingMessagesState);
-		removeReference(anyProcess);
 	}
 };

test/fixtures/ipc-disconnect-get.js

@@ -0,0 +1,7 @@
+#!/usr/bin/env node
+import process from 'node:process';
+import {getOneMessage} from '../../index.js';
+
+process.disconnect();
+console.log(process.channel);
+await getOneMessage();

test/ipc/reference.js

@@ -100,3 +100,22 @@ test('process.once("disconnect") does not keep the subprocess alive, after getOn
 	t.deepEqual(ipcOutput, ['.']);
 	t.is(stdout, '.');
 });
+
+test('Can call subprocess.disconnect() right away', async t => {
+	const subprocess = execa('ipc-send.js', {ipc: true});
+	subprocess.disconnect();
+	t.is(subprocess.channel, null);
+
+	await t.throwsAsync(subprocess.getOneMessage(), {
+		message: /subprocess.getOneMessage\(\) could not complete/,
+	});
+	await t.throwsAsync(subprocess, {
+		message: /Error: sendMessage\(\) cannot be used/,
+	});
+});
+
+test('Can call process.disconnect() right away', async t => {
+	const {stdout, stderr} = await t.throwsAsync(execa('ipc-disconnect-get.js', {ipc: true}));
+	t.is(stdout, 'null');
+	t.true(stderr.includes('Error: getOneMessage() cannot be used'));
+});