diff --git a/pgvectorscale/src/access_method/scan.rs b/pgvectorscale/src/access_method/scan.rs index a5262e8c..fb794263 100644 --- a/pgvectorscale/src/access_method/scan.rs +++ b/pgvectorscale/src/access_method/scan.rs @@ -54,6 +54,24 @@ impl TSVScanState { } } + /// Release the current iteration's [`StorageState`], if any. + /// + /// The state is owned by this struct: `initialize` publishes it with + /// [`Box::into_raw`] and this reclaims it with [`Box::from_raw`]. Keeping it a + /// raw pointer (rather than an `Option>`) preserves the existing call + /// sites, which need `&mut TSVScanState` at the same time as the state itself. + /// + /// The pointer is nulled out immediately, so calling this repeatedly -- as + /// `amendscan` and then `Drop` do -- can never double free. + fn release_storage(&mut self) { + if !self.storage.is_null() { + // SAFETY: `storage` is null or a pointer from `Box::into_raw` in + // `initialize`; it is nulled below before any other code can observe it. + drop(unsafe { Box::from_raw(self.storage) }); + self.storage = std::ptr::null_mut(); + } + } + fn initialize( &mut self, index: &PgRelation, @@ -83,11 +101,27 @@ impl TSVScanState { } }; - self.storage = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(store_type); + // Release the previous iteration before publishing the new one. `amrescan` + // re-enters `initialize` for every rescan (nested-loop / lateral joins); + // previously each call leaked a fresh `StorageState` -- a response iterator + // plus a cloned quantizer -- into the executor's per-query memory context + // without dropping the one it replaced, so those piled up for the lifetime + // of the whole query instead of the scan. + self.release_storage(); + self.storage = Box::into_raw(Box::new(store_type)); self.distance_fn = Some(distance); } } +impl Drop for TSVScanState { + fn drop(&mut self) { + // Backstop for scans torn down without `amendscan` (e.g. an aborted query): + // the state object itself is released when the per-query memory context is + // reset, and this releases the storage it owns. + self.release_storage(); + } +} + struct ResortData { heap_pointer: HeapPointer, index_pointer: IndexPointer, @@ -437,22 +471,29 @@ fn get_tuple( #[pg_guard] pub extern "C-unwind" fn amendscan(scan: pg_sys::IndexScanDesc) { + let scan: PgBox = unsafe { PgBox::from_pg(scan) }; + let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); + let min_level = unsafe { let l = pg_sys::log_min_messages; let c = pg_sys::client_min_messages; std::cmp::min(l, c) }; if min_level <= pg_sys::DEBUG1 as _ { - let scan: PgBox = unsafe { PgBox::from_pg(scan) }; - let state = - unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); - - let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); - match &mut storage { - StorageState::SbqSpeedup(_bq, iter) => end_scan::(iter), - StorageState::Plain(iter) => end_scan::(iter), + // A scan that is ended without ever having been rescanned has no storage, + // which is not an error -- only report stats when there are some. + if let Some(storage) = unsafe { state.storage.as_mut() } { + match storage { + StorageState::SbqSpeedup(_bq, iter) => end_scan::(iter), + StorageState::Plain(iter) => end_scan::(iter), + } } } + + // Release the scan's state at end of scan instead of deferring to the + // executor's per-query context teardown, so a query that opens many scans does + // not hold every iterator (and cloned quantizer) alive until it completes. + state.release_storage(); } fn end_scan( @@ -474,3 +515,61 @@ fn end_scan( debug_assert_eq!(iter.quantizer_stats.node_reads, 1); debug_assert_eq!(iter.quantizer_stats.node_writes, 0); } + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::{pg_test, spi, Spi}; + + /// Exercises repeated rescans of a diskann index inside a single statement. + /// + /// `amrescan` re-enters [`TSVScanState::initialize`] once per outer row of a + /// nested-loop / lateral join, and `amendscan` ends the scan. Each of those + /// must release the previous iteration's `StorageState` -- the response + /// iterator plus a cloned quantizer -- rather than letting one accumulate per + /// rescan until the whole query finishes. + /// + /// Beyond covering that path, this guards the release itself: an incorrect + /// ownership transfer would double free or use freed memory here and crash the + /// backend, not merely leak. The heap is deliberately churned first so the + /// rescore path also encounters dead / invisible tuples. + #[pg_test] + fn diskann_rescan_releases_state_each_iteration() -> spi::Result<()> { + Spi::run( + "CREATE TABLE rescan_test(id int, embedding vector(3)); + INSERT INTO rescan_test(id, embedding) + SELECT g, ('[' || g || ',' || (g % 7) || ',' || (g % 5) || ']')::vector + FROM generate_series(1, 200) g; + CREATE INDEX ON rescan_test USING diskann (embedding);", + )?; + + // Churn the heap so some index entries point at dead / superseded tuples. + Spi::run("UPDATE rescan_test SET embedding = embedding WHERE id % 3 = 0;")?; + Spi::run("DELETE FROM rescan_test WHERE id % 11 = 0;")?; + + // Make the planner prefer the index so the lateral really does rescan it. + Spi::run("SET enable_seqscan = off;")?; + + // One diskann rescan per outer row: 25 rescans in a single statement. + let neighbors = Spi::get_one::( + "SELECT count(*) FROM ( + SELECT o.id AS oid, n.id AS nid + FROM (SELECT id, embedding FROM rescan_test ORDER BY id LIMIT 25) o + CROSS JOIN LATERAL ( + SELECT i.id + FROM rescan_test i + ORDER BY i.embedding <=> o.embedding + LIMIT 3 + ) n + ) s", + )? + .expect("count should not be null"); + + assert_eq!( + neighbors, 75, + "each of the 25 outer rows should contribute 3 neighbours across rescans" + ); + + Ok(()) + } +} diff --git a/pgvectorscale/src/util/table_slot.rs b/pgvectorscale/src/util/table_slot.rs index 700036a8..53ab18c3 100644 --- a/pgvectorscale/src/util/table_slot.rs +++ b/pgvectorscale/src/util/table_slot.rs @@ -21,6 +21,18 @@ impl TableSlot { std::ptr::null_mut(), )); + // Arm the RAII wrapper IMMEDIATELY, before any fallible work. + // + // `table_slot_create` -> `MakeTupleTableSlot` calls `PinTupleDesc` on the + // heap relation's rowtype descriptor, and that pin is released only by + // `TableSlot::drop` -> `ExecDropSingleTupleTableSlot`. Constructing `Self` + // here means every exit path below — the `!valid` early return, and any + // unwind out of the fetch or the assert — drops the slot and releases the + // descriptor pin. Previously the wrapper was built only on the success + // path, so a heap pointer with no snapshot-visible tuple leaked both the + // slot and its TupleDesc reference on every rescored candidate. + let table_slot = Self { slot }; + let table_am = heap_rel.rd_tableam; let mut ctid: pg_sys::ItemPointerData = pg_sys::ItemPointerData { ..Default::default() @@ -35,7 +47,7 @@ impl TableSlot { scan, &mut ctid, snapshot, - slot.as_ptr(), + table_slot.slot.as_ptr(), &mut call_again, &mut all_dead, ); @@ -45,11 +57,13 @@ impl TableSlot { stats.record_heap_read(); if !valid { - /* no valid tuples found in HOT-chain */ + /* No visible tuple in the HOT chain (deleted / updated-away / not yet + * vacuumed). Dropping `table_slot` here frees the slot and releases its + * rowtype TupleDesc pin. */ return None; } - Some(Self { slot }) + Some(table_slot) } pub unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option { @@ -62,3 +76,97 @@ impl Drop for TableSlot { unsafe { pg_sys::ExecDropSingleTupleTableSlot(self.slot.as_ptr()) }; } } + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::{pg_sys, pg_test, Spi}; + + use super::*; + use crate::access_method::stats::GreedySearchStats; + + /// Parse a `ctid` in its text form, e.g. `"(0,1)"`. + fn parse_ctid(ctid: &str) -> (pg_sys::BlockNumber, pg_sys::OffsetNumber) { + let inner = ctid.trim_matches(|c| c == '(' || c == ')'); + let mut parts = inner.split(','); + let block = parts + .next() + .expect("ctid should have a block number") + .trim() + .parse() + .expect("block number should parse"); + let offset = parts + .next() + .expect("ctid should have an offset") + .trim() + .parse() + .expect("offset number should parse"); + (block, offset) + } + + /// Regression test for the per-scan `TupleDesc` + slot leak (issue #211). + /// + /// [`TableSlot::from_index_heap_pointer`] builds a `TupleTableSlot` with + /// `table_slot_create`, which pins the heap relation's rowtype `TupleDesc`; + /// that pin is released only by [`TableSlot`]'s `Drop`. When the heap pointer + /// has no snapshot-visible tuple the function returns `None`. If the RAII + /// wrapper is armed only on the success path, that early return leaks both the + /// slot and the descriptor pin -- once per dead/invisible rescored candidate -- + /// which surfaces as `resource was not closed: TupleDesc ... (,-1)` at + /// ResourceOwner release and as unbounded backend memory growth on tables with + /// ongoing updates/deletes. + /// + /// The descriptor's reference count must therefore be conserved across a call + /// that takes the no-visible-tuple path. + #[pg_test] + unsafe fn table_slot_releases_tupledesc_on_dead_tuple() { + Spi::run( + "CREATE TABLE slot_leak_test(encoding vector(3)); + INSERT INTO slot_leak_test(encoding) VALUES ('[1,2,3]');", + ) + .unwrap(); + + // Note where the row lives, then delete it, so that a snapshot taken + // afterwards finds no visible version at that TID -- the path under test. + let ctid = Spi::get_one::("SELECT ctid::text FROM slot_leak_test LIMIT 1") + .unwrap() + .expect("the inserted row should exist"); + let (block, offset) = parse_ctid(&ctid); + Spi::run("DELETE FROM slot_leak_test;").unwrap(); + + let heap_oid = Spi::get_one::("SELECT 'slot_leak_test'::regclass::oid") + .unwrap() + .expect("the relation should exist"); + // Open with a lock, exactly as the executor does before any table-AM + // fetch. `PgRelation::with_lock` also closes and unlocks on drop. + let heap_rel = PgRelation::with_lock(heap_oid, pg_sys::AccessShareLock as pg_sys::LOCKMODE); + + let tupdesc = heap_rel.rd_att; + assert!( + (*tupdesc).tdrefcount >= 0, + "the relcache descriptor must be reference-counted for this test to be meaningful" + ); + let refcount_before = (*tupdesc).tdrefcount; + + let mut stats = GreedySearchStats::default(); + let slot = TableSlot::from_index_heap_pointer( + &heap_rel, + HeapPointer::new(block, offset), + // The statement's own registered snapshot. It was taken after the + // DELETE above, so the row has no visible version at that TID. + pg_sys::GetActiveSnapshot(), + &mut stats, + ); + + assert!( + slot.is_none(), + "a deleted row must not yield a visible tuple" + ); + assert_eq!( + (*tupdesc).tdrefcount, + refcount_before, + "the rowtype TupleDesc pin must be released when there is no visible \ + tuple (leak: issue #211)" + ); + } +}