LUCENE-9535: Try to do larger flushes.

DWPTPool currently always returns the last DWPT that was added to the
pool. By returning the largest DWPT instead, we could try to do larger
flushes by finishing DWPTs that are close to being full instead of the
last one that was added to the pool, which might be close to being
empty.

When indexing wikimediumall, this change did not seem to improve the
indexing rate significantly, but it didn't slow things down either and
the number of flushes went from 224-226 to 216, about 4% less.

My expectation is that our nightly benchmarks are a best-case scenario
for DWPTPool as the same number of threads is dedicated to indexing over
time, but in the case when you have e.g. a single fixed threadpool that
is responsible for indexing into several indices, the number of indexing
threads that contribute to a given index might greatly vary over time.

Author: jpountz

lucene/CHANGES.txt

@@ -257,6 +257,9 @@ Improvements
 * LUCENE-9944: Allow DrillSideways users to provide their own CollectorManager without also requiring
   them to provide an ExecutorService. (Greg Miller)
 
+* LUCENE-9535: Improve DocumentsWriterPerThreadPool to prefer larger instances.
+  (Adrien Grand)
+
 Bug fixes
 
 * LUCENE-9686: Fix read past EOF handling in DirectIODirectory. (Zach Chen,

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

@@ -0,0 +1,138 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.function.Predicate;
+
+/**
+ * An approximate priority queue, which attempts to poll items by decreasing log of the weight,
+ * though exact ordering is not guaranteed. This class doesn't support null elements.
+ */
+final class ApproximatePriorityQueue<T> {
+
+  // Indexes between 0 and 63 are sparsely populated, and indexes that are
+  // greater than or equal to 64 are densely populated
+  // Items close to the beginning of this list are more likely to have a
+  // higher weight.
+  private final List<T> slots = new ArrayList<>(Long.SIZE);
+
+  // A bitset where ones indicate that the corresponding index in `slots` is taken.
+  private long usedSlots = 0L;
+
+  ApproximatePriorityQueue() {
+    for (int i = 0; i < Long.SIZE; ++i) {
+      slots.add(null);
+    }
+  }
+
+  /** Add an entry to this queue that has the provided weight. */
+  void add(T entry, long weight) {
+    assert entry != null;
+
+    // The expected slot of an item is the number of leading zeros of its weight,
+    // ie. the larger the weight, the closer an item is to the start of the array.
+    final int expectedSlot = Long.numberOfLeadingZeros(weight);
+
+    // If the slot is already taken, we look for the next one that is free.
+    // The above bitwise operation is equivalent to looping over slots until finding one that is
+    // free.
+    final long freeSlots = ~usedSlots;
+    final int destinationSlot =
+        expectedSlot + Long.numberOfTrailingZeros(freeSlots >>> expectedSlot);
+    assert destinationSlot >= expectedSlot;
+    if (destinationSlot < Long.SIZE) {
+      usedSlots |= 1L << destinationSlot;
+      T previous = slots.set(destinationSlot, entry);
+      assert previous == null;
+    } else {
+      slots.add(entry);
+    }
+  }
+
+  /**
+   * Return an entry matching the predicate. This will usually be one of the available entries that
+   * have the highest weight, though this is not guaranteed. This method returns {@code null} if no
+   * free entries are available.
+   */
+  T poll(Predicate<T> predicate) {
+    // Look at indexes 0..63 first, which are sparsely populated.
+    int nextSlot = 0;
+    do {
+      final int nextUsedSlot = nextSlot + Long.numberOfTrailingZeros(usedSlots >>> nextSlot);
+      if (nextUsedSlot >= Long.SIZE) {
+        break;
+      }
+      final T entry = slots.get(nextUsedSlot);
+      if (predicate.test(entry)) {
+        usedSlots &= ~(1L << nextUsedSlot);
+        slots.set(nextUsedSlot, null);
+        return entry;
+      } else {
+        nextSlot = nextUsedSlot + 1;
+      }
+    } while (nextSlot < Long.SIZE);
+
+    // Then look at indexes 64.. which are densely populated.
+    // Poll in descending order so that if the number of indexing threads
+    // decreases, we keep using the same entry over and over again.
+    // Resizing operations are also less costly on lists when items are closer
+    // to the end of the list.
+    for (ListIterator<T> lit = slots.listIterator(slots.size());
+        lit.previousIndex() >= Long.SIZE; ) {
+      final T entry = lit.previous();
+      if (predicate.test(entry)) {
+        lit.remove();
+        return entry;
+      }
+    }
+
+    // No entry matching the predicate was found.
+    return null;
+  }
+
+  // Only used for assertions
+  boolean contains(Object o) {
+    if (o == null) {
+      throw new NullPointerException();
+    }
+    return slots.contains(o);
+  }
+
+  boolean isEmpty() {
+    return usedSlots == 0 && slots.size() == Long.SIZE;
+  }
+
+  boolean remove(Object o) {
+    if (o == null) {
+      throw new NullPointerException();
+    }
+    int index = slots.indexOf(o);
+    if (index == -1) {
+      return false;
+    }
+    if (index >= Long.SIZE) {
+      slots.remove(index);
+    } else {
+      usedSlots &= ~(1L << index);
+      slots.set(index, null);
+    }
+    return true;
+  }
+}

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

@@ -17,10 +17,8 @@
 package org.apache.lucene.index;
 
 import java.io.Closeable;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Deque;
 import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -46,7 +44,8 @@ final class DocumentsWriterPerThreadPool implements Iterable<DocumentsWriterPerT
 
   private final Set<DocumentsWriterPerThread> dwpts =
       Collections.newSetFromMap(new IdentityHashMap<>());
-  private final Deque<DocumentsWriterPerThread> freeList = new ArrayDeque<>();
+  private final ApproximatePriorityQueue<DocumentsWriterPerThread> freeList =
+      new ApproximatePriorityQueue<>();
   private final Supplier<DocumentsWriterPerThread> dwptFactory;
   private int takenWriterPermits = 0;
   private boolean closed;
@@ -116,21 +115,12 @@ private synchronized DocumentsWriterPerThread newWriter() {
   DocumentsWriterPerThread getAndLock() {
     synchronized (this) {
       ensureOpen();
-      // Important that we are LIFO here! This way if number of concurrent indexing threads was once
-      // high,
-      // but has now reduced, we only use a limited number of DWPTs. This also guarantees that if we
-      // have suddenly
-      // a single thread indexing
-      final Iterator<DocumentsWriterPerThread> descendingIterator = freeList.descendingIterator();
-      while (descendingIterator.hasNext()) {
-        DocumentsWriterPerThread perThread = descendingIterator.next();
-        if (perThread.tryLock()) {
-          descendingIterator.remove();
-          return perThread;
-        }
+      DocumentsWriterPerThread dwpt = freeList.poll(DocumentsWriterPerThread::tryLock);
+      if (dwpt == null) {
+        dwpt = newWriter();
       }
       // DWPT is already locked before return by this method:
-      return newWriter();
+      return dwpt;
     }
   }
 
@@ -141,10 +131,11 @@ private void ensureOpen() {
   }
 
   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);
+      freeList.add(state, ramBytesUsed);
     }
     state.unlock();
   }

