naskya/napi-rs
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Merge pull request #1235 from napi-rs/buffer-leak

Browse source 
fix(napi): memory leak in Buffer/ArrayBuffer
This commit is contained in:
LongYinan 2022-07-11 21:48:42 +08:00 committed by GitHub
commit 40c132aefc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 127 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs
View file

@ -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);
}
}

48
crates/napi/src/bindgen_runtime/js_values/buffer.rs
View file

@ -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) }
}

14
crates/napi/src/bindgen_runtime/mod.rs
View file

@ -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,
));
}
}
}

26
memory-testing/buffer.mjs Normal file
View file

@ -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++
}

1
memory-testing/index.mjs
View file

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

20
memory-testing/src/lib.rs
View file

@ -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
}