fix: check mvcc in heap fetch (#336)

usamoi · usamoi · commit d2f87ba8b056 · 2025-09-15T10:27:30.000+08:00
#335 --------- Signed-off-by: usamoi <usamoi@outlook.com>
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
@@ -137,28 +137,28 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Clippy
-        run: cargo clippy --workspace --exclude vchord
+        run: cargo clippy --locked --workspace --exclude vchord
 
       - name: Cargo Test
-        run: cargo test --workspace --exclude vchord --exclude simd --no-fail-fast
+        run: cargo test --locked --workspace --exclude vchord --exclude simd --no-fail-fast
 
       - name: Cargo Test (simd)
         run: |
           if [ "$(uname -m)" == "x86_64" ]; then
             cargo \
               --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \
-              test -p simd -- --no-capture
+              test --locked -p simd -- --no-capture
           fi
           if [ "$(uname -m)" == "aarch64" ]; then
             cargo \
               --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=16"]' \
-              test -p simd -- --no-capture
+              test --locked -p simd -- --no-capture
             cargo \
               --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=32"]' \
-              test -p simd -- --no-capture
+              test --locked -p simd -- --no-capture
             cargo \
               --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=64"]' \
-              test -p simd -- --no-capture
+              test --locked -p simd -- --no-capture
           fi
 
   psql:
@@ -218,7 +218,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Clippy
-        run: cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps
+        run: cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps
 
       - name: Install
         run: |
@@ -301,7 +301,7 @@ jobs:
 
       - name: Clippy
         run: |
-          cargo clippy -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps
+          cargo clippy --locked -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps
 
       - name: Install
         run: |
@@ -407,7 +407,7 @@ jobs:
 
       - name: Clippy
         run: |
-          cargo clippy --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps
+          cargo clippy --locked --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps
 
       - name: Install
         run: |
@@ -514,7 +514,7 @@ jobs:
       - name: Clippy
         run: |
           . "$HOME/.cargo/env"
-          cargo clippy --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps
+          cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps
 
       - name: Install
         run: |
@@ -711,10 +711,10 @@ jobs:
       - name: Clippy & Test
         run: |
           PGRX_PG_CONFIG_PATH=pg_config \
-            cargo clippy --target ${{ matrix.rust_triple }} \
+            cargo clippy --locked --target ${{ matrix.rust_triple }} \
             --workspace --features pg${{ matrix.version }}
           PGRX_PG_CONFIG_PATH=pg_config \
-            cargo test --target ${{ matrix.rust_triple }} \
+            cargo test --locked --target ${{ matrix.rust_triple }} \
             --workspace --exclude vchord  --no-fail-fast \
             -- --no-capture
 
diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs
@@ -168,7 +168,8 @@ fn build(
 ) -> Result<PathBuf, Box<dyn Error>> {
     let mut command = Command::new("cargo");
     command
-        .args(["build", "-p", "vchord", "--lib"])
+        .args(["build", "--locked"])
+        .args(["-p", "vchord", "--lib"])
         .args(["--profile", profile])
         .args(["--target", target])
         .args(["--features", pg_version])
@@ -265,7 +266,8 @@ fn generate(
     )?;
     let mut command = Command::new("cargo");
     command
-        .args(["rustc", "-p", "vchord", "--bin", "pgrx_embed_vchord"])
+        .args(["rustc", "--locked"])
+        .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"])
         .args(["--profile", profile])
         .args(["--target", target])
         .args(["--features", pg_version])
diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs
@@ -48,6 +48,7 @@ pub struct HeapFetcher {
     econtext: *mut pgrx::pg_sys::ExprContext,
     heap_relation: pgrx::pg_sys::Relation,
     snapshot: pgrx::pg_sys::Snapshot,
+    heapfetch: *mut pgrx::pg_sys::IndexFetchTableData,
     slot: *mut pgrx::pg_sys::TupleTableSlot,
     values: [Datum; 32],
     is_nulls: [bool; 32],
@@ -59,6 +60,7 @@ impl HeapFetcher {
         index_relation: pgrx::pg_sys::Relation,
         heap_relation: pgrx::pg_sys::Relation,
         snapshot: pgrx::pg_sys::Snapshot,
+        heapfetch: *mut pgrx::pg_sys::IndexFetchTableData,
         hack: *mut pgrx::pg_sys::IndexScanState,
     ) -> Self {
         unsafe {
@@ -71,6 +73,7 @@ impl HeapFetcher {
                 econtext,
                 heap_relation,
                 snapshot,
+                heapfetch,
                 slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()),
                 values: [Datum::null(); 32],
                 is_nulls: [true; 32],
@@ -99,16 +102,49 @@ impl Fetcher for HeapFetcher {
             use pgrx::pg_sys::ffi::pg_guard_ffi_boundary;
             let mut ctid = key_to_ctid(key);
             let table_am = (*self.heap_relation).rd_tableam;
-            let fetch_row_version = (*table_am)
-                .tuple_fetch_row_version
+            let index_fetch_tuple = (*table_am)
+                .index_fetch_tuple
                 .expect("unsupported heap access method");
-            #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")]
-            if pg_guard_ffi_boundary(|| {
-                !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot)
-            }) {
-                return None;
+            let found = 'a: {
+                let mut call_again = false;
+                let mut all_dead = false;
+                #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")]
+                let found = pg_guard_ffi_boundary(|| {
+                    index_fetch_tuple(
+                        self.heapfetch,
+                        &mut ctid,
+                        self.snapshot,
+                        self.slot,
+                        &mut call_again,
+                        &mut all_dead,
+                    )
+                });
+                if found {
+                    break 'a true;
+                }
+                while call_again {
+                    #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")]
+                    let found = pg_guard_ffi_boundary(|| {
+                        index_fetch_tuple(
+                            self.heapfetch,
+                            &mut ctid,
+                            self.snapshot,
+                            self.slot,
+                            &mut call_again,
+                            &mut all_dead,
+                        )
+                    });
+                    if found {
+                        break 'a true;
+                    }
+                }
+                false
+            };
+            if found {
+                Some(HeapTuple { this: self })
+            } else {
+                None
             }
-            Some(HeapTuple { this: self })
         }
     }
 }
diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs
@@ -365,6 +365,7 @@ pub unsafe extern "C-unwind" fn amrescan(
                     (*scan).indexRelation,
                     (*scan).heapRelation,
                     (*scan).xs_snapshot,
+                    (*scan).xs_heapfetch,
                     if let Some(hack) = hack {
                         hack.as_ptr()
                     } else {
diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs
@@ -448,6 +448,7 @@ pub unsafe extern "C-unwind" fn amrescan(
                     (*scan).indexRelation,
                     (*scan).heapRelation,
                     (*scan).xs_snapshot,
+                    (*scan).xs_heapfetch,
                     if let Some(hack) = hack {
                         hack.as_ptr()
                     } else {
diff --git a/tests/vchordrq/issue_335.slt b/tests/vchordrq/issue_335.slt
@@ -0,0 +1,37 @@
+statement ok
+SET enable_seqscan = off;
+
+statement ok
+CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3), visible boolean default true);
+
+statement ok
+INSERT INTO items (embedding) SELECT ARRAY[i * 0.001, i * 0.001, i * 0.001]::real[] FROM generate_series(1, 1000) s(i);
+
+statement ok
+CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops);
+
+statement ok
+SET vchordrq.prefilter = on;
+
+query I
+SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3;
+----
+3
+4
+2
+
+statement ok
+UPDATE items SET visible = false WHERE id = 3;
+
+statement ok
+UPDATE items SET visible = true WHERE id = 3;
+
+query I
+SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3;
+----
+3
+4
+2
+
+statement ok
+DROP TABLE items;
