fix(napi): support for env cleanup hooks (#17324)

This commit adds support for "napi_add_env_cleanup_hook" and
"napi_remove_env_cleanup_hook" function for Node-API.

7 files changed, 143 insertions, 18 deletions

diff --git a/cli/napi/env.rs b/cli/napi/env.rs
index 1d59b2f95..6568c20c9 100644
--- a/cli/napi/env.rs
+++ b/cli/napi/env.rs
@@ -52,24 +52,49 @@ fn napi_fatal_exception(env: *mut Env, value: napi_value) -> Result {
   );
 }
 
-// TODO: properly implement
 #[napi_sym::napi_sym]
 fn napi_add_env_cleanup_hook(
-  _env: *mut Env,
-  _hook: extern "C" fn(*const c_void),
-  _data: *const c_void,
+  env: *mut Env,
+  hook: extern "C" fn(*const c_void),
+  data: *const c_void,
 ) -> Result {
-  log::info!("napi_add_env_cleanup_hook is currently not supported");
+  let env: &mut Env = env.as_mut().ok_or(Error::InvalidArg)?;
+
+  {
+    let mut env_cleanup_hooks = env.cleanup_hooks.borrow_mut();
+    if env_cleanup_hooks
+      .iter()
+      .any(|pair| pair.0 == hook && pair.1 == data)
+    {
+      panic!("Cleanup hook with this data already registered");
+    }
+    env_cleanup_hooks.push((hook, data));
+  }
   Ok(())
 }
 
 #[napi_sym::napi_sym]
 fn napi_remove_env_cleanup_hook(
-  _env: *mut Env,
-  _hook: extern "C" fn(*const c_void),
-  _data: *const c_void,
+  env: *mut Env,
+  hook: extern "C" fn(*const c_void),
+  data: *const c_void,
 ) -> Result {
-  log::info!("napi_remove_env_cleanup_hook is currently not supported");
+  let env: &mut Env = env.as_mut().ok_or(Error::InvalidArg)?;
+
+  {
+    let mut env_cleanup_hooks = env.cleanup_hooks.borrow_mut();
+    // Hooks are supposed to be removed in LIFO order
+    let maybe_index = env_cleanup_hooks
+      .iter()
+      .rposition(|&pair| pair.0 == hook && pair.1 == data);
+
+    if let Some(index) = maybe_index {
+      env_cleanup_hooks.remove(index);
+    } else {
+      panic!("Cleanup hook with this data not found");
+    }
+  }
+
   Ok(())
 }
 
diff --git a/ext/napi/lib.rs b/ext/napi/lib.rs
index 577545774..882f7c19d 100644
--- a/ext/napi/lib.rs
+++ b/ext/napi/lib.rs
@@ -18,6 +18,7 @@ use std::cell::RefCell;
 use std::ffi::CString;
 use std::path::Path;
 use std::path::PathBuf;
+use std::rc::Rc;
 use std::task::Poll;
 use std::thread_local;
 
@@ -333,8 +334,20 @@ pub struct NapiState {
     mpsc::UnboundedReceiver<ThreadSafeFunctionStatus>,
   pub threadsafe_function_sender:
     mpsc::UnboundedSender<ThreadSafeFunctionStatus>,
+  pub env_cleanup_hooks:
+    Rc<RefCell<Vec<(extern "C" fn(*const c_void), *const c_void)>>>,
 }
 
+impl Drop for NapiState {
+  fn drop(&mut self) {
+    let mut hooks = self.env_cleanup_hooks.borrow_mut();
+    // Hooks are supposed to be run in LIFO order
+    let hooks = hooks.drain(..).rev();
+    for (fn_ptr, data) in hooks {
+      (fn_ptr)(data);
+    }
+  }
+}
 #[repr(C)]
 #[derive(Debug)]
 /// Env that is shared between all contexts in same native module.
@@ -376,6 +389,8 @@ pub struct Env {
   pub async_work_sender: mpsc::UnboundedSender<PendingNapiAsyncWork>,
   pub threadsafe_function_sender:
     mpsc::UnboundedSender<ThreadSafeFunctionStatus>,
+  pub cleanup_hooks:
+    Rc<RefCell<Vec<(extern "C" fn(*const c_void), *const c_void)>>>,
 }
 
 unsafe impl Send for Env {}
