From ea9ce1c801dc28da409a98e915fea79af23d2404 Mon Sep 17 00:00:00 2001
From: Kevin Eady <8634912+KevinEady@users.noreply.github.com>
Date: Wed, 9 Oct 2019 20:18:56 +0200
Subject: [PATCH] tsfn: add wrappers for Ref and Unref

Ref: https://github.com/nodejs/node-addon-api/issues/556#issuecomment-539157399
PR-URL: https://github.com/nodejs/node-addon-api/pull/561
Reviewed-By: Michael Dawson <michael_dawson@ca.ibm.com>
---
 napi-inl.h                                    | 14 +++++
 napi.h                                        |  6 +++
 test/binding.cc                               |  2 +
 test/binding.gyp                              |  1 +
 test/index.js                                 |  1 +
 test/napi_child.js                            |  7 +++
 .../threadsafe_function_unref.cc              | 41 ++++++++++++++
 .../threadsafe_function_unref.js              | 53 +++++++++++++++++++
 8 files changed, 125 insertions(+)
 create mode 100644 test/threadsafe_function/threadsafe_function_unref.cc
 create mode 100644 test/threadsafe_function/threadsafe_function_unref.js

diff --git a/napi-inl.h b/napi-inl.h
index 366cb6cc3..271819944 100644
--- a/napi-inl.h
+++ b/napi-inl.h
@@ -4029,6 +4029,20 @@ inline napi_status ThreadSafeFunction::NonBlockingCall(
   return CallInternal(new CallbackWrapper(wrapper), napi_tsfn_nonblocking);
 }
 
+inline void ThreadSafeFunction::Ref(napi_env env) const {
+  if (_tsfn != nullptr) {
+    napi_status status = napi_ref_threadsafe_function(env, *_tsfn);
+    NAPI_THROW_IF_FAILED_VOID(env, status);
+  }
+}
+
+inline void ThreadSafeFunction::Unref(napi_env env) const {
+  if (_tsfn != nullptr) {
+    napi_status status = napi_unref_threadsafe_function(env, *_tsfn);
+    NAPI_THROW_IF_FAILED_VOID(env, status);
+  }
+}
+
 inline napi_status ThreadSafeFunction::Acquire() const {
   return napi_acquire_threadsafe_function(*_tsfn);
 }
