napi: keep threadsafe functions alive after their env is torn down

src/codegen/generate-js2native.ts

@@ -67,6 +67,7 @@ const rustIdentifierPaths: Record<string, string> = {
   "install_binding.rs": "install_jsc/install_binding.rs",
   "ipc.rs": "jsc/ipc.rs",
   "mysql.rs": "sql_jsc/mysql.rs",
+  "napi_body.rs": "runtime/napi/napi_body.rs",
   "node_assert_binding.rs": "runtime/node/node_assert_binding.rs",
   "node_cluster_binding.rs": "runtime/node/node_cluster_binding.rs",
   "node_crypto_binding.rs": "runtime/node/node_crypto_binding.rs",

src/js/internal-for-testing.ts

@@ -52,6 +52,8 @@ const shellParse = $newRustFunction("shell.rs", "TestingAPIs.shellParse", 2);
 
 export const sslCtxLiveCount = $newRustFunction("SecureContext.rs", "jsLiveCount", 0);
 
+export const napiThreadsafeFunctionLiveCount = $newRustFunction("napi_body.rs", "jsThreadsafeFunctionLiveCount", 0);
+
 export const escapeRegExp = $newRustFunction("escapeRegExp.rs", "jsEscapeRegExp", 1);
 export const escapeRegExpForPackageNameMatching = $newRustFunction(
   "escapeRegExp.rs",

src/jsc/bindings/napi.cpp

@@ -3039,6 +3039,16 @@ extern "C" void napi_internal_cleanup_env_cpp(napi_env env)
     env->cleanup();
 }
 
+extern "C" bool NapiEnv__registerThreadSafeFunction(napi_env env, void* tsfn)
+{
+    return env->registerThreadSafeFunction(tsfn);
+}
+
+extern "C" void NapiEnv__unregisterThreadSafeFunction(napi_env env, void* tsfn)
+{
+    env->unregisterThreadSafeFunction(tsfn);
+}
+
 extern "C" void napi_internal_remove_finalizer(napi_env env, napi_finalize callback, void* hint, void* data)
 {
     env->removeFinalizer(callback, hint, data);

src/jsc/bindings/napi.h

@@ -17,13 +17,16 @@
 #include "wtf/Assertions.h"
 #include "napi_macros.h"
 
+#include <wtf/HashSet.h>
 #include <wtf/ListHashSet.h>
+#include <wtf/Lock.h>
 
 #include <optional>
 #include <unordered_set>
 #include <variant>
 
 extern "C" void napi_internal_register_cleanup_zig(napi_env env);
+extern "C" void napi_internal_threadsafe_function_env_teardown(void* tsfn);
 extern "C" void napi_internal_suppress_crash_on_abort_if_desired();
 extern "C" void Bun__crashHandler(const char* message, size_t message_len);
 
@@ -211,6 +214,22 @@ struct NapiEnv : public WTF::RefCounted<NapiEnv> {
             drain();
         }
 
+        // Threadsafe functions hold a raw pointer to this env's event loop,
+        // which a worker's shutdown frees while addon threads keep running.
+        // Neutralize them all before that happens (Node: ThreadSafeFunction::Cleanup).
+        // Their finalizers are native callbacks like the ones below, so they
+        // start from a clean exception state too: a cleanup hook may have left one.
+        clearExceptionsBetweenFinalizers();
+        abortThreadSafeFunctions();
+
+        // A threadsafe function's finalizer can register a cleanup hook of its
+        // own, and the loop above has already run: drain again so it is not
+        // dropped on the floor.
+        while (!m_cleanupHooks.empty()) {
+            drain();
+        }
+        clearExceptionsBetweenFinalizers();
+
         // Defer GC during entire finalizer cleanup to prevent iterator invalidation.
         // This prevents any GC-triggered finalizer execution while m_finalizers is being iterated.
         JSC::DeferGCForAWhile deferGC(m_vm);
@@ -242,6 +261,40 @@ struct NapiEnv : public WTF::RefCounted<NapiEnv> {
         clearExceptionsBetweenFinalizers();
     }
 
+    // Threadsafe-function registry. Entries are raw ThreadSafeFunction* owned
+    // by the Rust side, added and removed on the JS thread only: at creation,
+    // by the destroy path (an event-loop task), and by teardown below. The
+    // lock guards the torn-down flag so the "created after teardown" case is
+    // decided in one critical section.
+    bool registerThreadSafeFunction(void* tsfn)
+    {
+        WTF::Locker locker { m_threadSafeFunctionsLock };
+        if (m_threadSafeFunctionsTornDown) {
+            return false;
+        }
+        m_threadSafeFunctions.add(tsfn);
+        return true;
+    }
+
+    void unregisterThreadSafeFunction(void* tsfn)
+    {
+        WTF::Locker locker { m_threadSafeFunctionsLock };
+        m_threadSafeFunctions.remove(tsfn);
+    }
+
+    void abortThreadSafeFunctions()
+    {
+        WTF::HashSet<void*> tsfns;
+        {
+            WTF::Locker locker { m_threadSafeFunctionsLock };
+            m_threadSafeFunctionsTornDown = true;
+            tsfns = std::exchange(m_threadSafeFunctions, WTF::HashSet<void*> {});
+        }
+        for (void* tsfn : tsfns) {
+            napi_internal_threadsafe_function_env_teardown(tsfn);
+        }
+    }
+
     void removeFinalizer(napi_finalize callback, void* hint, void* data)
     {
         m_finalizers.remove({ callback, hint, data });
@@ -533,6 +586,10 @@ struct NapiEnv : public WTF::RefCounted<NapiEnv> {
     JSC::Strong<JSC::Unknown> m_pendingException;
     size_t m_cleanupHookCounter = 0;
 
+    WTF::Lock m_threadSafeFunctionsLock;
+    WTF::HashSet<void*> m_threadSafeFunctions WTF_GUARDED_BY_LOCK(m_threadSafeFunctionsLock);
+    bool m_threadSafeFunctionsTornDown WTF_GUARDED_BY_LOCK(m_threadSafeFunctionsLock) = false;
+
     // Drop any pending exception -- VM-scope or env-scope -- between
     // finalizers run from cleanup(). Used by cleanup() only. Defined
     // out-of-line in napi.cpp so its uses of JSC::TopExceptionScope

src/runtime/napi/mod.rs

@@ -6,7 +6,7 @@
 //! `bun_threading::{Condvar, Mutex, WorkPool}`, `bun_output` macros.
 
 #[path = "napi_body.rs"]
-mod napi_body;
+pub(crate) mod napi_body;
 pub use napi_body::{
     NapiFinalizerTask, ThreadSafeFunction, fix_dead_code_elimination, napi_async_work,
 };