From decfc2ae5c1a5664eaf85f71829919cd4765fca9 Mon Sep 17 00:00:00 2001
From: rickyes <mail@zhoumq.cn>
Date: Wed, 29 Apr 2020 02:00:35 +0800
Subject: [PATCH] fs: add .ref() and .unref() methods to watcher classes

Backport-PR-URL: https://github.com/nodejs/node/pull/35555
PR-URL: https://github.com/nodejs/node/pull/33134
Fixes: https://github.com/nodejs/node/issues/33096
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Anna Henningsen <anna@addaleax.net>
---
 doc/api/fs.md                                | 67 ++++++++++++++++++++
 lib/fs.js                                    |  6 ++
 lib/internal/fs/watchers.js                  | 56 +++++++++++++++-
 src/fs_event_wrap.cc                         |  3 +-
 test/parallel/test-fs-watch-ref-unref.js     | 20 ++++++
 test/parallel/test-fs-watchfile-ref-unref.js | 35 ++++++++++
 tools/doc/type-parser.js                     |  1 +
 7 files changed, 184 insertions(+), 4 deletions(-)
 create mode 100644 test/parallel/test-fs-watch-ref-unref.js
 create mode 100644 test/parallel/test-fs-watchfile-ref-unref.js

diff --git a/doc/api/fs.md b/doc/api/fs.md
index 57af737a2f6444..f48f2e3b2ddf6a 100644
--- a/doc/api/fs.md
+++ b/doc/api/fs.md
@@ -598,6 +598,72 @@ added: v0.5.8
 Stop watching for changes on the given `fs.FSWatcher`. Once stopped, the
 `fs.FSWatcher` object is no longer usable.
 
+### `watcher.ref()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {fs.FSWatcher}
+
+When called, requests that the Node.js event loop *not* exit so long as the
+`FSWatcher` is active. Calling `watcher.ref()` multiple times will have
+no effect.
+
+By default, all `FSWatcher` objects are "ref'ed", making it normally
+unnecessary to call `watcher.ref()` unless `watcher.unref()` had been
+called previously.
+
+### `watcher.unref()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {fs.FSWatcher}
+
+When called, the active `FSWatcher` object will not require the Node.js
+event loop to remain active. If there is no other activity keeping the
+event loop running, the process may exit before the `FSWatcher` object's
+callback is invoked. Calling `watcher.unref()` multiple times will have
+no effect.
+
+## Class: `fs.StatWatcher`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Extends {EventEmitter}
+
+A successful call to `fs.watchFile()` method will return a new `fs.StatWatcher`
+object.
+
+### `watcher.ref()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {fs.StatWatcher}
+
+When called, requests that the Node.js event loop *not* exit so long as the
+`StatWatcher` is active. Calling `watcher.ref()` multiple times will have
+no effect.
+
+By default, all `StatWatcher` objects are "ref'ed", making it normally
+unnecessary to call `watcher.ref()` unless `watcher.unref()` had been
+called previously.
+
+### `watcher.unref()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {fs.StatWatcher}
+
+When called, the active `StatWatcher` object will not require the Node.js
+event loop to remain active. If there is no other activity keeping the
+event loop running, the process may exit before the `StatWatcher` object's
+callback is invoked. Calling `watcher.unref()` multiple times will have
+no effect.
+
 ## Class: `fs.ReadStream`
 <!-- YAML
 added: v0.1.93
@@ -4029,6 +4095,7 @@ changes:
 * `listener` {Function}
   * `current` {fs.Stats}
   * `previous` {fs.Stats}
+* Returns: {fs.StatWatcher}
 
 Watch for changes on `filename`. The callback `listener` will be called each
 time the file is accessed.
diff --git a/lib/fs.js b/lib/fs.js
index 03b5451c4c1123..eead8bb5a39176 100644
--- a/lib/fs.js
+++ b/lib/fs.js
@@ -1487,6 +1487,8 @@ function watchFile(filename, options, listener) {
     stat = new watchers.StatWatcher(options.bigint);
     stat.start(filename, options.persistent, options.interval);
     statWatchers.set(filename, stat);
+  } else {
+    stat[watchers.kFSStatWatcherAddOrCleanRef]('add');
   }
 
   stat.addListener('change', listener);
