Reduce contention in DocumentsWriterPerThreadPool.

lucene/core/src/java/org/apache/lucene/index/ConcurrentApproximatePriorityQueue.java

@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.index;
+
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Predicate;
+
+/**
+ * Concurrent version of {@link ApproximatePriorityQueue}, which trades a bit more of ordering for
+ * better concurrency by maintaining 8 sub {@link ApproximatePriorityQueue}s that are locked
+ * independently.
+ */
+final class ConcurrentApproximatePriorityQueue<T> {
+
+  /** Keeping 8 queues should already help a lot compared to a single one. */
+  private static final int CONCURRENCY = 8;
+
+  private static final int MASK = 0x07;
+
+  private final Lock[] locks;
+  private final ApproximatePriorityQueue<T>[] queues;
+
+  ConcurrentApproximatePriorityQueue() {
+    locks = new Lock[CONCURRENCY];
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    ApproximatePriorityQueue<T>[] queues = new ApproximatePriorityQueue[CONCURRENCY];
+    this.queues = queues;
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      locks[i] = new ReentrantLock();
+      queues[i] = new ApproximatePriorityQueue<>();
+    }
+  }
+
+  void add(T entry, long weight) {
+    // Seed the order in which to look at entries based on the current thread. This helps distribute
+    // entries across queues and gives a bit of thread affinity between entries and threads, which
+    // can't hurt.
+    final int threadHash = Thread.currentThread().hashCode();
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      final int index = (threadHash + i) & MASK;
+      final Lock lock = locks[index];
+      final ApproximatePriorityQueue<T> queue = queues[index];
+      if (lock.tryLock()) {
+        try {
+          queue.add(entry, weight);
+          return;
+        } finally {
+          lock.unlock();
+        }
+      }
+    }
+    final int index = threadHash & MASK;
+    final Lock lock = locks[index];
+    final ApproximatePriorityQueue<T> queue = queues[index];
+    lock.lock();
+    try {
+      queue.add(entry, weight);
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  T poll(Predicate<T> predicate) {
+    final int threadHash = Thread.currentThread().hashCode();
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      final int index = (threadHash + i) & MASK;
+      final Lock lock = locks[index];
+      final ApproximatePriorityQueue<T> queue = queues[index];
+      if (lock.tryLock()) {
+        try {
+          T entry = queue.poll(predicate);
+          if (entry != null) {
+            return entry;
+          }
+        } finally {
+          lock.unlock();
+        }
+      }
+    }
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      final int index = (threadHash + i) & MASK;
+      final Lock lock = locks[index];
+      final ApproximatePriorityQueue<T> queue = queues[index];
+      lock.lock();
+      try {
+        T entry = queue.poll(predicate);
+        if (entry != null) {
+          return entry;
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+    return null;
+  }
+
+  // Only used for assertions
+  boolean contains(Object o) {
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      locks[i].lock();
+      try {
+        if (queues[i].contains(o)) {
+          return true;
+        }
+      } finally {
+        locks[i].unlock();
+      }
+    }
+    return false;
+  }
+
+  boolean remove(Object o) {
+    for (int i = 0; i < CONCURRENCY; ++i) {
+      locks[i].lock();
+      try {
+        if (queues[i].remove(o)) {
+          return true;
+        }
+      } finally {
+        locks[i].unlock();
+      }
+    }
+    return false;
+  }
+}

lucene/core/src/java/org/apache/lucene/index/DocumentsWriterPerThreadPool.java

@@ -44,11 +44,11 @@ final class DocumentsWriterPerThreadPool implements Iterable<DocumentsWriterPerT
 
   private final Set<DocumentsWriterPerThread> dwpts =
       Collections.newSetFromMap(new IdentityHashMap<>());
-  private final ApproximatePriorityQueue<DocumentsWriterPerThread> freeList =
-      new ApproximatePriorityQueue<>();
+  private final ConcurrentApproximatePriorityQueue<DocumentsWriterPerThread> freeList =
+      new ConcurrentApproximatePriorityQueue<>();
   private final Supplier<DocumentsWriterPerThread> dwptFactory;
   private int takenWriterPermits = 0;
-  private boolean closed;
+  private volatile boolean closed;
 
   DocumentsWriterPerThreadPool(Supplier<DocumentsWriterPerThread> dwptFactory) {
     this.dwptFactory = dwptFactory;
@@ -113,15 +113,17 @@ private synchronized DocumentsWriterPerThread newWriter() {
    * operation (add/updateDocument).
    */
   DocumentsWriterPerThread getAndLock() {
-    synchronized (this) {
-      ensureOpen();
-      DocumentsWriterPerThread dwpt = freeList.poll(DocumentsWriterPerThread::tryLock);
-      if (dwpt == null) {
-        dwpt = newWriter();
-      }
-      // DWPT is already locked before return by this method:
-      return dwpt;
+    ensureOpen();
+    DocumentsWriterPerThread dwpt = freeList.poll(DocumentsWriterPerThread::tryLock);
+    if (dwpt == null) {
+      // newWriter() adds the DWPT to the `dwpts` set as a side-effect. However it is not added to
+      // `freeList` at this point, it will be added later on once DocumentsWriter has indexed a
+      // document into this DWPT and then gives it back to the pool by calling
+      // #marksAsFreeAndUnlock.
+      dwpt = newWriter();
     }
+    // DWPT is already locked before return by this method:
+    return dwpt;
   }
 
   private void ensureOpen() {
@@ -130,13 +132,15 @@ private void ensureOpen() {
     }
   }
 
+  private synchronized boolean contains(DocumentsWriterPerThread state) {
+    return dwpts.contains(state);
+  }
+
   void marksAsFreeAndUnlock(DocumentsWriterPerThread state) {
     final long ramBytesUsed = state.ramBytesUsed();
-    synchronized (this) {
-      assert dwpts.contains(state)
-          : "we tried to add a DWPT back to the pool but the pool doesn't know aobut this DWPT";
-      freeList.add(state, ramBytesUsed);
-    }
+    assert contains(state)
+        : "we tried to add a DWPT back to the pool but the pool doesn't know about this DWPT";
+    freeList.add(state, ramBytesUsed);
     state.unlock();
   }
 
@@ -175,6 +179,9 @@ List<DocumentsWriterPerThread> filterAndLock(Predicate<DocumentsWriterPerThread>
    * @return <code>true</code> iff the given DWPT has been removed. Otherwise <code>false</code>
    */
   synchronized boolean checkout(DocumentsWriterPerThread perThread) {
+    // The DWPT must be held by the current thread. This guarantees that concurrent calls to
+    // #getAndLock cannot pull this DWPT out of the pool since #getAndLock does a DWPT#tryLock to
+    // check if the DWPT is available.
     assert perThread.isHeldByCurrentThread();
     if (dwpts.remove(perThread)) {
       freeList.remove(perThread);