Skip to content

Commit

Permalink
Reduce allocation rate in HNSW concurrent merge (#14011)
Browse files Browse the repository at this point in the history
The PR removes the allocation of a new `LockedRow` for each locking operation in `HnswLock`. Even though the object was very quickly released, and JIT supports on-stack allocation, it didn't happen in my experiments on OpenJDK 21 - it's easy to avoid the allocation, rather than rely on the JIT.

(cherry picked from commit 532d267)
  • Loading branch information
viliam-durina committed Dec 4, 2024
1 parent 4783c05 commit 45bdbe9
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 24 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import org.apache.lucene.search.TaskExecutor;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.FixedBitSet;
Expand Down Expand Up @@ -56,7 +57,7 @@ public HnswConcurrentMergeBuilder(
this.taskExecutor = taskExecutor;
AtomicInteger workProgress = new AtomicInteger(0);
workers = new ConcurrentMergeWorker[numWorker];
hnswLock = new HnswLock(hnsw);
hnswLock = new HnswLock();
for (int i = 0; i < numWorker; i++) {
workers[i] =
new ConcurrentMergeWorker(
Expand Down Expand Up @@ -221,13 +222,16 @@ private MergeSearcher(NeighborQueue candidates, HnswLock hnswLock, BitSet visite

@Override
void graphSeek(HnswGraph graph, int level, int targetNode) {
try (HnswLock.LockedRow rowLock = hnswLock.read(level, targetNode)) {
NeighborArray neighborArray = rowLock.row();
Lock lock = hnswLock.read(level, targetNode);
try {
NeighborArray neighborArray = ((OnHeapHnswGraph) graph).getNeighbors(level, targetNode);
if (nodeBuffer == null || nodeBuffer.length < neighborArray.size()) {
nodeBuffer = new int[neighborArray.size()];
}
size = neighborArray.size();
if (size >= 0) System.arraycopy(neighborArray.nodes(), 0, nodeBuffer, 0, size);
System.arraycopy(neighborArray.nodes(), 0, nodeBuffer, 0, size);
} finally {
lock.unlock();
}
upto = -1;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.util.Objects;
import java.util.SplittableRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.util.FixedBitSet;
Expand Down Expand Up @@ -338,9 +339,12 @@ private void addDiverseNeighbors(int level, int node, NeighborArray candidates)
}
int nbr = candidates.nodes()[i];
if (hnswLock != null) {
try (HnswLock.LockedRow rowLock = hnswLock.write(level, nbr)) {
NeighborArray nbrsOfNbr = rowLock.row();
Lock lock = hnswLock.write(level, nbr);
try {
NeighborArray nbrsOfNbr = getGraph().getNeighbors(level, nbr);
nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorerSupplier);
} finally {
lock.unlock();
}
} else {
NeighborArray nbrsOfNbr = hnsw.getNeighbors(level, nbr);
Expand Down
26 changes: 8 additions & 18 deletions lucene/core/src/java/org/apache/lucene/util/hnsw/HnswLock.java
Original file line number Diff line number Diff line change
Expand Up @@ -17,49 +17,39 @@

package org.apache.lucene.util.hnsw;

import java.io.Closeable;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
* Provide (read-and-write) locked access to rows of an OnHeapHnswGraph. For use by
* HnswConcurrentMerger and its HnswGraphBuilders.
* Provide (read-and-write) striped locks for access to nodes of an {@link OnHeapHnswGraph}. For use
* by {@link HnswConcurrentMergeBuilder} and its HnswGraphBuilders.
*/
final class HnswLock {
private static final int NUM_LOCKS = 512;
private final ReentrantReadWriteLock[] locks;
private final OnHeapHnswGraph graph;

HnswLock(OnHeapHnswGraph graph) {
this.graph = graph;
HnswLock() {
locks = new ReentrantReadWriteLock[NUM_LOCKS];
for (int i = 0; i < NUM_LOCKS; i++) {
locks[i] = new ReentrantReadWriteLock();
}
}

LockedRow read(int level, int node) {
Lock read(int level, int node) {
int lockid = hash(level, node) % NUM_LOCKS;
Lock lock = locks[lockid].readLock();
lock.lock();
return new LockedRow(graph.getNeighbors(level, node), lock);
return lock;
}

LockedRow write(int level, int node) {
Lock write(int level, int node) {
int lockid = hash(level, node) % NUM_LOCKS;
Lock lock = locks[lockid].writeLock();
lock.lock();
return new LockedRow(graph.getNeighbors(level, node), lock);
return lock;
}

record LockedRow(NeighborArray row, Lock lock) implements Closeable {
@Override
public void close() {
lock.unlock();
}
}

static int hash(int v1, int v2) {
private static int hash(int v1, int v2) {
return v1 * 31 + v2;
}
}

0 comments on commit 45bdbe9

Please sign in to comment.