Add async support (#20)

Co-authored-by: Sindre Sorhus <sindresorhus@gmail.com>

fixture-async.js

@@ -0,0 +1,32 @@
+import exitHook, {asyncExitHook, gracefulExit} from './index.js';
+
+exitHook(() => {
+	console.log('foo');
+});
+
+exitHook(() => {
+	console.log('bar');
+});
+
+const unsubscribe = exitHook(() => {
+	console.log('baz');
+});
+
+unsubscribe();
+
+asyncExitHook(
+	async () => {
+		await new Promise(resolve => {
+			setTimeout(() => {
+				resolve();
+			}, 100);
+		});
+
+		console.log('quux');
+	},
+	{
+		minimumWait: 200,
+	},
+);
+
+gracefulExit();

index.d.ts

@@ -3,7 +3,7 @@ Run some code when the process exits.
 
 The `process.on('exit')` event doesn't catch all the ways a process can exit.
 
-This package is useful for cleaning up before exiting.
+This is useful for cleaning synchronously before exiting.
 
 @param onExit - The callback function to execute when the process exits.
 @returns A function that removes the hook when called.
@@ -33,3 +33,61 @@ unsubscribe();
 ```
 */
 export default function exitHook(onExit: () => void): () => void;
+
+/**
+Run code asynchronously when the process exits.
+
+@see https://github.com/sindresorhus/exit-hook/blob/main/readme.md#asynchronous-exit-notes
+@param onExit - The callback function to execute when the process exits via `gracefulExit`, and will be wrapped in `Promise.resolve`.
+@returns A function that removes the hook when called.
+
+@example
+```
+import {asyncExitHook} from 'exit-hook';
+
+asyncExitHook(() => {
+	console.log('Exiting');
+}, {
+	minimumWait: 500
+});
+
+throw new Error('🦄');
+
+//=> 'Exiting'
+
+// Removing an exit hook:
+const unsubscribe = asyncExitHook(() => {}, {});
+
+unsubscribe();
+```
+*/
+export function asyncExitHook(onExit: () => (void | Promise<void>), options: Options): () => void;
+
+/**
+Exit the process and make a best-effort to complete all asynchronous hooks.
+
+If you are using `asyncExitHook`, consider using `gracefulExit()` instead of `process.exit()` to ensure all asynchronous tasks are given an opportunity to run.
+
+@param signal - The exit code to use. Same as the argument to `process.exit()`.
+@see https://github.com/sindresorhus/exit-hook/blob/main/readme.md#asynchronous-exit-notes
+
+@example
+```
+import {asyncExitHook, gracefulExit} from 'exit-hook';
+
+asyncExitHook(() => {
+	console.log('Exiting');
+}, 500);
+
+// Instead of `process.exit()`
+gracefulExit();
+```
+*/
+export function gracefulExit(signal?: number): void;
+
+export interface Options {
+	/**
+	The amount of time in milliseconds that the `onExit` function is expected to take.
+	*/
+	minimumWait: number;
+}

index.js

@@ -1,46 +1,132 @@
 import process from 'node:process';
 
+const asyncCallbacks = new Set();
 const callbacks = new Set();
+
 let isCalled = false;
 let isRegistered = false;
 
-function exit(shouldManuallyExit, signal) {
+async function exit(shouldManuallyExit, isSynchronous, signal) {
+	if (asyncCallbacks.size > 0 && isSynchronous) {
+		console.error([
+			'SYNCHRONOUS TERMINATION NOTICE:',
+			'When explicitly exiting the process via process.exit or via a parent process,',
+			'asynchronous tasks in your exitHooks will not run. Either remove these tasks,',
+			'use gracefulExit() instead of process.exit(), or ensure your parent process',
+			'sends a SIGINT to the process running this code.',
+		].join(' '));
+	}
+
 	if (isCalled) {
 		return;
 	}
 
 	isCalled = true;
 
+	const done = (force = false) => {
+		if (force === true || shouldManuallyExit === true) {
+			process.exit(128 + signal); // eslint-disable-line unicorn/no-process-exit
+		}
+	};
+
 	for (const callback of callbacks) {
 		callback();
 	}
 
-	if (shouldManuallyExit === true) {
-		process.exit(128 + signal); // eslint-disable-line unicorn/no-process-exit
+	if (isSynchronous) {
+		done();
+		return;
 	}
+
+	const promises = [];
+	let forceAfter = 0;
+	for (const [callback, wait] of asyncCallbacks) {
+		forceAfter = Math.max(forceAfter, wait);
+		promises.push(Promise.resolve(callback()));
+	}
+
+	// Force exit if we exceeded our wait value
+	const asyncTimer = setTimeout(() => {
+		done(true);
+	}, forceAfter);
+
+	await Promise.all(promises);
+	clearTimeout(asyncTimer);
+	done();
 }
 
-export default function exitHook(onExit) {
-	callbacks.add(onExit);
+function addHook(options) {
+	const {onExit, minimumWait, isSynchronous} = options;
+	const asyncCallbackConfig = [onExit, minimumWait];
+
+	if (isSynchronous) {
+		callbacks.add(onExit);
+	} else {
+		asyncCallbacks.add(asyncCallbackConfig);
+	}
 
 	if (!isRegistered) {
 		isRegistered = true;
 
-		process.once('exit', exit);
-		process.once('SIGINT', exit.bind(undefined, true, 2));
-		process.once('SIGTERM', exit.bind(undefined, true, 15));
+		// Exit cases that support asynchronous handling
+		process.once('beforeExit', exit.bind(undefined, true, false, 0));
+		process.once('SIGINT', exit.bind(undefined, true, false, 2));
+		process.once('SIGTERM', exit.bind(undefined, true, false, 15));
+
+		// Explicit exit events. Calling will force an immediate exit and run all
+		// synchronous hooks. Explicit exits must not extend the node process
+		// artificially. Will log errors if asynchronous calls exist.
+		process.once('exit', exit.bind(undefined, false, true, 0));
 
-		// PM2 Cluster shutdown message. Caught to support async handlers with pm2, needed because
-		// explicitly calling process.exit() doesn't trigger the beforeExit event, and the exit
-		// event cannot support async handlers, since the event loop is never called after it.
+		// PM2 Cluster shutdown message. Caught to support async handlers with pm2,
+		// needed because explicitly calling process.exit() doesn't trigger the
+		// beforeExit event, and the exit event cannot support async handlers,
+		// since the event loop is never called after it.
 		process.on('message', message => {
 			if (message === 'shutdown') {
-				exit(true, -128);
+				exit(true, true, -128);
 			}
 		});
 	}
 
 	return () => {
-		callbacks.delete(onExit);
+		if (isSynchronous) {
+			callbacks.delete(onExit);
+		} else {
+			asyncCallbacks.delete(asyncCallbackConfig);
+		}
 	};
 }
+
+function exitHook(onExit) {
+	if (typeof onExit !== 'function') {
+		throw new TypeError('onExit must be a function');
+	}
+
+	return addHook({
+		onExit,
+		isSynchronous: true,
+	});
+}
+
+export function asyncExitHook(onExit, options) {
+	if (typeof onExit !== 'function') {
+		throw new TypeError('onExit must be a function');
+	}
+
+	if (typeof options?.minimumWait !== 'number' || options.minimumWait <= 0) {
+		throw new TypeError('minimumWait must be set to a positive numeric value');
+	}
+
+	return addHook({
+		onExit,
+		minimumWait: options.minimumWait,
+		isSynchronous: false,
+	});
+}
+
+export function gracefulExit(signal = 0) {
+	exit(true, false, -128 + signal);
+}
+
+export default exitHook;

index.test-d.ts

@@ -1,7 +1,14 @@
 import {expectType} from 'tsd';
-import exitHook from './index.js';
+import exitHook, {asyncExitHook} from './index.js';
 
 const unsubscribe = exitHook(() => {}); // eslint-disable-line @typescript-eslint/no-empty-function
 
+const asyncUnsubscribe = asyncExitHook(async () => {}, // eslint-disable-line @typescript-eslint/no-empty-function
+	{minimumWait: 300},
+);
+
 expectType<() => void>(unsubscribe);
 unsubscribe();
+
+expectType<() => void>(asyncUnsubscribe);
+asyncUnsubscribe();

readme.md

@@ -46,14 +46,89 @@ unsubscribe();
 
 ### exitHook(onExit)
 
-Returns a function that removes the hook when called.
+Register a function to run during `process.exit`. Returns a function that removes the hook when called.
 
 #### onExit
 
-Type: `Function`
+Type: `function(): void`
 
 The callback function to execute when the process exits.
 
+### asyncExitHook(onExit, minimumWait)
+
+Register a function to run during `gracefulExit`. Returns a function that removes the hook when called.
+
+Please see [Async Notes](#async-notes) for considerations when using the asynchronous API.
+
+```js
+import {asyncExitHook} from 'exit-hook';
+
+asyncExitHook(async () => {
+	console.log('Exiting');
+}, 300);
+
+throw new Error('🦄');
+
+//=> 'Exiting'
+```
+
+Removing an asynchronous exit hook:
+
+```js
+import {asyncExitHook} from 'exit-hook';
+
+const unsubscribe = asyncExitHook(async () => {
+	console.log('Exiting');
+}, {
+	minimumWait: 300
+});
+
+unsubscribe();
+```
+
+#### onExit
+
+Type: `function(): void | Promise<void>`
+
+The callback function to execute when the process exits via `gracefulExit`, and will be wrapped in `Promise.resolve`.
+
+#### options
+
+##### minimumWait
+
+Type: `number`
+
+The amount of time in milliseconds that the `onExit` function is expected to take.
+
+### gracefulExit(signal?: number): void
+
+Exit the process and make a best-effort to complete all asynchronous hooks.
+
+If you are using `asyncExitHook`, consider using `gracefulExit()` instead of `process.exit()` to ensure all asynchronous tasks are given an opportunity to run.
+
+```js
+import {gracefulExit} from 'exit-hook';
+
+gracefulExit();
+```
+
+#### signal
+
+Type: `number`\
+Default: `0`
+
+The exit code to use. Same as the argument to `process.exit()`.
+
+## Asynchronous Exit Notes
+
+**tl;dr** If you have 100% control over how your process terminates, then you can swap `exitHook` and `process.exit` for `asyncExitHook` and `gracefulExit` respectively. Otherwise, keep reading to understand important tradeoffs if you're using `asyncExitHook`.
+
+Node.js does not offer an asynchronous shutdown API by default [#1](https://github.com/nodejs/node/discussions/29480#discussioncomment-99213) [#2](https://github.com/nodejs/node/discussions/29480#discussioncomment-99217), so `asyncExitHook` and `gracefulExit` will make a "best effort" attempt to shut down the process and run your asynchronous tasks.
+
+If you have asynchronous hooks registered and your node.js process is terminated in a synchronous manner, a `SYNCHRONOUS TERMINATION NOTICE` error will be logged to the console. To avoid this, ensure you're only exiting via `gracefulExit` or that an upstream process manager is sending a `SIGINT` or `SIGTERM` signal to node.js.
+
+Asynchronous hooks should make a "best effort" to perform their tasks within the `minimumWait` time, but also be written to assume they may not complete their tasks before termination.
+
 ---
 
 <div align="center">