From cad48d92036c7b587022b6cadf2fa45bc457bdae Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 16 Jan 2024 13:34:55 -0800 Subject: [PATCH] Improved locking --- src/hnsw.c | 64 +++++++++++++++++++++++-------------------------- src/hnsw.h | 8 +++---- src/hnswbuild.c | 18 +++++++------- src/hnswutils.c | 4 ++-- 4 files changed, 44 insertions(+), 50 deletions(-) diff --git a/src/hnsw.c b/src/hnsw.c index 2ea5895..5734c6c 100644 --- a/src/hnsw.c +++ b/src/hnsw.c @@ -16,9 +16,35 @@ int hnsw_ef_search; static relopt_kind hnsw_relopt_kind; -int entryLockTrancheId; -int allocatorLockTrancheId; -int flushLockTrancheId; +int hnsw_lock_tranche_id; + +/* + * Assign tranche IDs for our LWLocks. This only needs to be done by one + * backend, the tranche IDs are remembered in shared memory. + * + * This shared memory area is very small, so we just allocate it from the + * "slop" that PostgreSQL reserves for small allocations like this. If + * this grows bigger, we should use a shmem_request_hook and + * RequestAddinShmemSpace() to pre-reserve space for this. + */ +static void +HnswInitLockTranche(void) +{ + int *tranche_ids; + bool found; + + LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); + tranche_ids = ShmemInitStruct("hnsw LWLock ids", + sizeof(int) * 1, + &found); + if (!found) + tranche_ids[0] = LWLockNewTrancheId(); + hnsw_lock_tranche_id = tranche_ids[0]; + LWLockRelease(AddinShmemInitLock); + + /* Per-backend registration of the tranche IDs */ + LWLockRegisterTranche(hnsw_lock_tranche_id, "HnswBuild"); +} /* * Initialize index options and variables @@ -26,37 +52,7 @@ int flushLockTrancheId; void HnswInit(void) { - int *tranche_ids; - bool found; - - /* - * Assign tranche IDs for our LWLocks. This only needs to be done by one - * backend, the tranche IDs are remembered in shared memory. - * - * This shared memory area is very small, so we just allocate it from the - * "slop" that PostgreSQL reserves for small allocations like this. If - * this grows bigger, we should use a shmem_request_hook and - * RequestAddinShmemSpace() to pre-reserve space for this. - */ - LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); - tranche_ids = ShmemInitStruct("pgvector LWLock ids", - sizeof(int) * 3, - &found); - if (!found) - { - tranche_ids[0] = LWLockNewTrancheId(); - tranche_ids[1] = LWLockNewTrancheId(); - tranche_ids[2] = LWLockNewTrancheId(); - } - entryLockTrancheId = tranche_ids[0]; - allocatorLockTrancheId = tranche_ids[1]; - flushLockTrancheId = tranche_ids[2]; - LWLockRelease(AddinShmemInitLock); - - /* Per-backend registration of the tranche IDs */ - LWLockRegisterTranche(entryLockTrancheId, "pgvector entryLock"); - LWLockRegisterTranche(allocatorLockTrancheId, "pgvector allocatorLock"); - LWLockRegisterTranche(flushLockTrancheId, "pgvector flushLock"); + HnswInitLockTranche(); hnsw_relopt_kind = add_reloption_kind(); add_int_reloption(hnsw_relopt_kind, "m", "Max number of connections", diff --git a/src/hnsw.h b/src/hnsw.h index 9627c71..55986e2 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -116,10 +116,8 @@ /* Variables */ extern int hnsw_ef_search; -/* These are initialized when the module is loaded */ -extern int entryLockTrancheId; -extern int allocatorLockTrancheId; -extern int flushLockTrancheId; +/* This is initialized when the module is loaded */ +extern int hnsw_lock_tranche_id; typedef struct HnswElementData HnswElementData; typedef struct HnswNeighborArray HnswNeighborArray; @@ -149,7 +147,7 @@ typedef struct HnswElementData OffsetNumber neighborOffno; BlockNumber neighborPage; DatumPtr value; - slock_t lock; + LWLock lock; } HnswElementData; typedef HnswElementData * HnswElement; diff --git a/src/hnswbuild.c b/src/hnswbuild.c index fcc22b3..b3ce6ea 100644 --- a/src/hnswbuild.c +++ b/src/hnswbuild.c @@ -297,17 +297,17 @@ FlushPages(HnswBuildState * buildstate) static bool HnswAddDuplicateInMemory(HnswElement element, HnswElement dup) { - SpinLockAcquire(&dup->lock); + LWLockAcquire(&dup->lock, LW_EXCLUSIVE); if (dup->heaptidsLength == HNSW_HEAPTIDS) { - SpinLockRelease(&dup->lock); + LWLockRelease(&dup->lock); return false; } HnswAddHeapTid(dup, &element->heaptids[0]); - SpinLockRelease(&dup->lock); + LWLockRelease(&dup->lock); return true; } @@ -371,9 +371,9 @@ HnswUpdateNeighborPagesInMemory(char *base, FmgrInfo *procinfo, Oid collation, H Assert(neighborElement); /* Use element for lock instead of hc since hc can be replaced */ - SpinLockAcquire(&neighborElement->lock); + LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE); HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation); - SpinLockRelease(&neighborElement->lock); + LWLockRelease(&neighborElement->lock); } } } @@ -482,7 +482,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn HnswPtrStore(base, element->value, valuePtr); /* Create element lock */ - SpinLockInit(&element->lock); + LWLockInitialize(&element->lock, hnsw_lock_tranche_id); /* Get entry point */ LWLockAcquire(&graph->entryLock, LW_EXCLUSIVE); @@ -558,9 +558,9 @@ InitGraph(HnswGraph * graph, char *base, long memoryTotal) graph->flushed = false; graph->indtuples = 0; SpinLockInit(&graph->lock); - LWLockInitialize(&graph->entryLock, entryLockTrancheId); - LWLockInitialize(&graph->allocatorLock, allocatorLockTrancheId); - LWLockInitialize(&graph->flushLock, flushLockTrancheId); + LWLockInitialize(&graph->entryLock, hnsw_lock_tranche_id); + LWLockInitialize(&graph->allocatorLock, hnsw_lock_tranche_id); + LWLockInitialize(&graph->flushLock, hnsw_lock_tranche_id); } /* diff --git a/src/hnswutils.c b/src/hnswutils.c index d291640..811f2ec 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -780,9 +780,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F /* Copy neighborhood to local memory if needed */ if (index == NULL) { - SpinLockAcquire(&cElement->lock); + LWLockAcquire(&cElement->lock, LW_EXCLUSIVE); memcpy(neighborhoodData, neighborhood, neighborhoodSize); - SpinLockRelease(&cElement->lock); + LWLockRelease(&cElement->lock); neighborhood = neighborhoodData; }