Reduced lock contention with parallel HNSW index builds

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.6.2 (unreleased)
+
+- Reduced lock contention with parallel HNSW index builds
+
 ## 0.6.1 (2024-03-04)
 
 - Fixed error with `ANALYZE` and vectors with different dimensions

src/hnsw.c

@@ -17,10 +17,7 @@
 #endif
 
 int			hnsw_ef_search;
-int			hnsw_entry_lock_tranche_id;
+int			hnsw_lock_tranche_id;
-int			hnsw_allocator_lock_tranche_id;
-int			hnsw_flush_lock_tranche_id;
-int			hnsw_element_lock_tranche_id;
 static relopt_kind hnsw_relopt_kind;
 
 /*
@@ -39,27 +36,16 @@ HnswInitLockTranche(void)
 	bool		found;
 
 	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-	tranche_ids = ShmemInitStruct("hnsw LWLock ids v2",
+	tranche_ids = ShmemInitStruct("hnsw LWLock ids",
-								  sizeof(int) * 4,
+								  sizeof(int) * 1,
 								  &found);
 	if (!found)
-	{
 		tranche_ids[0] = LWLockNewTrancheId();
-		tranche_ids[1] = LWLockNewTrancheId();
+	hnsw_lock_tranche_id = tranche_ids[0];
-		tranche_ids[2] = LWLockNewTrancheId();
-		tranche_ids[3] = LWLockNewTrancheId();
-	}
-	hnsw_entry_lock_tranche_id = tranche_ids[0];
-	hnsw_allocator_lock_tranche_id = tranche_ids[1];
-	hnsw_flush_lock_tranche_id = tranche_ids[2];
-	hnsw_element_lock_tranche_id = tranche_ids[3];
 	LWLockRelease(AddinShmemInitLock);
 
 	/* Per-backend registration of the tranche ID */
-	LWLockRegisterTranche(hnsw_entry_lock_tranche_id, "HnswEntry");
+	LWLockRegisterTranche(hnsw_lock_tranche_id, "HnswBuild");
-	LWLockRegisterTranche(hnsw_allocator_lock_tranche_id, "HnswAllocator");
-	LWLockRegisterTranche(hnsw_flush_lock_tranche_id, "HnswFlush");
-	LWLockRegisterTranche(hnsw_element_lock_tranche_id, "HnswElement");
 }
 
 /*

src/hnsw.h

@@ -113,10 +113,7 @@
 
 /* Variables */
 extern int	hnsw_ef_search;
-extern int	hnsw_entry_lock_tranche_id;
+extern int	hnsw_lock_tranche_id;
-extern int	hnsw_allocator_lock_tranche_id;
-extern int	hnsw_flush_lock_tranche_id;
-extern int	hnsw_element_lock_tranche_id;
 
 typedef struct HnswElementData HnswElementData;
 typedef struct HnswNeighborArray HnswNeighborArray;
@@ -188,6 +185,7 @@ typedef struct HnswGraph
 
 	/* Entry state */
 	LWLock		entryLock;
+	LWLock		entryWaitLock;
 	HnswElementPtr entryPoint;
 
 	/* Allocations state */

src/hnswbuild.c

@@ -431,10 +431,15 @@ 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;
 
+	/* Wait if another process needs exclusive lock */
+	if (LWLockAcquireOrWait(entryWaitLock, LW_EXCLUSIVE))
+		LWLockRelease(entryWaitLock);
+
 	/* Get entry point */
 	LWLockAcquire(entryLock, LW_SHARED);
 	entryPoint = HnswPtrAccess(base, graph->entryPoint);
@@ -446,7 +451,9 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 		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);
@@ -548,7 +555,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	HnswPtrStore(base, element->value, valuePtr);
 
 	/* Create a lock for the element */
-	LWLockInitialize(&element->lock, hnsw_element_lock_tranche_id);
+	LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
 
 	/* Insert tuple */
 	InsertTupleInMemory(buildstate, element);
@@ -611,9 +618,10 @@ InitGraph(HnswGraph * graph, char *base, long memoryTotal)
 	graph->flushed = false;
 	graph->indtuples = 0;
 	SpinLockInit(&graph->lock);
-	LWLockInitialize(&graph->entryLock, hnsw_entry_lock_tranche_id);
+	LWLockInitialize(&graph->entryLock, hnsw_lock_tranche_id);
-	LWLockInitialize(&graph->allocatorLock, hnsw_allocator_lock_tranche_id);
+	LWLockInitialize(&graph->entryWaitLock, hnsw_lock_tranche_id);
-	LWLockInitialize(&graph->flushLock, hnsw_flush_lock_tranche_id);
+	LWLockInitialize(&graph->allocatorLock, hnsw_lock_tranche_id);
+	LWLockInitialize(&graph->flushLock, hnsw_lock_tranche_id);
 }
 
 /*