From 72c20fc14f7668b04e86ad2c9b8184fadf7970f9 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 11 Mar 2024 17:24:42 -0700
Subject: [PATCH] Improved performance of parallel HNSW index builds

---
 CHANGELOG.md    |  4 ++++
 src/hnswbuild.c | 23 ++++++++++-------------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 34af92e..a86b626 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.6.2 (unreleased)
+
+- Improved performance of parallel HNSW index builds
+
 ## 0.6.1 (2024-03-04)
 
 - Fixed error with `ANALYZE` and vectors with different dimensions
diff --git a/src/hnswbuild.c b/src/hnswbuild.c
index 1a2a96c..b0f00d1 100644
--- a/src/hnswbuild.c
+++ b/src/hnswbuild.c
@@ -434,24 +434,20 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 	int			efConstruction = buildstate->efConstruction;
 	int			m = buildstate->m;
 	char	   *base = buildstate->hnswarea;
+	bool		updateEntryPoint;
 
 	/* Get entry point */
-	LWLockAcquire(entryLock, LW_SHARED);
+	LWLockAcquire(entryLock, LW_EXCLUSIVE);
 	entryPoint = HnswPtrAccess(base, graph->entryPoint);
 
-	/* Prevent concurrent inserts when likely updating entry point */
-	if (entryPoint == NULL || element->level > entryPoint->level)
-	{
-		/* Release shared lock */
+	/* Pause new inserts when updating entry point */
+	/* May still be in-flight inserts */
+	updateEntryPoint = entryPoint == NULL || element->level > entryPoint->level;
+
+	/* Release entry lock */
+	if (!updateEntryPoint)
 		LWLockRelease(entryLock);
 
-		/* Get exclusive lock */
-		LWLockAcquire(entryLock, LW_EXCLUSIVE);
-
-		/* Get latest entry point after lock is acquired */
-		entryPoint = HnswPtrAccess(base, graph->entryPoint);
-	}
-
 	/* Find neighbors for element */
 	HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
 
@@ -459,7 +455,8 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 	UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
 
 	/* Release entry lock */
-	LWLockRelease(entryLock);
+	if (updateEntryPoint)
+		LWLockRelease(entryLock);
 }
 
 /*