@@ -1501,9 +1503,13 @@ function unwatchFile(filename, listener) {
   if (stat === undefined) return;
 
   if (typeof listener === 'function') {
+    const beforeListenerCount = stat.listenerCount('change');
     stat.removeListener('change', listener);
+    if (stat.listenerCount('change') < beforeListenerCount)
+      stat[watchers.kFSStatWatcherAddOrCleanRef]('clean');
   } else {
     stat.removeAllListeners('change');
+    stat[watchers.kFSStatWatcherAddOrCleanRef]('cleanAll');
   }
 
   if (stat.listenerCount('change') === 0) {
diff --git a/lib/internal/fs/watchers.js b/lib/internal/fs/watchers.js
index 9539b46877ed3e..1ffe70323f3bf3 100644
--- a/lib/internal/fs/watchers.js
+++ b/lib/internal/fs/watchers.js
@@ -29,6 +29,10 @@ const assert = require('internal/assert');
 const kOldStatus = Symbol('kOldStatus');
 const kUseBigint = Symbol('kUseBigint');
 
+const KFSStatWatcherRefCount = Symbol('KFSStatWatcherRefCount');
+const KFSStatWatcherMaxRefCount = Symbol('KFSStatWatcherMaxRefCount');
+const kFSStatWatcherAddOrCleanRef = Symbol('kFSStatWatcherAddOrCleanRef');
+
 function emitStop(self) {
   self.emit('stop');
 }
@@ -39,6 +43,8 @@ function StatWatcher(bigint) {
   this._handle = null;
   this[kOldStatus] = -1;
   this[kUseBigint] = bigint;
+  this[KFSStatWatcherRefCount] = 1;
+  this[KFSStatWatcherMaxRefCount] = 1;
 }
 ObjectSetPrototypeOf(StatWatcher.prototype, EventEmitter.prototype);
 ObjectSetPrototypeOf(StatWatcher, EventEmitter);
@@ -70,7 +76,7 @@ StatWatcher.prototype.start = function(filename, persistent, interval) {
   this._handle[owner_symbol] = this;
   this._handle.onchange = onchange;
   if (!persistent)
-    this._handle.unref();
+    this.unref();
 
   // uv_fs_poll is a little more powerful than ev_stat but we curb it for
   // the sake of backwards compatibility
@@ -106,6 +112,41 @@ StatWatcher.prototype.stop = function() {
   this._handle = null;
 };
 
+// Clean up or add ref counters.
+StatWatcher.prototype[kFSStatWatcherAddOrCleanRef] = function(operate) {
+  if (operate === 'add') {
+    // Add a Ref
+    this[KFSStatWatcherRefCount]++;
+    this[KFSStatWatcherMaxRefCount]++;
+  } else if (operate === 'clean') {
+    // Clean up a single
+    this[KFSStatWatcherMaxRefCount]--;
+    this.unref();
+  } else if (operate === 'cleanAll') {
+    // Clean up all
+    this[KFSStatWatcherMaxRefCount] = 0;
+    this[KFSStatWatcherRefCount] = 0;
+    this._handle && this._handle.unref();
+  }
+};
+
+StatWatcher.prototype.ref = function() {
+  // Avoid refCount calling ref multiple times causing unref to have no effect.
+  if (this[KFSStatWatcherRefCount] === this[KFSStatWatcherMaxRefCount])
+    return this;
+  if (this._handle && this[KFSStatWatcherRefCount]++ === 0)
+    this._handle.ref();
+  return this;
+};
+
+StatWatcher.prototype.unref = function() {
+  // Avoid refCount calling unref multiple times causing ref to have no effect.
+  if (this[KFSStatWatcherRefCount] === 0) return this;
+  if (this._handle && --this[KFSStatWatcherRefCount] === 0)
+    this._handle.unref();
+  return this;
+};
+
 
 function FSWatcher() {
   EventEmitter.call(this);
@@ -193,6 +234,16 @@ FSWatcher.prototype.close = function() {
   process.nextTick(emitCloseNT, this);
 };
 
+FSWatcher.prototype.ref = function() {
+  if (this._handle) this._handle.ref();
+  return this;
+};
+
+FSWatcher.prototype.unref = function() {
+  if (this._handle) this._handle.unref();
+  return this;
+};
+
 function emitCloseNT(self) {
   self.emit('close');
 }
@@ -206,5 +257,6 @@ ObjectDefineProperty(FSEvent.prototype, 'owner', {
 
 module.exports = {
   FSWatcher,
-  StatWatcher
+  StatWatcher,
+  kFSStatWatcherAddOrCleanRef,
 };
diff --git a/src/fs_event_wrap.cc b/src/fs_event_wrap.cc
index 858455bff2d42d..e16dc84ce8e5ee 100644
--- a/src/fs_event_wrap.cc
+++ b/src/fs_event_wrap.cc
@@ -101,9 +101,8 @@ void FSEventWrap::Initialize(Local<Object> target,
       FSEventWrap::kInternalFieldCount);
   t->SetClassName(fsevent_string);
 
-  t->Inherit(AsyncWrap::GetConstructorTemplate(env));
+  t->Inherit(HandleWrap::GetConstructorTemplate(env));
   env->SetProtoMethod(t, "start", Start);
-  env->SetProtoMethod(t, "close", Close);
 
   Local<FunctionTemplate> get_initialized_templ =
       FunctionTemplate::New(env->isolate(),
diff --git a/test/parallel/test-fs-watch-ref-unref.js b/test/parallel/test-fs-watch-ref-unref.js
new file mode 100644
index 00000000000000..e51ecaf5b2fcc1
--- /dev/null
+++ b/test/parallel/test-fs-watch-ref-unref.js
@@ -0,0 +1,20 @@
+'use strict';
+
+const common = require('../common');
+
+if (common.isIBMi)
+  common.skip('IBMi does not support `fs.watch()`');
+
+const fs = require('fs');
+
+const watcher = fs.watch(__filename, common.mustNotCall());
+
+watcher.unref();
+
+setTimeout(
+  common.mustCall(() => {
+    watcher.ref();
+    watcher.unref();
+  }),
+  common.platformTimeout(100)
+);
diff --git a/test/parallel/test-fs-watchfile-ref-unref.js b/test/parallel/test-fs-watchfile-ref-unref.js
new file mode 100644
index 00000000000000..4ac2691ea9d702
--- /dev/null
+++ b/test/parallel/test-fs-watchfile-ref-unref.js
@@ -0,0 +1,35 @@
+'use strict';
+
+const common = require('../common');
+
+const fs = require('fs');
+const assert = require('assert');
+
+const uncalledListener = common.mustNotCall();
+const uncalledListener2 = common.mustNotCall();
+const watcher = fs.watchFile(__filename, uncalledListener);
+
+watcher.unref();
+watcher.unref();
+watcher.ref();
+watcher.unref();
+watcher.ref();
+watcher.ref();
+watcher.unref();
+
+fs.unwatchFile(__filename, uncalledListener);
+
+// Watch the file with two different listeners.
+fs.watchFile(__filename, uncalledListener);
+const watcher2 = fs.watchFile(__filename, uncalledListener2);
+
+setTimeout(
+  common.mustCall(() => {
+    fs.unwatchFile(__filename, common.mustNotCall());
+    assert.strictEqual(watcher2.listenerCount('change'), 2);
+    fs.unwatchFile(__filename, uncalledListener);
+    assert.strictEqual(watcher2.listenerCount('change'), 1);
+    watcher2.unref();
+  }),
+  common.platformTimeout(100)
+);
diff --git a/tools/doc/type-parser.js b/tools/doc/type-parser.js
index 67b3a85c2c54c2..c2ad56655f8602 100644
--- a/tools/doc/type-parser.js
+++ b/tools/doc/type-parser.js
@@ -82,6 +82,7 @@ const customTypesMap = {
   'fs.FSWatcher': 'fs.html#fs_class_fs_fswatcher',
   'fs.ReadStream': 'fs.html#fs_class_fs_readstream',
   'fs.Stats': 'fs.html#fs_class_fs_stats',
+  'fs.StatWatcher': 'fs.html#fs_class_fs_statwatcher',
   'fs.WriteStream': 'fs.html#fs_class_fs_writestream',
 
   'http.Agent': 'http.html#http_class_http_agent',