@@ -387,6 +402,9 @@ impl Env {
     context: v8::Global<v8::Context>,
     sender: mpsc::UnboundedSender<PendingNapiAsyncWork>,
     threadsafe_function_sender: mpsc::UnboundedSender<ThreadSafeFunctionStatus>,
+    cleanup_hooks: Rc<
+      RefCell<Vec<(extern "C" fn(*const c_void), *const c_void)>>,
+    >,
   ) -> Self {
     let sc = sender.clone();
     ASYNC_WORK_SENDER.with(|s| {
@@ -404,6 +422,7 @@ impl Env {
       open_handle_scopes: 0,
       async_work_sender: sender,
       threadsafe_function_sender,
+      cleanup_hooks,
     }
   }
 
@@ -507,6 +526,7 @@ pub fn init<P: NapiPermissions + 'static>(unstable: bool) -> Extension {
         threadsafe_function_sender,
         threadsafe_function_receiver,
         active_threadsafe_functions: 0,
+        env_cleanup_hooks: Rc::new(RefCell::new(vec![])),
       });
       state.put(Unstable(unstable));
       Ok(())
@@ -543,13 +563,14 @@ where
   let permissions = op_state.borrow_mut::<NP>();
   permissions.check(Some(&PathBuf::from(&path)))?;
 
-  let (async_work_sender, tsfn_sender, isolate_ptr) = {
+  let (async_work_sender, tsfn_sender, isolate_ptr, cleanup_hooks) = {
     let napi_state = op_state.borrow::<NapiState>();
     let isolate_ptr = op_state.borrow::<*mut v8::OwnedIsolate>();
     (
       napi_state.async_work_sender.clone(),
       napi_state.threadsafe_function_sender.clone(),
       *isolate_ptr,
+      napi_state.env_cleanup_hooks.clone(),
     )
   };
 
@@ -571,6 +592,7 @@ where
     v8::Global::new(scope, ctx),
     async_work_sender,
     tsfn_sender,
+    cleanup_hooks,
   );
   env.shared = Box::into_raw(Box::new(env_shared));
   let env_ptr = Box::into_raw(Box::new(env)) as _;
diff --git a/test_napi/cleanup_hook_test.js b/test_napi/cleanup_hook_test.js
new file mode 100644
index 000000000..30ceae470
--- /dev/null
+++ b/test_napi/cleanup_hook_test.js
@@ -0,0 +1,33 @@
+// Copyright 2018-2023 the Deno authors. All rights reserved. MIT license.
+
+import { assertEquals, loadTestLibrary } from "./common.js";
+
+const properties = loadTestLibrary();
+
+if (import.meta.main) {
+  properties.installCleanupHook();
+  console.log("installed cleanup hook");
+} else {
+  Deno.test("napi cleanup hook", async () => {
+    const { stdout, stderr, code } = await new Deno.Command(Deno.execPath(), {
+      args: [
+        "run",
+        "--allow-read",
+        "--allow-run",
+        "--allow-ffi",
+        "--unstable",
+        import.meta.url,
+      ],
+    }).output();
+
+    assertEquals(code, 0);
+    assertEquals(new TextDecoder().decode(stderr), "");
+
+    const stdoutText = new TextDecoder().decode(stdout);
+    const stdoutLines = stdoutText.split("\n");
+    assertEquals(stdoutLines.length, 4);
+    assertEquals(stdoutLines[0], "installed cleanup hook");
+    assertEquals(stdoutLines[1], "cleanup(18)");
+    assertEquals(stdoutLines[2], "cleanup(42)");
+  });
+}
diff --git a/test_napi/src/arraybuffer.rs b/test_napi/src/arraybuffer.rs
index ce50f914f..6765df481 100644
--- a/test_napi/src/arraybuffer.rs
+++ b/test_napi/src/arraybuffer.rs
@@ -2,7 +2,6 @@
 
 use napi_sys::Status::napi_ok;
 use napi_sys::*;
