Fixed locking for non-MVCC snapshots

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 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)

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",
@@ -179,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)
@@ -189,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);
 
@@ -201,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;
@@ -229,10 +223,6 @@ 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);
 

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,6 +249,9 @@ typedef struct IvfflatScanOpaqueData
 	int			probes;
 	int			dimensions;
 	bool		first;
+	bool		hasLock;
+	Buffer		buf;
+	ItemPointerData heaptid;
 
 	/* 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
@@ -143,6 +144,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				bool		isnull;
 				ItemId		itemid = PageGetItemId(page, offno);
 
+				/* Skip dead tuples */
+				if (scan->ignore_killed_tuples && ItemIdIsDead(itemid))
+					continue;
+
 				itup = (IndexTuple) PageGetItem(page, itemid);
 				datum = index_getattr(itup, 1, tupdesc, &isnull);
 
@@ -157,6 +162,8 @@ 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] = Int32GetDatum((int) searchPage);
+				slot->tts_isnull[2] = false;
 				ExecStoreVirtualTuple(slot);
 
 				tuplesort_puttupleslot(so->sortstate, slot);
@@ -181,6 +188,55 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	tuplesort_performsort(so->sortstate);
 }
 
+/*
+ * Mark prior tuple as dead
+ */
+static void
+MarkPriorTupleDead(IndexScanDesc scan)
+{
+	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
+	Buffer		buf = so->buf;
+	Page		page;
+	OffsetNumber maxoffno;
+
+	/* Safety check */
+	if (!BufferIsValid(so->buf) || !ItemPointerIsValid(&so->heaptid))
+		return;
+
+	/* Only a shared locked is needed for ItemIdMarkDead */
+	LockBuffer(buf, BUFFER_LOCK_SHARE);
+	page = BufferGetPage(buf);
+	maxoffno = PageGetMaxOffsetNumber(page);
+
+	for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
+	{
+		ItemId		itemid = PageGetItemId(page, offno);
+		IndexTuple	itup = (IndexTuple) PageGetItem(page, itemid);
+
+		/*
+		 * Find tuple. Since buffer has been pinned, tuple cannot have been
+		 * vacuumed (and heap TID reused).
+		 */
+		if (ItemPointerEquals(&itup->t_tid, &so->heaptid))
+		{
+			/*
+			 * Make sure tuple has not already been marked dead to avoid extra
+			 * WAL if wal_log_hints or data checksums enabled
+			 */
+			if (!ItemIdIsDead(itemid))
+			{
+				ItemIdMarkDead(itemid);
+				MarkBufferDirtyHint(buf, true);
+			}
+
+			break;
+		}
+	}
+
+	/* Unlock buffer */
+	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
+}
+
 /*
  * Prepare for an index scan
  */
@@ -206,7 +262,10 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 		probes = lists;
 
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
+	so->buf = InvalidBuffer;
 	so->first = true;
+	so->hasLock = false;
+	ItemPointerSetInvalid(&so->heaptid);
 	so->probes = probes;
 	so->dimensions = dimensions;
 
@@ -217,12 +276,13 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 
 	/* Create tuple description for sorting */
 #if PG_VERSION_NUM >= 120000
-	so->tupdesc = CreateTemplateTupleDesc(2);
+	so->tupdesc = CreateTemplateTupleDesc(3);
 #else
-	so->tupdesc = CreateTemplateTupleDesc(2, false);
+	so->tupdesc = CreateTemplateTupleDesc(3, false);
 #endif
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", 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);
@@ -254,6 +314,7 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 #endif
 
 	so->first = true;
+	ItemPointerSetInvalid(&so->heaptid);
 	pairingheap_reset(so->listQueue);
 
 	if (keys && scan->numberOfKeys > 0)
@@ -288,9 +349,12 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		if (scan->orderByData == NULL)
 			elog(ERROR, "cannot scan ivfflat index without order");
 
-		/* Safety check */
-		if (!IsMVCCSnapshot(scan->xs_snapshot))
-			elog(ERROR, "non-MVCC snapshots are not supported with ivfflat");
+		/* 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));
@@ -315,10 +379,17 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		if (value != scan->orderByData->sk_argument)
 			pfree(DatumGetPointer(value));
 	}
+	else
+	{
+		/* Mark prior tuple as dead */
+		if (scan->kill_prior_tuple)
+			MarkPriorTupleDead(scan);
+	}
 
 	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
 	{
 		ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
+		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
 		scan->xs_heaptid = *heaptid;
@@ -326,11 +397,20 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *heaptid;
 #endif
 
+		/* Keep track of info needed to mark tuple as dead */
+		so->heaptid = *heaptid;
+
+		/* Unpin buffer */
+		if (BufferIsValid(so->buf))
+			ReleaseBuffer(so->buf);
+
 		/*
-		 * Since no pinning during sorting, requires MVCC-compliant snapshot.
-		 * Also, cannot kill prior tuples, since heap TID may be recycled.
+		 * An index scan must maintain a pin on the index page holding the
+		 * item last returned by amgettuple
+		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
+		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -347,6 +427,14 @@ ivfflatendscan(IndexScanDesc scan)
 {
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
+	/* Release pin */
+	if (BufferIsValid(so->buf))
+		ReleaseBuffer(so->buf);
+
+	/* 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);
 			}
 
 			/*