Fixed locking for non-MVCC snapshots

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.5.1 (unreleased)
 
 - Improved performance of index scans for IVFFlat after updates and deletes
+- Fixed locking for index scans for HNSW
 
 ## 0.5.0 (2023-08-28)
 

Makefile

@@ -8,7 +8,7 @@ HEADERS = src/vector.h
 
 TESTS = $(wildcard test/sql/*.sql)
 REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))
-REGRESS_OPTS = --inputdir=test --load-extension=vector
+REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
 
 OPTFLAGS = -march=native
 

Makefile.win

@@ -5,7 +5,7 @@ OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hn
 HEADERS = src\vector.h
 
 REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
-REGRESS_OPTS = --inputdir=test --load-extension=vector
+REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
 
 # For /arch flags
 # https://learn.microsoft.com/en-us/cpp/build/reference/arch-minimum-cpu-architecture

src/hnsw.h

@@ -218,7 +218,6 @@ typedef HnswNeighborTupleData * HnswNeighborTuple;
 typedef struct HnswScanOpaqueData
 {
 	bool		first;
-	Buffer		buf;
 	List	   *w;
 	MemoryContext tmpCtx;
 

src/hnswscan.c

@@ -101,7 +101,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 	scan = RelationGetIndexScan(index, nkeys, norderbys);
 
 	so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
-	so->buf = InvalidBuffer;
 	so->first = true;
 	so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 									   "Hnsw scan temporary context",
@@ -114,6 +113,12 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
+	/*
+	 * Get a shared lock. This allows vacuum to ensure no in-flight scans
+	 * before marking tuples as deleted.
+	 */
+	LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
+
 	return scan;
 }
 
@@ -164,17 +169,8 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		/* Get scan value */
 		value = GetScanValue(scan);
 
-		/*
-		 * Get a shared lock. This allows vacuum to ensure no in-flight scans
-		 * before marking tuples as deleted.
-		 */
-		LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
-
 		so->w = GetScanItems(scan, value);
 
-		/* Release shared lock */
-		UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
-
 		so->first = false;
 	}
 
@@ -182,7 +178,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 	{
 		HnswCandidate *hc = llast(so->w);
 		ItemPointer heaptid;
-		BlockNumber indexblkno;
 
 		/* Move to next element if no valid heap TIDs */
 		if (list_length(hc->element->heaptids) == 0)
@@ -192,7 +187,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		}
 
 		heaptid = llast(hc->element->heaptids);
-		indexblkno = hc->element->blkno;
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
@@ -204,17 +198,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *heaptid;
 #endif
 
-		/* Unpin buffer */
-		if (BufferIsValid(so->buf))
-			ReleaseBuffer(so->buf);
-
 		/*
-		 * An index scan must maintain a pin on the index page holding the
-		 * item last returned by amgettuple
+		 * Typically, an index scan must maintain a pin on the index page
+		 * holding the item last returned by amgettuple. However, this is not
+		 * needed with the current vacuum strategy, which ensures scans do not
+		 * visit tuples in danger of being marked as deleted.
 		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
-		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -232,9 +223,8 @@ hnswendscan(IndexScanDesc scan)
 {
 	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
 
-	/* Release pin */
-	if (BufferIsValid(so->buf))
-		ReleaseBuffer(so->buf);
+	/* Release shared lock */
+	UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
 
 	MemoryContextDelete(so->tmpCtx);
 

src/hnswvacuum.c

@@ -330,7 +330,10 @@ RepairGraph(HnswVacuumState * vacuumstate)
 	BufferAccessStrategy bas = vacuumstate->bas;
 	BlockNumber blkno = HNSW_HEAD_BLKNO;
 
-	/* Wait for inserts to complete */
+	/*
+	 * Wait for inserts to complete. Inserts before this point may have
+	 * neighbors about to be deleted. Inserts after this point will not.
+	 */
 	LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
 	UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
 
@@ -443,7 +446,11 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 	Relation	index = vacuumstate->index;
 	BufferAccessStrategy bas = vacuumstate->bas;
 
-	/* Wait for selects to complete */
+	/*
+	 * Wait for index scans to complete. Scans before this point may contain
+	 * tuples about to be deleted. Scans after this point will not, since the
+	 * graph has been repaired.
+	 */
 	LockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
 	UnlockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
 

src/ivfflat.c

@@ -182,15 +182,6 @@ ivfflatvalidate(Oid opclassoid)
 	return true;
 }
 
-/*
- * Checks if index-only scan is supported
- */
-static bool
-ivfflatcanreturn(Relation index, int attno)
-{
-	return attno == 1 && IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC) == NULL;
-}
-
 /*
  * Define index handler
  *
@@ -232,7 +223,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
 	amroutine->aminsert = ivfflatinsert;
 	amroutine->ambulkdelete = ivfflatbulkdelete;
 	amroutine->amvacuumcleanup = ivfflatvacuumcleanup;
-	amroutine->amcanreturn = ivfflatcanreturn;
+	amroutine->amcanreturn = NULL;	/* tuple not included in heapsort */
 	amroutine->amcostestimate = ivfflatcostestimate;
 	amroutine->amoptions = ivfflatoptions;
 	amroutine->amproperty = NULL;	/* TODO AMPROP_DISTANCE_ORDERABLE */

src/ivfflat.h

@@ -40,6 +40,9 @@
 #define IVFFLAT_METAPAGE_BLKNO	0
 #define IVFFLAT_HEAD_BLKNO		1	/* first list page */
 
+/* Must correspond to page numbers since page lock is used */
+#define IVFFLAT_SCAN_LOCK		0
+
 /* IVFFlat parameters */
 #define IVFFLAT_DEFAULT_LISTS	100
 #define IVFFLAT_MIN_LISTS		1