-use std::ptr;
 
 extern "C" fn test_detached(
   env: napi_env,
diff --git a/test_napi/src/lib.rs b/test_napi/src/lib.rs
index 0bdea0280..3ab1f3c84 100644
--- a/test_napi/src/lib.rs
+++ b/test_napi/src/lib.rs
@@ -2,6 +2,9 @@
 #![allow(clippy::all)]
 #![allow(clippy::undocumented_unsafe_blocks)]
 
+use napi_sys::Status::napi_ok;
+use std::ffi::c_void;
+
 use napi_sys::*;
 
 pub mod array;
@@ -20,9 +23,9 @@ pub mod typedarray;
 #[macro_export]
 macro_rules! get_callback_info {
   ($env: expr, $callback_info: expr, $size: literal) => {{
-    let mut args = [ptr::null_mut(); $size];
+    let mut args = [std::ptr::null_mut(); $size];
     let mut argc = $size;
-    let mut this = ptr::null_mut();
+    let mut this = std::ptr::null_mut();
     unsafe {
       assert!(
         napi_get_cb_info(
@@ -31,7 +34,7 @@ macro_rules! get_callback_info {
           &mut argc,
           args.as_mut_ptr(),
           &mut this,
-          ptr::null_mut(),
+          std::ptr::null_mut(),
         ) == napi_ok,
       )
     };
@@ -44,17 +47,58 @@ macro_rules! new_property {
   ($env: expr, $name: expr, $value: expr) => {
     napi_property_descriptor {
       utf8name: $name.as_ptr() as *const std::os::raw::c_char,
-      name: ptr::null_mut(),
+      name: std::ptr::null_mut(),
       method: Some($value),
       getter: None,
       setter: None,
-      data: ptr::null_mut(),
+      data: std::ptr::null_mut(),
       attributes: 0,
-      value: ptr::null_mut(),
+      value: std::ptr::null_mut(),
     }
   };
 }
 
+extern "C" fn cleanup(arg: *mut c_void) {
+  println!("cleanup({})", arg as i64);
+}
+
+static SECRET: i64 = 42;
+static WRONG_SECRET: i64 = 17;
+static THIRD_SECRET: i64 = 18;
+
+extern "C" fn install_cleanup_hook(
+  env: napi_env,
+  info: napi_callback_info,
+) -> napi_value {
+  let (_args, argc, _) = get_callback_info!(env, info, 1);
+  assert_eq!(argc, 0);
+
+  unsafe {
+    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_remove_env_cleanup_hook(
+      env,
+      Some(cleanup),
+      WRONG_SECRET as *mut c_void,
+    );
+  }
+
+  std::ptr::null_mut()
+}
+
+pub fn init_cleanup_hook(env: napi_env, exports: napi_value) {
+  let properties = &[new_property!(
+    env,
+    "installCleanupHook\0",
+    install_cleanup_hook
+  )];
+
+  unsafe {
+    napi_define_properties(env, exports, properties.len(), properties.as_ptr())
+  };
+}
+
 #[no_mangle]
 unsafe extern "C" fn napi_register_module_v1(
   env: napi_env,
@@ -77,6 +121,7 @@ unsafe extern "C" fn napi_register_module_v1(
   object_wrap::init(env, exports);
   callback::init(env, exports);
   r#async::init(env, exports);
+  init_cleanup_hook(env, exports);
 
   exports
 }
diff --git a/test_napi/src/strings.rs b/test_napi/src/strings.rs
index 5029944da..af6f84189 100644
--- a/test_napi/src/strings.rs
+++ b/test_napi/src/strings.rs
@@ -3,7 +3,6 @@
 use napi_sys::Status::napi_ok;
 use napi_sys::ValueType::napi_string;
 use napi_sys::*;
-use std::ptr;
 
 extern "C" fn test_utf8(env: napi_env, info: napi_callback_info) -> napi_value {
   let (args, argc, _) = crate::get_callback_info!(env, info, 1);
diff --git a/test_napi/tests/napi_tests.rs b/test_napi/tests/napi_tests.rs
index 1c0f59fdf..3351d74c1 100644
--- a/test_napi/tests/napi_tests.rs
+++ b/test_napi/tests/napi_tests.rs
@@ -30,6 +30,7 @@ fn napi_tests() {
     .arg("--allow-read")
     .arg("--allow-env")
     .arg("--allow-ffi")
+    .arg("--allow-run")
     .arg("--unstable")
     .spawn()
     .unwrap()
@@ -37,6 +38,7 @@ fn napi_tests() {
     .unwrap();
   let stdout = std::str::from_utf8(&output.stdout).unwrap();
   let stderr = std::str::from_utf8(&output.stderr).unwrap();
+
   if !output.status.success() {
     println!("stdout {}", stdout);
     println!("stderr {}", stderr);