handle graceful shutdown to wait for open workers message channel before disconnecting

Author: AbanoubGhadban

packages/react-on-rails-pro-node-renderer/src/worker/handleGracefulShutdown.ts

@@ -2,48 +2,121 @@ import cluster from 'cluster';
 import { FastifyInstance } from './types.js';
 import { SHUTDOWN_WORKER_MESSAGE } from '../shared/utils.js';
 import log from '../shared/log.js';
+import { onMessageEnded, onMessageInitiated } from '../workerMessagesRouter.js';
+
+type GracefulShutdownController = {
+  readonly activeRequestsCount: number;
+  readonly activeMessageChannelsCount: number;
+  readonly isShuttingDown: boolean;
+  markNewRequestReceived: () => void;
+  markRequestHandled: () => void;
+  markNewMessageChannelOpened: () => void;
+  markMessageChannelClosed: () => void;
+};
+
+let gracefulShutdownController: GracefulShutdownController;
+
+const setupGracefulShutdownHandler = () => {
+  if (gracefulShutdownController) {
+    return gracefulShutdownController;
+  }
 
-const handleGracefulShutdown = (app: FastifyInstance) => {
   const { worker } = cluster;
   if (!worker) {
-    log.error('handleGracefulShutdown is called on master, expected to call it on worker only');
-    return;
+    log.error('setupGracefulShutdownHandler is called on master, expected to call it on worker only');
+    return undefined;
   }
 
   let activeRequestsCount = 0;
+  let activeMessageChannelsCount = 0;
   let isShuttingDown = false;
 
+  const handleCloseEvent = () => {
+    if (!isShuttingDown) {
+      return;
+    }
+
+    if (activeMessageChannelsCount > 0) {
+      log.info(
+        'Worker #%d has "%d" active message channels, keep the worker connected',
+        worker.id,
+        activeMessageChannelsCount,
+      );
+    } else if (activeRequestsCount > 0) {
+      log.info(
+        'Worker #%d has "%d" active requests, disconnecting the worker',
+        worker.id,
+        activeRequestsCount,
+      );
+      worker.disconnect();
+    } else {
+      log.info('Worker #%d has no active requests, killing the worker', worker.id);
+      worker.destroy();
+    }
+  };
+
+  gracefulShutdownController = {
+    get activeRequestsCount() {
+      return activeRequestsCount;
+    },
+    get activeMessageChannelsCount() {
+      return activeMessageChannelsCount;
+    },
+    get isShuttingDown() {
+      return isShuttingDown;
+    },
+    markNewRequestReceived: () => {
+      activeRequestsCount += 1;
+    },
+    markRequestHandled: () => {
+      activeRequestsCount -= 1;
+      handleCloseEvent();
+    },
+    markNewMessageChannelOpened: () => {
+      activeMessageChannelsCount += 1;
+    },
+    markMessageChannelClosed: () => {
+      activeMessageChannelsCount -= 1;
+      handleCloseEvent();
+    },
+  };
+
   process.on('message', (msg) => {
     if (msg === SHUTDOWN_WORKER_MESSAGE) {
-      log.debug('Worker #%d received graceful shutdown message', worker.id);
+      log.info('Worker #%d received graceful shutdown message', worker.id);
       isShuttingDown = true;
-      if (activeRequestsCount === 0) {
-        log.debug('Worker #%d has no active requests, killing the worker', worker.id);
-        worker.destroy();
-      } else {
-        log.debug(
-          'Worker #%d has "%d" active requests, disconnecting the worker',
-          worker.id,
-          activeRequestsCount,
-        );
-        worker.disconnect();
-      }
+      handleCloseEvent();
     }
   });
 
+  return gracefulShutdownController;
+};
+
+const handleGracefulShutdown = (app: FastifyInstance) => {
+  const controller = setupGracefulShutdownHandler();
+  if (!controller) {
+    return;
+  }
+
   app.addHook('onRequest', (_req, _reply, done) => {
-    activeRequestsCount += 1;
+    controller.markNewRequestReceived();
     done();
   });
 
   app.addHook('onResponse', (_req, _reply, done) => {
-    activeRequestsCount -= 1;
-    if (isShuttingDown && activeRequestsCount === 0) {
-      log.debug('Worker #%d served all active requests and going to be killed', worker.id);
-      worker.destroy();
-    }
+    controller.markRequestHandled();
     done();
   });
 };
 
+onMessageInitiated(() => {
+  const controller = setupGracefulShutdownHandler();
+  controller?.markNewMessageChannelOpened();
+});
+
+onMessageEnded(() => {
+  const controller = setupGracefulShutdownHandler();
+  controller?.markMessageChannelClosed();
+});
+
 export default handleGracefulShutdown;

packages/react-on-rails-pro-node-renderer/src/workerMessagesRouter.ts

@@ -25,17 +25,21 @@ export const routeMessagesFromWorker = (worker: Worker) => {
     if (!isInterWorkerMessage(msg)) {
       return;
     }
-    const { workers } = cluster;
-    if (!workers) {
+    const workers = Object.values(cluster.workers ?? {}).filter((w) => {
+      return worker.id === w?.id || (w?.isConnected() && !w.isScheduledRestart);
+    });
+    if (workers.length === 0) {
       return;
     }
 
     const { requestId, messageType } = msg;
     if (messageType === 'initiate') {
-      const workerIds = Object.keys(workers);
-      const workersCount = workerIds.length;
-      const workerIndex = workerIds.indexOf(worker.id.toString());
+      const workersCount = workers.length;
+      const workerIndex = workers.findIndex((w) => w?.id === worker.id);
       const otherWorker = workers[workersCount - 1 - workerIndex];
+      if (!otherWorker) {
+        throw new Error("Can't find a worker to forward the message to");
+      }
 
       if (!otherWorker) {
         return;
@@ -65,6 +69,17 @@ type ReceivedMessage = {
   reply: (payload: unknown, close?: boolean) => void;
 };
 
+const onMessageInitiatedCallbacks: (() => void)[] = [];
+const onMessageEndedCallbacks: (() => void)[] = [];
+
+export const onMessageInitiated = (callback: () => void) => {
+  onMessageInitiatedCallbacks.push(callback);
+};
+
+export const onMessageEnded = (callback: () => void) => {
+  onMessageEndedCallbacks.push(callback);
+};
+
 export const onMessageReceived = (callback: (receivedMessage: ReceivedMessage) => void) => {
   process.on('message', (msg) => {
     if (!isInterWorkerMessage(msg)) {
@@ -83,8 +98,18 @@ export const onMessageReceived = (callback: (receivedMessage: ReceivedMessage) =
             requestId,
           };
           process.send?.(replyMsg);
+
+          if (close) {
+            onMessageEndedCallbacks.forEach((onMessageEndedCallback) => {
+              setTimeout(onMessageEndedCallback, 0);
+            });
+          }
         },
       });
+
+      onMessageInitiatedCallbacks.forEach((onMessageInitiatedCallback) => {
+        setTimeout(onMessageInitiatedCallback, 0);
+      });
     }
   });
 };
@@ -93,8 +118,14 @@ export const sendMessage = (initialPayload: unknown): Readable => {
   const requestId = randomUUID();
   const reveivedStream = new PassThrough();
 
+  onMessageInitiatedCallbacks.forEach((onMessageInitiatedCallback) => {
+    setTimeout(onMessageInitiatedCallback, 0);
+  });
+
   process.on('message', function messageReceivedCallback(msg) {
-    if (!isInterWorkerMessage(msg) || msg.requestId !== requestId) {
+    // If number of workers are small (1 or 2), the request can be redirected to the same worker
+    // So, ignore initiate message, it should be handled by the onMessageReceived not the sendMessage function
+    if (!isInterWorkerMessage(msg) || msg.requestId !== requestId || msg.messageType === 'initiate') {
       return;
     }
 
@@ -105,6 +136,9 @@ export const sendMessage = (initialPayload: unknown): Readable => {
     if (messageType === 'end') {
       reveivedStream.push(null);
       process.off('message', messageReceivedCallback);
+      onMessageEndedCallbacks.forEach((onMessageEndedCallback) => {
+        setTimeout(onMessageEndedCallback, 0);
+      });
     }
   });