fix(napi): memory leak in Buffer/ArrayBuffer

crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs

@@ -2,7 +2,10 @@ use std::ffi::c_void;
 use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::ptr;
-use std::sync::Arc;
+use std::sync::{
+  atomic::{AtomicBool, Ordering},
+  Arc,
+};
 
 pub use crate::js_values::TypedArrayType;
 use crate::{check_status, sys, Error, Result, Status};
@@ -24,13 +27,14 @@ trait Finalizer {
 macro_rules! impl_typed_array {
   ($name:ident, $rust_type:ident, $typed_array_type:expr) => {
     pub struct $name {
-      /// Used for `clone_reference` Owned | External TypedArray
-      data_reference: Option<Arc<()>>,
       data: *mut $rust_type,
       length: usize,
       data_managed_type: DataManagedType,
       byte_offset: usize,
       raw: Option<(crate::sys::napi_ref, crate::sys::napi_env)>,
+      // Use `Arc` for ref count
+      // Use `AtomicBool` for flag to indicate whether the value is dropped in VM
+      drop_in_vm: Arc<AtomicBool>,
       finalizer_notify: Box<dyn FnOnce(*mut $rust_type, usize)>,
     }
 
@@ -52,22 +56,35 @@ macro_rules! impl_typed_array {
       }
 
       fn ref_count(&self) -> usize {
-        if let Some(inner) = &self.data_reference {
-          Arc::strong_count(inner)
-        } else {
-          0
-        }
+        Arc::strong_count(&self.drop_in_vm)
       }
     }
 
     impl Drop for $name {
       fn drop(&mut self) {
-        if let Some((ref_, env)) = self.raw {
-          crate::check_status_or_throw!(
-            env,
-            unsafe { sys::napi_reference_unref(env, ref_, &mut 0) },
-            "Failed to delete Buffer reference in drop"
-          );
+        if Arc::strong_count(&self.drop_in_vm) == 1 {
+          if let Some((ref_, env)) = self.raw {
+            crate::check_status_or_throw!(
+              env,
+              unsafe { sys::napi_reference_unref(env, ref_, &mut 0) },
+              "Failed to delete Buffer reference in drop"
+            );
+            return;
+          }
+          if !self.drop_in_vm.load(Ordering::Acquire) {
+            match &self.data_managed_type {
+              DataManagedType::Owned => {
+                let length = self.length;
+                unsafe { Vec::from_raw_parts(self.data, length, length) };
+              }
+              DataManagedType::External => {
+                let mut finalizer: Box<dyn FnOnce(*mut $rust_type, usize)> = Box::new(|_a, _b| {});
+                std::mem::swap(&mut finalizer, &mut self.finalizer_notify);
+                (finalizer)(self.data, self.length);
+              }
+              _ => {}
+            }
+          }
         }
       }
     }
@@ -77,12 +94,12 @@ macro_rules! impl_typed_array {
 
       pub fn new(mut data: Vec<$rust_type>) -> Self {
         let ret = $name {
-          data_reference: Some(Arc::new(())),
           data: data.as_mut_ptr(),
           length: data.len(),
           data_managed_type: DataManagedType::Owned,
           byte_offset: 0,
           raw: None,
+          drop_in_vm: Arc::new(AtomicBool::new(false)),
           finalizer_notify: Box::new(Self::noop_finalize),
         };
         mem::forget(data);
@@ -95,12 +112,12 @@ macro_rules! impl_typed_array {
       {
         let mut data_copied = data.as_ref().to_vec();
         let ret = $name {
-          data_reference: Some(Arc::new(())),
           data: data_copied.as_mut_ptr(),
           length: data.as_ref().len(),
           data_managed_type: DataManagedType::Owned,
           finalizer_notify: Box::new(Self::noop_finalize),
           raw: None,
+          drop_in_vm: Arc::new(AtomicBool::new(false)),
           byte_offset: 0,
         };
         mem::forget(data_copied);
@@ -115,39 +132,29 @@ macro_rules! impl_typed_array {
         F: 'static + FnOnce(*mut $rust_type, usize),
       {
         $name {
-          data_reference: Some(Arc::new(())),
           data,
           length,
           data_managed_type: DataManagedType::External,
           finalizer_notify: Box::new(notify),
           raw: None,
+          drop_in_vm: Arc::new(AtomicBool::new(false)),
           byte_offset: 0,
         }
       }
+    }
 
+    impl Clone for $name {
       /// Clone reference, the inner data is not copied nor moved
-      pub fn clone(&mut self, env: crate::Env) -> crate::Result<$name> {
-        let raw = if let Some((ref_, _)) = self.raw {
-          crate::check_status!(
-            unsafe { sys::napi_reference_ref(env.0, ref_, &mut 0) },
-            "Failed to ref Buffer reference in TypedArray::clone"
-          )?;
-          Some((ref_, env.0))
-        } else {
-          None
-        };
-        Ok(Self {
-          data_reference: match self.data_managed_type {
-            DataManagedType::Owned | DataManagedType::External => self.data_reference.clone(),
-            _ => None,
-          },
+      fn clone(&self) -> $name {
+        Self {
           data: self.data,
           length: self.length,
           data_managed_type: self.data_managed_type,
           finalizer_notify: Box::new(Self::noop_finalize),
-          raw,
+          raw: self.raw,
+          drop_in_vm: self.drop_in_vm.clone(),
           byte_offset: self.byte_offset,
-        })
+        }
       }
     }
 
@@ -244,11 +251,11 @@ macro_rules! impl_typed_array {
           ));
         }
         Ok($name {
-          data_reference: None,
           data: data as *mut $rust_type,
           length,
           byte_offset,
           raw: Some((ref_, env)),
+          drop_in_vm: Arc::new(AtomicBool::new(true)),
           data_managed_type: DataManagedType::Vm,
           finalizer_notify: Box::new(Self::noop_finalize),
         })
@@ -256,14 +263,13 @@ macro_rules! impl_typed_array {
     }
 
     impl ToNapiValue for $name {
-      unsafe fn to_napi_value(env: sys::napi_env, mut val: Self) -> Result<sys::napi_value> {
+      unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
         if let Some((ref_, _)) = val.raw {
           let mut napi_value = std::ptr::null_mut();
           check_status!(
             unsafe { sys::napi_get_reference_value(env, ref_, &mut napi_value) },
             "Failed to delete reference from Buffer"
           )?;
-          val.raw = None;
           return Ok(napi_value);
         }
         let mut arraybuffer_value = ptr::null_mut();
@@ -271,6 +277,7 @@ macro_rules! impl_typed_array {
         let length = val.length * ratio;
         let val_data = val.data;
         let val_length = val.length;
+        val.drop_in_vm.store(true, Ordering::Release);
         let hint_ptr = Box::into_raw(Box::new(val));
         check_status!(
           if length == 0 {
@@ -323,10 +330,7 @@ macro_rules! impl_typed_array {
           Ok(napi_value)
         } else {
           let cloned_value = $name {
-            data_reference: match val.data_managed_type {
-              DataManagedType::Owned | DataManagedType::External => val.data_reference.clone(),
-              _ => None,
-            },
+            drop_in_vm: val.drop_in_vm.clone(),
             data: val.data,
             length: val.length,
             data_managed_type: val.data_managed_type,
@@ -355,13 +359,13 @@ unsafe extern "C" fn finalizer<Data, T: Finalizer<RustType = Data>>(
       // do nothing
     }
     DataManagedType::Owned => {
-      if data.ref_count() == 0 {
+      if data.ref_count() == 1 {
         let length = length;
         unsafe { Vec::from_raw_parts(finalize_data as *mut Data, length, length) };
       }
     }
     DataManagedType::External => {
-      if data.ref_count() == 0 {
+      if data.ref_count() == 1 {
         (finalizer_notify)(finalize_data as *mut Data, length);
       }
     }

crates/napi/src/bindgen_runtime/js_values/buffer.rs

@@ -5,6 +5,7 @@ use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::ptr;
 use std::slice;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 #[cfg(debug_assertions)]
 use std::sync::Mutex;
@@ -21,20 +22,33 @@ thread_local! {
 /// So it is safe to use it in `async fn`, the `&[u8]` under the hood will not be dropped until the `drop` called.
 /// Clone will create a new `Reference` to the same underlying `JavaScript Buffer`.
 pub struct Buffer {
-  pub(crate) data_reference: Option<Arc<()>>,
   pub(crate) inner: &'static mut [u8],
   pub(crate) capacity: usize,
   raw: Option<(sys::napi_ref, sys::napi_env)>,
+  // use it as ref count
+  pub(crate) drop_in_vm: Arc<AtomicBool>,
 }
 
 impl Drop for Buffer {
   fn drop(&mut self) {
-    if let Some((ref_, env)) = self.raw {
-      check_status_or_throw!(
-        env,
-        unsafe { sys::napi_reference_unref(env, ref_, &mut 0) },
-        "Failed to delete Buffer reference in drop"
-      );
+    if Arc::strong_count(&self.drop_in_vm) == 1 {
+      if let Some((ref_, env)) = self.raw {
+        check_status_or_throw!(
+          env,
+          unsafe { sys::napi_reference_unref(env, ref_, &mut 0) },
+          "Failed to unref Buffer reference in drop"
+        );
+        return;
+      }
+      // Drop in Rust side
+      // ```rust
+      // #[napi]
+      // fn buffer_len() -> u32 {
+      //     Buffer::from(vec![1, 2, 3]).len() as u32
+      // }
+      if !self.drop_in_vm.load(Ordering::Acquire) {
+        unsafe { Vec::from_raw_parts(self.inner.as_mut_ptr(), self.inner.len(), self.capacity) };
+      }
     }
   }
 }
@@ -53,10 +67,10 @@ impl Buffer {
       None
     };
     Ok(Self {
-      data_reference: self.data_reference.clone(),
       inner: unsafe { slice::from_raw_parts_mut(self.inner.as_mut_ptr(), self.inner.len()) },
       capacity: self.capacity,
       raw,
+      drop_in_vm: self.drop_in_vm.clone(),
     })
   }
 }
@@ -78,10 +92,10 @@ impl From<Vec<u8>> for Buffer {
     let capacity = data.capacity();
     mem::forget(data);
     Buffer {
-      data_reference: Some(Arc::new(())),
       inner: unsafe { slice::from_raw_parts_mut(inner_ptr, len) },
       capacity,
       raw: None,
+      drop_in_vm: Arc::new(AtomicBool::new(false)),
     }
   }
 }
@@ -149,16 +163,16 @@ impl FromNapiValue for Buffer {
     )?;
 
     Ok(Self {
-      data_reference: None,
       inner: unsafe { slice::from_raw_parts_mut(buf as *mut _, len) },
       capacity: len,
       raw: Some((ref_, env)),
+      drop_in_vm: Arc::new(AtomicBool::new(true)),
     })
   }
 }
 
 impl ToNapiValue for Buffer {
-  unsafe fn to_napi_value(env: sys::napi_env, mut val: Self) -> Result<sys::napi_value> {
+  unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
     // From Node.js value, not from `Vec<u8>`
     if let Some((ref_, _)) = val.raw {
       let mut buf = ptr::null_mut();
@@ -166,14 +180,10 @@ impl ToNapiValue for Buffer {
         unsafe { sys::napi_get_reference_value(env, ref_, &mut buf) },
         "Failed to get Buffer value from reference"
       )?;
-      check_status!(
-        unsafe { sys::napi_delete_reference(env, ref_) },
-        "Failed to delete Buffer reference"
-      )?;
-      val.raw = None; // Prevent double free
       return Ok(buf);
     }
     let len = val.inner.len();
+    val.drop_in_vm.store(true, Ordering::Release);
     let mut ret = ptr::null_mut();
     check_status!(
       if len == 0 {
@@ -210,17 +220,13 @@ impl ToNapiValue for &mut Buffer {
         unsafe { sys::napi_get_reference_value(env, ref_, &mut buf) },
         "Failed to get Buffer value from reference"
       )?;
-      check_status!(
-        unsafe { sys::napi_delete_reference(env, ref_) },
-        "Failed to delete Buffer reference"
-      )?;
       Ok(buf)
     } else {
       let buf = Buffer {
-        data_reference: val.data_reference.clone(),
         inner: unsafe { slice::from_raw_parts_mut(val.inner.as_mut_ptr(), val.capacity) },
         capacity: val.capacity,
         raw: None,
+        drop_in_vm: Arc::new(AtomicBool::new(true)),
       };
       unsafe { ToNapiValue::to_napi_value(env, buf) }
     }

crates/napi/src/bindgen_runtime/mod.rs

@@ -76,14 +76,12 @@ pub unsafe extern "C" fn drop_buffer(
   }
   unsafe {
     let buf = Box::from_raw(finalize_hint as *mut Buffer);
-    if let Some(data_reference) = buf.data_reference.as_ref() {
-      if Arc::strong_count(data_reference) == 0 {
-        mem::drop(Vec::from_raw_parts(
-          finalize_data as *mut u8,
-          buf.inner.len(),
-          buf.capacity,
-        ));
-      }
+    if Arc::strong_count(&buf.drop_in_vm) == 1 {
+      mem::drop(Vec::from_raw_parts(
+        finalize_data as *mut u8,
+        buf.inner.len(),
+        buf.capacity,
+      ));
     }
   }
 }

memory-testing/buffer.mjs

@@ -0,0 +1,26 @@
+import { createRequire } from 'module'
+import { setTimeout } from 'timers/promises'
+
+import { displayMemoryUsageFromNode } from './util.mjs'
+
+const initialMemoryUsage = process.memoryUsage()
+
+const require = createRequire(import.meta.url)
+
+const api = require(`./index.node`)
+
+let i = 1
+// eslint-disable-next-line no-constant-condition
+while (true) {
+  api.bufferLen()
+  api.arrayBufferLen()
+  api.bufferConvert(Buffer.from(Array.from({ length: 1024 * 10240 }).fill(1)))
+  api.arrayBufferConvert(
+    Uint8Array.from(Array.from({ length: 1024 * 10240 }).fill(1)),
+  )
+  if (i % 10 === 0) {
+    await setTimeout(100)
+    displayMemoryUsageFromNode(initialMemoryUsage)
+  }
+  i++
+}

memory-testing/index.mjs

@@ -4,3 +4,4 @@ await createSuite('reference')
 await createSuite('tokio-future')
 await createSuite('serde')
 await createSuite('tsfn')
+await createSuite('buffer')

memory-testing/src/lib.rs

@@ -141,3 +141,23 @@ pub fn leaking_func(env: Env, func: JsFunction) -> napi::Result<()> {
 
   Ok(())
 }
+
+#[napi]
+pub fn buffer_convert(buffer: Buffer) -> Buffer {
+  Buffer::from(vec![0; buffer.len()])
+}
+
+#[napi]
+pub fn buffer_len() -> u32 {
+  Buffer::from(vec![0; 1024 * 10240]).len() as u32
+}
+
+#[napi]
+pub fn array_buffer_convert(array_buffer: Uint8Array) -> Uint8Array {
+  Uint8Array::new(vec![1; array_buffer.len()])
+}
+
+#[napi]
+pub fn array_buffer_len() -> u32 {
+  Uint8Array::new(vec![1; 1024 * 10240]).len() as u32
+}