From 3944ff3b52370e3903a3c8142ff59b64b2aa46e2 Mon Sep 17 00:00:00 2001
From: sinkhaha <1468709106@qq.com>
Date: Fri, 24 May 2024 09:44:08 +0800
Subject: [PATCH] feat: add setDisableRefork helper (#118)

closes https://github.com/node-modules/cfork/issues/117

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced the ability to dynamically disable reforking of worker
processes.
- Added functions to set and get the `disableRefork` status of worker
processes.

- **Documentation**
- Updated README with examples on how to use the new `setDisableRefork`
function.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 README.md                      |  7 +++++-
 fixtures/kill_worker/master.js |  2 +-
 index.js                       | 39 ++++++++++++++++++++++++++++------
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 94d7ec4..619c586 100644
--- a/README.md
+++ b/README.md
@@ -33,7 +33,7 @@ $ npm install cfork --save
 const cfork = require('cfork');
 const util = require('util');
 
-cfork({
+const clusterWorker = cfork({
   exec: '/your/app/worker.js',
   // slaves: ['/your/app/slave.js'],
   // count: require('os').cpus().length,
@@ -66,6 +66,11 @@ cfork({
 process.on('uncaughtException', err => {
   // do what you want
 });
+
+// if you want to dynamically disable refork, you can call the setDisableRefork, priority over the refork parameter
+cfork.setDisableRefork(clusterWorker, true);
+
+
 ```
 
 ### Options
diff --git a/fixtures/kill_worker/master.js b/fixtures/kill_worker/master.js
index d0f59dc..aee6d5f 100644
--- a/fixtures/kill_worker/master.js
+++ b/fixtures/kill_worker/master.js
@@ -52,7 +52,7 @@ http.createServer(function (req, res) {
   var count = 0;
   for (var id in cluster.workers) {
     var worker = cluster.workers[id];
-    worker.disableRefork = true;
+    cfork.setDisableRefork(worker, true);
     worker.process.kill('SIGTERM');
     count++;
   }
diff --git a/index.js b/index.js
index b91e878..fee62d4 100644
--- a/index.js
+++ b/index.js
@@ -9,6 +9,8 @@ var defer = global.setImmediate || process.nextTick;
 
 module.exports = fork;
 
+fork.setDisableRefork = setDisableRefork;
+
 /**
  * cluster fork
  *
@@ -89,20 +91,21 @@ function fork(options) {
   var unexpectedCount = 0;
 
   cluster.on('disconnect', function (worker) {
-    var log = console[worker.disableRefork ? 'info' : 'error'];
+    var disableRefork = getDisableRefork(worker);
+    var log = console[ disableRefork ? 'info' : 'error'];
     disconnectCount++;
     var isDead = worker.isDead && worker.isDead();
     var propertyName = worker.hasOwnProperty('exitedAfterDisconnect') ? 'exitedAfterDisconnect' : 'suicide';
     log('[%s] [cfork:master:%s] worker:%s disconnect (%s: %s, state: %s, isDead: %s, worker.disableRefork: %s)',
       utility.logDate(), process.pid, worker.process.pid, propertyName, worker[propertyName],
-      worker.state, isDead, worker.disableRefork);
+      worker.state, isDead, disableRefork);
     if (isDead) {
       // worker has terminated before disconnect
       log('[%s] [cfork:master:%s] don\'t fork, because worker:%s exit event emit before disconnect',
         utility.logDate(), process.pid, worker.process.pid);
       return;
     }
-    if (worker.disableRefork) {
+    if (disableRefork) {
       // worker has terminated by master, like egg-cluster master will set disableRefork to true
       log('[%s] [cfork:master:%s] don\'t fork, because worker:%s will be kill soon',
         utility.logDate(), process.pid, worker.process.pid);
@@ -123,19 +126,20 @@ function fork(options) {
   });
 
   cluster.on('exit', function (worker, code, signal) {
-    var log = console[worker.disableRefork ? 'info' : 'error'];
+    var disableRefork = getDisableRefork(worker);
+    var log = console[disableRefork ? 'info' : 'error'];
     var isExpected = !!disconnects[worker.process.pid];
     var isDead = worker.isDead && worker.isDead();
     var propertyName = worker.hasOwnProperty('exitedAfterDisconnect') ? 'exitedAfterDisconnect' : 'suicide';
     log('[%s] [cfork:master:%s] worker:%s exit (code: %s, %s: %s, state: %s, isDead: %s, isExpected: %s, worker.disableRefork: %s)',
       utility.logDate(), process.pid, worker.process.pid, code, propertyName, worker[propertyName],
-      worker.state, isDead, isExpected, worker.disableRefork);
+      worker.state, isDead, isExpected, disableRefork);
     if (isExpected) {
       delete disconnects[worker.process.pid];
       // worker disconnect first, exit expected
       return;
     }
-    if (worker.disableRefork) {
+    if (disableRefork) {
       // worker is killed by master
       return;
     }
@@ -291,3 +295,26 @@ function fork(options) {
     }
   }
 }
+
+/**
+ * set disableRefork
+ *
+ * @param {import('cluster').Worker} worker - worker, Node.js cluster.Worker object
+ * @param {boolean} isDisableRefork - the worker process is not refork
+ * @returns {void}
+ */
+function setDisableRefork(worker, isDisableRefork) {
+  if (worker) {
+    worker.disableRefork = isDisableRefork;
+  }
+}
+
+/**
+ * get disableRefork
+ *
+ * @param {import('cluster').Worker} worker - worker, Node.js cluster.Worker object
+ * @returns {boolean} the worker process is not refork
+ */
+function getDisableRefork(worker) {
+  return (worker && worker.disableRefork) || false;
+}