BUGFIX: per-scan TupleDesc/slot leak and per-rescan state leak

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<Box<_>>`) 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<pg_sys::IndexScanDescData> = 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<pg_sys::IndexScanDescData> = 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::<SbqSpeedupStorage>(iter),
-            StorageState::Plain(iter) => end_scan::<PlainStorage>(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::<SbqSpeedupStorage>(iter),
+                StorageState::Plain(iter) => end_scan::<PlainStorage>(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<S: Storage>(
@@ -474,3 +515,61 @@ fn end_scan<S: Storage>(
     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::<i64>(
+            "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(())
+    }
+}

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<Datum> {
@@ -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 ... (<oid>,-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::<String>("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::<pg_sys::Oid>("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)"
+        );
+    }
+}