From 977b6b331aa287b397c0688bbf5aa045314d4550 Mon Sep 17 00:00:00 2001 From: Piotr Rzysko Date: Sun, 4 May 2025 09:58:59 +0200 Subject: [PATCH] Fix memory leak caused by unclosed virtual table cursors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The following code reproduces the leak, with memory usage increasing over time: ``` #[tokio::main] async fn main() { let db = Builder::new_local(":memory:").build().await.unwrap(); let conn = db.connect().unwrap(); conn.execute("SELECT load_extension('./target/debug/liblimbo_series');", ()) .await .unwrap(); loop { conn.execute("SELECT * FROM generate_series(1,10,2);", ()) .await .unwrap(); } } ``` After switching to the system allocator, the leak becomes detectable with Valgrind: ``` 32,000 bytes in 1,000 blocks are definitely lost in loss record 24 of 24 at 0x538580F: malloc (vg_replace_malloc.c:446) by 0x62E15FA: alloc::alloc::alloc (alloc.rs:99) by 0x62E172C: alloc::alloc::Global::alloc_impl (alloc.rs:192) by 0x62E1530: allocate (alloc.rs:254) by 0x62E1530: alloc::alloc::exchange_malloc (alloc.rs:349) by 0x62E0271: new (boxed.rs:257) by 0x62E0271: open_GenerateSeriesVTab (lib.rs:19) by 0x425D8FA: limbo_core::VirtualTable::open (lib.rs:732) by 0x4285DDA: limbo_core::vdbe::execute::op_vopen (execute.rs:890) by 0x42351E8: limbo_core::vdbe::Program::step (mod.rs:396) by 0x425C638: limbo_core::Statement::step (lib.rs:610) by 0x40DB238: limbo::Statement::execute::{{closure}} (lib.rs:181) by 0x40D9EAF: limbo::Connection::execute::{{closure}} (lib.rs:109) by 0x40D54A1: example::main::{{closure}} (example.rs:26) ``` Interestingly, when using mimalloc, neither Valgrind nor mimalloc’s internal statistics report the leak. --- core/lib.rs | 2 +- core/vdbe/mod.rs | 22 ++++++++++++++++++---- extensions/core/src/vtabs.rs | 6 ++++++ macros/src/lib.rs | 13 +++++++++++++ 4 files changed, 38 insertions(+), 5 deletions(-) diff --git a/core/lib.rs b/core/lib.rs index 2f7ab0577..f36072e3a 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -730,7 +730,7 @@ impl VirtualTable { pub fn open(&self) -> crate::Result { let cursor = unsafe { (self.implementation.open)(self.implementation.ctx) }; - VTabOpaqueCursor::new(cursor) + VTabOpaqueCursor::new(cursor, self.implementation.close) } #[tracing::instrument(skip(cursor))] diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 2adf438b6..d90504df4 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -203,20 +203,34 @@ impl Bitfield { } } -pub struct VTabOpaqueCursor(*const c_void); +type VTabOpaqueCursorCloseFn = unsafe extern "C" fn(*const c_void) -> limbo_ext::ResultCode; + +pub struct VTabOpaqueCursor { + cursor: *const c_void, + close: VTabOpaqueCursorCloseFn, +} impl VTabOpaqueCursor { - pub fn new(cursor: *const c_void) -> Result { + pub fn new(cursor: *const c_void, close: VTabOpaqueCursorCloseFn) -> Result { if cursor.is_null() { return Err(LimboError::InternalError( "VTabOpaqueCursor: cursor is null".into(), )); } - Ok(Self(cursor)) + Ok(Self { cursor, close }) } pub fn as_ptr(&self) -> *const c_void { - self.0 + self.cursor + } +} + +impl Drop for VTabOpaqueCursor { + fn drop(&mut self) { + let result = unsafe { (self.close)(self.cursor) }; + if !result.is_ok() { + tracing::error!("Failed to close virtual table cursor"); + } } } diff --git a/extensions/core/src/vtabs.rs b/extensions/core/src/vtabs.rs index 5d86457f7..de5794c6e 100644 --- a/extensions/core/src/vtabs.rs +++ b/extensions/core/src/vtabs.rs @@ -15,6 +15,7 @@ pub struct VTabModuleImpl { pub name: *const c_char, pub create_schema: VtabFnCreateSchema, pub open: VtabFnOpen, + pub close: VtabFnClose, pub filter: VtabFnFilter, pub column: VtabFnColumn, pub next: VtabFnNext, @@ -44,6 +45,8 @@ pub type VtabFnCreateSchema = unsafe extern "C" fn(args: *const Value, argc: i32 pub type VtabFnOpen = unsafe extern "C" fn(*const c_void) -> *const c_void; +pub type VtabFnClose = unsafe extern "C" fn(cursor: *const c_void) -> ResultCode; + pub type VtabFnFilter = unsafe extern "C" fn( cursor: *const c_void, argc: i32, @@ -134,6 +137,9 @@ pub trait VTabCursor: Sized { fn column(&self, idx: u32) -> Result; fn eof(&self) -> bool; fn next(&mut self) -> ResultCode; + fn close(&self) -> ResultCode { + ResultCode::OK + } } #[repr(u8)] diff --git a/macros/src/lib.rs b/macros/src/lib.rs index d47101589..2d6694e10 100644 --- a/macros/src/lib.rs +++ b/macros/src/lib.rs @@ -448,6 +448,7 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream { let register_fn_name = format_ident!("register_{}", struct_name); let create_schema_fn_name = format_ident!("create_schema_{}", struct_name); let open_fn_name = format_ident!("open_{}", struct_name); + let close_fn_name = format_ident!("close_{}", struct_name); let filter_fn_name = format_ident!("filter_{}", struct_name); let column_fn_name = format_ident!("column_{}", struct_name); let next_fn_name = format_ident!("next_{}", struct_name); @@ -486,6 +487,17 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream { } } + #[no_mangle] + unsafe extern "C" fn #close_fn_name( + cursor: *const ::std::ffi::c_void + ) -> ::limbo_ext::ResultCode { + if cursor.is_null() { + return ::limbo_ext::ResultCode::Error; + } + let boxed_cursor = ::std::boxed::Box::from_raw(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor); + boxed_cursor.close() + } + #[no_mangle] unsafe extern "C" fn #filter_fn_name( cursor: *const ::std::ffi::c_void, @@ -649,6 +661,7 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream { name: name_c, create_schema: Self::#create_schema_fn_name, open: Self::#open_fn_name, + close: Self::#close_fn_name, filter: Self::#filter_fn_name, column: Self::#column_fn_name, next: Self::#next_fn_name,