Reduced lock contention with parallel HNSW index builds

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.6.2 (unreleased)
 
-- Improved performance of parallel HNSW index builds
+- Reduced lock contention with parallel HNSW index builds
 
 ## 0.6.1 (2024-03-04)
 

src/hnsw.h

@@ -185,6 +185,7 @@ typedef struct HnswGraph
 
 	/* Entry state */
 	LWLock		entryLock;
+	LWLock		entryWaitLock;
 	HnswElementPtr entryPoint;
 
 	/* Allocations state */

src/hnswbuild.c

@@ -400,7 +400,7 @@ UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswEleme
  * Update graph in memory
  */
 static void
-UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, bool updateEntryPoint, HnswBuildState * buildstate)
+UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
 {
 	HnswGraph  *graph = buildstate->graph;
 	char	   *base = buildstate->hnswarea;
@@ -416,7 +416,7 @@ UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int
 	UpdateNeighborsInMemory(base, procinfo, collation, element, m);
 
 	/* Update entry point if needed (already have lock) */
-	if (updateEntryPoint)
+	if (entryPoint == NULL || element->level > entryPoint->level)
 		HnswPtrStore(base, graph->entryPoint, element);
 }
 
@@ -431,32 +431,42 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 	HnswGraph  *graph = buildstate->graph;
 	HnswElement entryPoint;
 	LWLock	   *entryLock = &graph->entryLock;
+	LWLock	   *entryWaitLock = &graph->entryWaitLock;
 	int			efConstruction = buildstate->efConstruction;
 	int			m = buildstate->m;
 	char	   *base = buildstate->hnswarea;
-	bool		updateEntryPoint;
+
+	/* Wait if another process needs exclusive lock */
+	if (LWLockAcquireOrWait(entryWaitLock, LW_EXCLUSIVE))
+		LWLockRelease(entryWaitLock);
 
 	/* Get entry point */
-	LWLockAcquire(entryLock, LW_EXCLUSIVE);
+	LWLockAcquire(entryLock, LW_SHARED);
 	entryPoint = HnswPtrAccess(base, graph->entryPoint);
 
-	/* Pause new inserts when updating entry point */
+	/* Prevent concurrent inserts when likely updating entry point */
-	/* May still be in-flight inserts */
+	if (entryPoint == NULL || element->level > entryPoint->level)
-	updateEntryPoint = entryPoint == NULL || element->level > entryPoint->level;
+	{
-
+		/* Release shared lock */
-	/* Release entry lock */
-	if (!updateEntryPoint)
 		LWLockRelease(entryLock);
 
+		/* Get exclusive lock */
+		LWLockAcquire(entryWaitLock, LW_EXCLUSIVE);
+		LWLockAcquire(entryLock, LW_EXCLUSIVE);
+		LWLockRelease(entryWaitLock);
+
+		/* 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);
 
 	/* Update graph in memory */
-	UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, updateEntryPoint, buildstate);
+	UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
 
 	/* Release entry lock */
-	if (updateEntryPoint)
+	LWLockRelease(entryLock);
-		LWLockRelease(entryLock);
 }
 
 /*
@@ -609,6 +619,7 @@ InitGraph(HnswGraph * graph, char *base, long memoryTotal)
 	graph->indtuples = 0;
 	SpinLockInit(&graph->lock);
 	LWLockInitialize(&graph->entryLock, hnsw_lock_tranche_id);
+	LWLockInitialize(&graph->entryWaitLock, hnsw_lock_tranche_id);
 	LWLockInitialize(&graph->allocatorLock, hnsw_lock_tranche_id);
 	LWLockInitialize(&graph->flushLock, hnsw_lock_tranche_id);
 }