From 0e1de45463a8c805c45b49892bd1a3e7f9147fb4 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 21 Aug 2023 15:42:59 -0700
Subject: [PATCH] Improved locking code [skip ci]

---
 src/hnsw.h       |  4 ++++
 src/hnswinsert.c |  8 ++++----
 src/hnswscan.c   | 12 ++++++++++--
 src/hnswvacuum.c |  8 ++++----
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/hnsw.h b/src/hnsw.h
index 3acbe02..07b5b08 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -28,6 +28,10 @@
 #define HNSW_METAPAGE_BLKNO	0
 #define HNSW_HEAD_BLKNO		1	/* first element page */
 
+/* Must correspond to page numbers since page lock is used */
+#define HNSW_UPDATE_LOCK 	0
+#define HNSW_SCAN_LOCK		1
+
 #define HNSW_DEFAULT_M	16
 #define HNSW_MIN_M	2
 #define HNSW_MAX_M		100
diff --git a/src/hnswinsert.c b/src/hnswinsert.c
index ab5b98a..05facb9 100644
--- a/src/hnswinsert.c
+++ b/src/hnswinsert.c
@@ -519,7 +519,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 	 * Get a shared lock for the duration of the insert. Use a page lock so it
 	 * does not interfere with buffer lock (or reads when vacuuming).
 	 */
-	LockPage(index, HNSW_METAPAGE_BLKNO, lockmode);
+	LockPage(index, HNSW_UPDATE_LOCK, lockmode);
 
 	/* Get entry point */
 	entryPoint = HnswGetEntryPoint(index);
@@ -528,11 +528,11 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 	if (entryPoint == NULL || element->level > entryPoint->level)
 	{
 		/* Release shared lock */
-		UnlockPage(index, HNSW_METAPAGE_BLKNO, lockmode);
+		UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
 
 		/* Get exclusive lock */
 		lockmode = ExclusiveLock;
-		LockPage(index, HNSW_METAPAGE_BLKNO, lockmode);
+		LockPage(index, HNSW_UPDATE_LOCK, lockmode);
 
 		/* Get latest entry point after lock is acquired */
 		entryPoint = HnswGetEntryPoint(index);
@@ -548,7 +548,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 	WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint);
 
 	/* Release shared lock */
-	UnlockPage(index, HNSW_METAPAGE_BLKNO, lockmode);
+	UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
 
 	return true;
 }
diff --git a/src/hnswscan.c b/src/hnswscan.c
index 040738b..dba3ca9 100644
--- a/src/hnswscan.c
+++ b/src/hnswscan.c
@@ -145,9 +145,17 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 				HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
 		}
 
-		LockPage(scan->indexRelation, HNSW_HEAD_BLKNO, ShareLock);
+		/*
+		 * 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);
-		UnlockPage(scan->indexRelation, HNSW_HEAD_BLKNO, ShareLock);
+
+		/* Release shared lock */
+		UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
+
 		so->first = false;
 	}
 
diff --git a/src/hnswvacuum.c b/src/hnswvacuum.c
index c6dd91b..37b4ebe 100644
--- a/src/hnswvacuum.c
+++ b/src/hnswvacuum.c
@@ -316,8 +316,8 @@ RepairGraph(HnswVacuumState * vacuumstate)
 	RepairGraphEntryPoint(vacuumstate);
 
 	/* Wait for inserts to complete */
-	LockPage(index, HNSW_METAPAGE_BLKNO, ExclusiveLock);
-	UnlockPage(index, HNSW_METAPAGE_BLKNO, ExclusiveLock);
+	LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
+	UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
 
 	while (BlockNumberIsValid(blkno))
 	{
@@ -405,8 +405,8 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 	BufferAccessStrategy bas = vacuumstate->bas;
 
 	/* Wait for selects to complete */
-	LockPage(index, HNSW_HEAD_BLKNO, ExclusiveLock);
-	UnlockPage(index, HNSW_HEAD_BLKNO, ExclusiveLock);
+	LockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
+	UnlockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
 
 	while (BlockNumberIsValid(blkno))
 	{