From b5a672eabb3b1681b88ac313f01011401ee8a07b Mon Sep 17 00:00:00 2001
From: Fedor Indutny <fedor@indutny.com>
Date: Wed, 19 May 2021 17:25:33 -0700
Subject: [PATCH] Handle early finalization of TSFN by Node.js

Thread-safe functions can get released early by Node.js itself during
the teardown of the Environment that holds them. When this happens -
`finalize_cb` of the thread-safe function is invoked to notify us. This
change makes us remember if the thread-safe function was finalized
before we got into the `Drop` implementation for it so that we don't
deallocate it twice.

Additionally, we should not invoke `call_threadsafe_function` after
either `finalize_cb` call or `napi_closing` return value from previous
`call_threadsafe_function`. Handle this by holding a mutex around the
call, which doesn't introduce any additional locking because
`call_threadsafe_function` itself holds another mutex while active.
---
 crates/neon-runtime/src/napi/tsfn.rs | 44 +++++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/crates/neon-runtime/src/napi/tsfn.rs b/crates/neon-runtime/src/napi/tsfn.rs
index 271a46511..313a7d032 100644
--- a/crates/neon-runtime/src/napi/tsfn.rs
+++ b/crates/neon-runtime/src/napi/tsfn.rs
@@ -2,6 +2,7 @@
 
 use std::ffi::c_void;
 use std::mem::MaybeUninit;
+use std::sync::{Arc, Mutex};
 
 use crate::napi::bindings as napi;
 use crate::raw::{Env, Local};
@@ -34,6 +35,7 @@ unsafe impl Sync for Tsfn {}
 /// function for scheduling tasks to execute on a JavaScript thread.
 pub struct ThreadsafeFunction<T> {
     tsfn: Tsfn,
+    is_finalized: Arc<Mutex<bool>>,
     callback: fn(Option<Env>, T),
 }
 
@@ -76,6 +78,7 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
         callback: fn(Option<Env>, T),
     ) -> Self {
         let mut result = MaybeUninit::uninit();
+        let is_finalized = Arc::new(Mutex::new(false));
 
         assert_eq!(
             napi::create_threadsafe_function(
@@ -87,8 +90,8 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
                 // Always set the reference count to 1. Prefer using
                 // Rust `Arc` to maintain the struct.
                 1,
-                std::ptr::null_mut(),
-                None,
+                Arc::into_raw(is_finalized.clone()) as *mut _,
+                Some(Self::finalize),
                 std::ptr::null_mut(),
                 Some(Self::callback),
                 result.as_mut_ptr(),
@@ -98,6 +101,7 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
 
         Self {
             tsfn: Tsfn(result.assume_init()),
+            is_finalized: is_finalized,
             callback,
         }
     }
@@ -115,12 +119,28 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
             data,
         }));
 
-        let status =
-            unsafe { napi::call_threadsafe_function(self.tsfn.0, callback as *mut _, is_blocking) };
+        // Hold the lock before entering `call_threadsafe_function` so that
+        // `finalize_cb` would never complete.
+        let mut is_finalized = self.is_finalized.lock().unwrap();
+
+        let status = {
+            if *is_finalized {
+                napi::Status::Closing
+            } else {
+                unsafe {
+                    napi::call_threadsafe_function(self.tsfn.0, callback as *mut _, is_blocking)
+                }
+            }
+        };
 
         if status == napi::Status::Ok {
             Ok(())
         } else {
+            // Prevent further calls to `call_threadsafe_function`
+            if status == napi::Status::Closing {
+                *is_finalized = true;
+            }
+
             // If the call failed, the callback won't execute
             let callback = unsafe { Box::from_raw(callback) };
 
@@ -149,6 +169,14 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
         );
     }
 
+    // Provides a C ABI wrapper for a napi callback notifying us about tsfn
+    // being finalized.
+    unsafe extern "C" fn finalize(_env: Env, data: *mut c_void, _hint: *mut c_void) {
+        let is_finalized = Arc::from_raw(data as *mut Mutex<bool>);
+
+        *is_finalized.lock().unwrap() = true;
+    }
+
     // Provides a C ABI wrapper for invoking the user supplied function pointer
     unsafe extern "C" fn callback(
         env: Env,
@@ -167,6 +195,14 @@ impl<T: Send + 'static> ThreadsafeFunction<T> {
 
 impl<T> Drop for ThreadsafeFunction<T> {
     fn drop(&mut self) {
+        let is_finalized = self.is_finalized.lock().unwrap();
+
+        // tsfn was already finalized by `Environment::CleanupHandles()` in
+        // Node.js
+        if *is_finalized {
+            return;
+        }
+
         unsafe {
             napi::release_threadsafe_function(
                 self.tsfn.0,