lucene/core/src/test/org/apache/lucene/index/TestApproximatePriorityQueue.java

@@ -0,0 +1,105 @@
+/*
+ * 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 org.apache.lucene.util.LuceneTestCase;
+
+public class TestApproximatePriorityQueue extends LuceneTestCase {
+
+  public void testBasics() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(8L, 8L);
+    pq.add(32L, 32L);
+    pq.add(0L, 0L);
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(32L), pq.poll(x -> true));
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(8L), pq.poll(x -> true));
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(0L), pq.poll(x -> true));
+    assertTrue(pq.isEmpty());
+    assertNull(pq.poll(x -> true));
+  }
+
+  public void testPollThenAdd() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(8L, 8L);
+    assertEquals(Long.valueOf(8L), pq.poll(x -> true));
+    assertNull(pq.poll(x -> true));
+    pq.add(0L, 0L);
+    assertEquals(Long.valueOf(0L), pq.poll(x -> true));
+    assertNull(pq.poll(x -> true));
+    pq.add(0L, 0L);
+    assertEquals(Long.valueOf(0L), pq.poll(x -> true));
+    assertNull(pq.poll(x -> true));
+  }
+
+  public void testCollision() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(2L, 2L);
+    pq.add(1L, 1L);
+    pq.add(0L, 0L);
+    pq.add(3L, 3L); // Same nlz as 2
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(2L), pq.poll(x -> true));
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(1L), pq.poll(x -> true));
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(3L), pq.poll(x -> true));
+    assertFalse(pq.isEmpty());
+    assertEquals(Long.valueOf(0L), pq.poll(x -> true));
+    assertTrue(pq.isEmpty());
+    assertNull(pq.poll(x -> true));
+  }
+
+  public void testPollWithPredicate() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(8L, 8L);
+    pq.add(32L, 32L);
+    pq.add(0L, 0L);
+    assertEquals(Long.valueOf(8L), pq.poll(x -> x == 8));
+    assertNull(pq.poll(x -> x == 8));
+    assertFalse(pq.isEmpty());
+  }
+
+  public void testCollisionPollWithPredicate() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(2L, 2L);
+    pq.add(1L, 1L);
+    pq.add(0L, 0L);
+    pq.add(3L, 3L); // Same nlz as 2
+    assertEquals(Long.valueOf(1L), pq.poll(x -> x % 2 == 1));
+    assertEquals(Long.valueOf(3L), pq.poll(x -> x % 2 == 1));
+    assertNull(pq.poll(x -> x % 2 == 1));
+    assertFalse(pq.isEmpty());
+  }
+
+  public void testRemove() {
+    ApproximatePriorityQueue<Long> pq = new ApproximatePriorityQueue<>();
+    pq.add(8L, 8L);
+    pq.add(32L, 32L);
+    pq.add(0L, 0L);
+
+    assertFalse(pq.remove(16L));
+    assertFalse(pq.remove(9L));
+    assertTrue(pq.remove(8L));
+    assertTrue(pq.remove(0L));
+    assertFalse(pq.remove(0L));
+    assertTrue(pq.remove(32L));
+    assertTrue(pq.isEmpty());
+  }
+}