diff --git a/napi.h b/napi.h
index 921aaa8bc..05b2914f4 100644
--- a/napi.h
+++ b/napi.h
@@ -2011,6 +2011,12 @@ namespace Napi {
     template <typename DataType, typename Callback>
     napi_status NonBlockingCall(DataType* data, Callback callback) const;
 
+    // This API may only be called from the main thread.
+    void Ref(napi_env env) const;
+
+    // This API may only be called from the main thread.
+    void Unref(napi_env env) const;
+
     // This API may be called from any thread.
     napi_status Acquire() const;
 
diff --git a/test/binding.cc b/test/binding.cc
index dca00771a..9286a87af 100644
--- a/test/binding.cc
+++ b/test/binding.cc
@@ -38,6 +38,7 @@ Object InitObjectDeprecated(Env env);
 Object InitPromise(Env env);
 #if (NAPI_VERSION > 3)
 Object InitThreadSafeFunctionPtr(Env env);
+Object InitThreadSafeFunctionUnref(Env env);
 Object InitThreadSafeFunction(Env env);
 #endif
 Object InitTypedArray(Env env);
@@ -83,6 +84,7 @@ Object Init(Env env, Object exports) {
   exports.Set("promise", InitPromise(env));
 #if (NAPI_VERSION > 3)
   exports.Set("threadsafe_function_ptr", InitThreadSafeFunctionPtr(env));
+  exports.Set("threadsafe_function_unref", InitThreadSafeFunctionUnref(env));
   exports.Set("threadsafe_function", InitThreadSafeFunction(env));
 #endif
   exports.Set("typedarray", InitTypedArray(env));
diff --git a/test/binding.gyp b/test/binding.gyp
index f8c1cb0a8..710beff52 100644
--- a/test/binding.gyp
+++ b/test/binding.gyp
@@ -34,6 +34,7 @@
         'object/set_property.cc',
         'promise.cc',
         'threadsafe_function/threadsafe_function_ptr.cc',
+        'threadsafe_function/threadsafe_function_unref.cc',
         'threadsafe_function/threadsafe_function.cc',
         'typedarray.cc',
         'objectwrap.cc',
diff --git a/test/index.js b/test/index.js
index d03c2e5b7..2c3092e93 100644
--- a/test/index.js
+++ b/test/index.js
@@ -38,6 +38,7 @@ let testModules = [
   'object/set_property',
   'promise',
   'threadsafe_function/threadsafe_function_ptr',
+  'threadsafe_function/threadsafe_function_unref',
   'threadsafe_function/threadsafe_function',
   'typedarray',
   'typedarray-bigint',
diff --git a/test/napi_child.js b/test/napi_child.js
index 29a80a115..76f0bc56a 100644
--- a/test/napi_child.js
+++ b/test/napi_child.js
@@ -5,3 +5,10 @@ exports.spawnSync = function(command, args, options) {
   }
   return require('child_process').spawnSync(command, args, options);
 };
+
+exports.spawn = function(command, args, options) {
+  if (require('../index').needsFlag) {
+    args.splice(0, 0, '--napi-modules');
+  }
+  return require('child_process').spawn(command, args, options);
+};
diff --git a/test/threadsafe_function/threadsafe_function_unref.cc b/test/threadsafe_function/threadsafe_function_unref.cc
new file mode 100644
index 000000000..6877e50f6
--- /dev/null
+++ b/test/threadsafe_function/threadsafe_function_unref.cc
@@ -0,0 +1,41 @@
+#include "napi.h"
+
+#if (NAPI_VERSION > 3)
+
+using namespace Napi;
+
+namespace {
+
+static Value TestUnref(const CallbackInfo& info) {
+  Napi::Env env = info.Env();
+  Object global = env.Global();
+  Object resource = info[0].As<Object>();
+  Function cb = info[1].As<Function>();
+  Function setTimeout = global.Get("setTimeout").As<Function>();
+  ThreadSafeFunction* tsfn = new ThreadSafeFunction;
+
+  *tsfn = ThreadSafeFunction::New(info.Env(), cb, resource, "Test", 1, 1, [tsfn](Napi::Env /* env */) {
+    delete tsfn;
+  });
+
+  tsfn->BlockingCall();
+
+  setTimeout.Call( global, {
+    Function::New(env, [tsfn](const CallbackInfo& info) {
+      tsfn->Unref(info.Env());
+    }),
+    Number::New(env, 100)
+  });
+
+  return info.Env().Undefined();
+}
+
+}
+
+Object InitThreadSafeFunctionUnref(Env env) {
+  Object exports = Object::New(env);
+  exports["testUnref"] = Function::New(env, TestUnref);
+  return exports;
+}
+
+#endif
diff --git a/test/threadsafe_function/threadsafe_function_unref.js b/test/threadsafe_function/threadsafe_function_unref.js
new file mode 100644
index 000000000..37ce56b4d
--- /dev/null
+++ b/test/threadsafe_function/threadsafe_function_unref.js
@@ -0,0 +1,53 @@
+'use strict';
+
+const assert = require('assert');
+const buildType = process.config.target_defaults.default_configuration;
+
+const isMainProcess = process.argv[1] != __filename;
+
+/**
+ * In order to test that the event loop exits even with an active TSFN, we need
+ * to spawn a new process for the test.
+ *   - Main process: spawns new node instance, executing this script
+ *   - Child process: creates TSFN. Native module Unref's via setTimeout after some time but does NOT call Release.
+ *
+ * Main process should expect child process to exit.
+ */
+
+if (isMainProcess) {
+  test(`../build/${buildType}/binding.node`);
+  test(`../build/${buildType}/binding_noexcept.node`);
+} else {
+  test(process.argv[2]);
+}
+
+function test(bindingFile) {
+  if (isMainProcess) {
+    // Main process
+    const child = require('../napi_child').spawn(process.argv[0], [ '--expose-gc', __filename, bindingFile ], {
+      stdio: 'inherit',
+    });
+
+    let timeout = setTimeout( function() {
+      child.kill();
+      timeout = 0;
+      throw new Error("Expected child to die");
+    }, 5000);
+
+    child.on("error", (err) => {
+      clearTimeout(timeout);
+      timeout = 0;
+      throw new Error(err);
+    })
+
+    child.on("close", (code) => {
+      if (timeout) clearTimeout(timeout);
+      assert(!code, "Expected return value 0");
+    });
+
+  } else {
+    // Child process
+    const binding = require(bindingFile);
+    binding.threadsafe_function_unref.testUnref({}, () => { });
+  }
+}