From 225114166aa7426d4b93fa13635559029c5ba65d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bartek=20Iwa=C5=84czuk?= <biwanczuk@gmail.com>
Date: Fri, 13 Jan 2023 22:17:25 +0100
Subject: [PATCH] fix(napi): allow cleanup hook to remove itself (#17402)

This commit fixes "cleanup hooks" in NAPI integration in two ways:
- don't hold to RefCell's borrow while iterating over hooks
- allow a hook to remove itself when being called
---
 ext/napi/lib.rs      | 29 +++++++++++++++++++++++++----
 test_napi/src/lib.rs |  6 ++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/ext/napi/lib.rs b/ext/napi/lib.rs
index 57f73a0ca2..22b0e3a088 100644
--- a/ext/napi/lib.rs
+++ b/ext/napi/lib.rs
@@ -343,11 +343,32 @@ pub struct NapiState {
 
 impl Drop for NapiState {
   fn drop(&mut self) {
-    let mut hooks = self.env_cleanup_hooks.borrow_mut();
+    let hooks = {
+      let h = self.env_cleanup_hooks.borrow_mut();
+      h.clone()
+    };
+
     // Hooks are supposed to be run in LIFO order
-    let hooks = hooks.drain(..).rev();
-    for (fn_ptr, data) in hooks {
-      (fn_ptr)(data);
+    let hooks_to_run = hooks.into_iter().rev();
+
+    for hook in hooks_to_run {
+      // This hook might have been removed by a previous hook, in such case skip it here.
+      if !self
+        .env_cleanup_hooks
+        .borrow()
+        .iter()
+        .any(|pair| pair.0 == hook.0 && pair.1 == hook.1)
+      {
+        continue;
+      }
+
+      (hook.0)(hook.1);
+      {
+        self
+          .env_cleanup_hooks
+          .borrow_mut()
+          .retain(|pair| !(pair.0 == hook.0 && pair.1 == hook.1));
+      }
     }
   }
 }
diff --git a/test_napi/src/lib.rs b/test_napi/src/lib.rs
index 2d43b0aea3..b54b3886b8 100644
--- a/test_napi/src/lib.rs
+++ b/test_napi/src/lib.rs
@@ -66,6 +66,11 @@ extern "C" fn cleanup(arg: *mut c_void) {
   println!("cleanup({})", arg as i64);
 }
 
+extern "C" fn remove_this_hook(arg: *mut c_void) {
+  let env = arg as napi_env;
+  unsafe { napi_remove_env_cleanup_hook(env, Some(remove_this_hook), arg) };
+}
+
 static SECRET: i64 = 42;
 static WRONG_SECRET: i64 = 17;
 static THIRD_SECRET: i64 = 18;
@@ -81,6 +86,7 @@ extern "C" fn install_cleanup_hook(
     napi_add_env_cleanup_hook(env, Some(cleanup), WRONG_SECRET as *mut c_void);
     napi_add_env_cleanup_hook(env, Some(cleanup), SECRET as *mut c_void);
     napi_add_env_cleanup_hook(env, Some(cleanup), THIRD_SECRET as *mut c_void);
+    napi_add_env_cleanup_hook(env, Some(remove_this_hook), env as *mut c_void);
     napi_remove_env_cleanup_hook(
       env,
       Some(cleanup),