diff --git a/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs b/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs index 22f0be4c..332789e7 100644 --- a/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs +++ b/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>, 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, finalizer_notify: Box, } @@ -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 = 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 { + unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result { 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>( // 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); } } diff --git a/crates/napi/src/bindgen_runtime/js_values/buffer.rs b/crates/napi/src/bindgen_runtime/js_values/buffer.rs index c5879510..07157955 100644 --- a/crates/napi/src/bindgen_runtime/js_values/buffer.rs +++ b/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>, 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, } 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> 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 { + unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result { // From Node.js value, not from `Vec` 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) } } diff --git a/crates/napi/src/bindgen_runtime/mod.rs b/crates/napi/src/bindgen_runtime/mod.rs index 7b418472..60630802 100644 --- a/crates/napi/src/bindgen_runtime/mod.rs +++ b/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, + )); } } } diff --git a/memory-testing/buffer.mjs b/memory-testing/buffer.mjs new file mode 100644 index 00000000..83926a0d --- /dev/null +++ b/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++ +} diff --git a/memory-testing/index.mjs b/memory-testing/index.mjs index 6dddeb9f..332a5477 100644 --- a/memory-testing/index.mjs +++ b/memory-testing/index.mjs @@ -4,3 +4,4 @@ await createSuite('reference') await createSuite('tokio-future') await createSuite('serde') await createSuite('tsfn') +await createSuite('buffer') diff --git a/memory-testing/src/lib.rs b/memory-testing/src/lib.rs index 9720747f..31befcf9 100644 --- a/memory-testing/src/lib.rs +++ b/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 +}