@@ -246,9 +249,9 @@ typedef struct IvfflatScanOpaqueData
 	int			probes;
 	int			dimensions;
 	bool		first;
+	bool		hasLock;
 	Buffer		buf;
 	ItemPointerData heaptid;
-	IndexTuple	itup;
 
 	/* Sorting */
 	Tuplesortstate *sortstate;

src/ivfscan.c

@@ -9,6 +9,7 @@
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "storage/bufmgr.h"
+#include "storage/lmgr.h"
 
 /*
  * Compare list distances
@@ -141,7 +142,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				IndexTuple	itup;
 				Datum		datum;
 				bool		isnull;
-				ItemPointerData indextid;
 				ItemId		itemid = PageGetItemId(page, offno);
 
 				/* Skip dead tuples */
@@ -150,7 +150,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 
 				itup = (IndexTuple) PageGetItem(page, itemid);
 				datum = index_getattr(itup, 1, tupdesc, &isnull);
-				ItemPointerSet(&indextid, searchPage, offno);
 
 				/*
 				 * Add virtual tuple
@@ -163,7 +162,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				slot->tts_isnull[0] = false;
 				slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
 				slot->tts_isnull[1] = false;
-				slot->tts_values[2] = PointerGetDatum(&indextid);
+				slot->tts_values[2] = Int32GetDatum((int) searchPage);
 				slot->tts_isnull[2] = false;
 				ExecStoreVirtualTuple(slot);
 
@@ -238,34 +237,6 @@ MarkPriorTupleDead(IndexScanDesc scan)
 	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
 }
 
-/*
- * Set tuple for index-only scan
- */
-static void
-SetIndexTuple(IndexScanDesc scan, ItemPointer indextid)
-{
-	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
-	Buffer		buf = so->buf;
-	Page		page;
-	OffsetNumber offno = ItemPointerGetOffsetNumber(indextid);
-	IndexTuple	itup;
-	Size		itupSize;
-
-	LockBuffer(buf, BUFFER_LOCK_SHARE);
-	page = BufferGetPage(buf);
-	itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
-	itupSize = IndexTupleSize(itup);
-
-	if (so->itup == NULL)
-		so->itup = palloc(BLCKSZ);
-
-	memcpy(so->itup, itup, itupSize);
-
-	scan->xs_itup = so->itup;
-
-	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
-}
-
 /*
  * Prepare for an index scan
  */
@@ -293,8 +264,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
 	so->first = true;
+	so->hasLock = false;
 	ItemPointerSetInvalid(&so->heaptid);
-	so->itup = NULL;
 	so->probes = probes;
 	so->dimensions = dimensions;
 
@@ -311,7 +282,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 #endif
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indextid", TIDOID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
 
 	/* Prep sort */
 	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -326,8 +297,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
-	scan->xs_itupdesc = RelationGetDescr(index);
-
 	return scan;
 }
 
@@ -380,6 +349,13 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		if (scan->orderByData == NULL)
 			elog(ERROR, "cannot scan ivfflat index without order");
 
+		/* Get a shared lock for non-MVCC snapshots */
+		if (!so->hasLock && !IsMVCCSnapshot(scan->xs_snapshot))
+		{
+			so->hasLock = true;
+			LockPage(scan->indexRelation, IVFFLAT_SCAN_LOCK, ShareLock);
+		}
+
 		if (scan->orderByData->sk_flags & SK_ISNULL)
 			value = PointerGetDatum(InitVector(so->dimensions));
 		else
@@ -413,7 +389,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
 	{
 		ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-		ItemPointer indextid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 3, &so->isnull));
+		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
 		scan->xs_heaptid = *heaptid;
@@ -434,11 +410,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
-		so->buf = ReadBuffer(scan->indexRelation, ItemPointerGetBlockNumber(indextid));
-
-		/* Set tuple for index-only scan */
-		if (scan->xs_want_itup)
-			SetIndexTuple(scan, indextid);
+		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -459,8 +431,9 @@ ivfflatendscan(IndexScanDesc scan)
 	if (BufferIsValid(so->buf))
 		ReleaseBuffer(so->buf);
 
-	if (so->itup != NULL)
-		pfree(so->itup);
+	/* Release lock */
+	if (so->hasLock)
+		UnlockPage(scan->indexRelation, IVFFLAT_SCAN_LOCK, ShareLock);
 
 	pairingheap_free(so->listQueue);
 	tuplesort_end(so->sortstate);

src/ivfvacuum.c

@@ -3,6 +3,7 @@
 #include "commands/vacuum.h"
 #include "ivfflat.h"
 #include "storage/bufmgr.h"
+#include "storage/lmgr.h"
 
 /*
  * Bulk delete tuples from the index
@@ -65,14 +66,10 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 
 				vacuum_delay_point();
 
-				buf = ReadBufferExtended(index, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
+				/* Ensure no in-flight index scans for non-MVCC snapshots */
+				LockPage(index, IVFFLAT_SCAN_LOCK, ExclusiveLock);
 
-				/*
-				 * ambulkdelete cannot delete entries from pages that are
-				 * pinned by other backends
-				 *
-				 * https://www.postgresql.org/docs/current/index-locking.html
-				 */
+				buf = ReadBufferExtended(index, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
 				LockBufferForCleanup(buf);
 
 				state = GenericXLogStart(index);
@@ -114,6 +111,8 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 					GenericXLogAbort(state);
 
 				UnlockReleaseBuffer(buf);
+
+				UnlockPage(index, IVFFLAT_SCAN_LOCK, ExclusiveLock);
 			}
 
 			/*