Ensure forcefulShutdown and gracefulShutdown yield the same promises when called a second time

Author: benjie

src/main.ts

@@ -14,6 +14,7 @@ import {
   FailJobFunction,
   GetJobFunction,
   Job,
+  PromiseOrDirect,
   RunOnceOptions,
   TaskList,
   WorkerEventMap,
@@ -709,6 +710,9 @@ export function _runTaskList(
   // Make sure Node doesn't get upset about unhandled rejection
   abortPromise.then(null, () => /* noop */ void 0);
 
+  let gracefulShutdownPromise: PromiseOrDirect<void>;
+  let forcefulShutdownPromise: PromiseOrDirect<void>;
+
   // This is a representation of us that can be interacted with externally
   const workerPool: WorkerPool = {
     // "otpool" - "one time pool"
@@ -748,22 +752,17 @@ export function _runTaskList(
         logger.error(
           `gracefulShutdown called when forcefulShutdown is already in progress`,
         );
-        return;
+        return forcefulShutdownPromise;
       }
       if (workerPool._shuttingDown) {
         logger.error(
           `gracefulShutdown called when gracefulShutdown is already in progress`,
         );
-        return;
+        return gracefulShutdownPromise;
       }
-      workerPool._shuttingDown = true;
-
-      const abortTimer = setTimeout(() => {
-        abortController.abort();
-      }, gracefulShutdownAbortTimeout);
-      abortTimer.unref();
 
-      return middleware.run(
+      workerPool._shuttingDown = true;
+      gracefulShutdownPromise = middleware.run(
         "poolGracefulShutdown",
         { ctx: compiledSharedOptions, workerPool, message },
         async ({ message }) => {
@@ -872,6 +871,13 @@ export function _runTaskList(
           }
         },
       );
+
+      const abortTimer = setTimeout(() => {
+        abortController.abort();
+      }, gracefulShutdownAbortTimeout);
+      abortTimer.unref();
+
+      return gracefulShutdownPromise;
     },
 
     /**
@@ -882,13 +888,14 @@ export function _runTaskList(
         logger.error(
           `forcefulShutdown called when forcefulShutdown is already in progress`,
         );
-        return;
+        return forcefulShutdownPromise;
       }
-      return middleware.run(
+
+      workerPool._forcefulShuttingDown = true;
+      forcefulShutdownPromise = middleware.run(
         "poolForcefulShutdown",
         { ctx: compiledSharedOptions, workerPool: this, message },
         async ({ message }) => {
-          workerPool._forcefulShuttingDown = true;
           events.emit("pool:forcefulShutdown", {
             pool: workerPool,
             workerPool,
@@ -970,6 +977,8 @@ export function _runTaskList(
           }
         },
       );
+
+      return forcefulShutdownPromise;
     },
 
     promise,