Refactor ElementWithReleasableBudget

Author: albertzaharovits

server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorService.java

@@ -571,7 +571,7 @@ MergeTask peekQueue() {
     static class PriorityBlockingQueueWithBudget<E> {
         private final ToLongFunction<? super E> budgetFunction;
         protected final PriorityQueue<E> enqueuedByBudget;
-        private final IdentityHashMap<Wrap<E>, Long> unreleasedBudgetPerElement;
+        private final IdentityHashMap<ElementWithReleasableBudget, Long> unreleasedBudgetPerElement;
         private final ReentrantLock lock;
         private final Condition elementAvailable;
         protected long availableBudget;
@@ -612,7 +612,7 @@ ElementWithReleasableBudget take() throws InterruptedException {
                     elementAvailable.await();
                 }
                 // deducts and holds up that element's budget from the available budget
-                return new ElementWithReleasableBudget(enqueuedByBudget.poll(), peekBudget);
+                return newElementWithReleasableBudget(enqueuedByBudget.poll(), peekBudget);
             } finally {
                 lock.unlock();
             }
@@ -647,23 +647,40 @@ int queueSize() {
             return enqueuedByBudget.size();
         }
 
+        private ElementWithReleasableBudget newElementWithReleasableBudget(E element, long budget) {
+            ElementWithReleasableBudget elementWithReleasableBudget = new ElementWithReleasableBudget(element);
+            assert this.lock.isHeldByCurrentThread();
+            // the taken element holds up some budget
+            var prev = this.unreleasedBudgetPerElement.put(elementWithReleasableBudget, budget);
+            assert prev == null;
+            this.availableBudget -= budget;
+            assert this.availableBudget >= 0L;
+            return elementWithReleasableBudget;
+        }
+
+        private void release(ElementWithReleasableBudget elementWithReleasableBudget) {
+            final ReentrantLock lock = this.lock;
+            lock.lock();
+            try {
+                assert elementWithReleasableBudget.isClosed() == false;
+                // when the taken element is not used anymore, it will not influence subsequent computations for available budget,
+                // but its allotted budget is not yet released
+                var val = unreleasedBudgetPerElement.remove(elementWithReleasableBudget);
+                assert val != null;
+            } finally {
+                lock.unlock();
+            }
+        }
+
+        private boolean isReleased(ElementWithReleasableBudget elementWithReleasableBudget) {
+            return unreleasedBudgetPerElement.containsKey(elementWithReleasableBudget) == false;
+        }
+
         class ElementWithReleasableBudget implements Releasable {
-            private final Wrap<E> wrappedElement;
-
-            private ElementWithReleasableBudget(E element, long budget) {
-                // Wrap the element in a brand-new instance that's used as the key in the
-                // {@link PriorityBlockingQueueWithBudget#unreleasedBudgetPerElement} identity map.
-                // This allows the same exact "element" instance to hold budgets multiple times concurrently.
-                // This way we allow to re-enqueue and re-take an element before a previous take completed and
-                // released the budget.
-                this.wrappedElement = new Wrap<>(element);
-                assert PriorityBlockingQueueWithBudget.this.lock.isHeldByCurrentThread();
-                // the taken element holds up some budget
-                var prev = PriorityBlockingQueueWithBudget.this.unreleasedBudgetPerElement.put(wrappedElement, budget);
-                assert prev == null;
-                assert isClosed() == false;
-                PriorityBlockingQueueWithBudget.this.availableBudget -= budget;
-                assert PriorityBlockingQueueWithBudget.this.availableBudget >= 0L;
+            private final E element;
+
+            private ElementWithReleasableBudget(E element) {
+                this.element = element;
             }
 
             /**
@@ -673,27 +690,17 @@ private ElementWithReleasableBudget(E element, long budget) {
              */
             @Override
             public void close() {
-                final ReentrantLock lock = PriorityBlockingQueueWithBudget.this.lock;
-                lock.lock();
-                try {
-                    assert isClosed() == false;
-                    // when the taken element is not used anymore, it will not influence subsequent available budget computations
-                    unreleasedBudgetPerElement.remove(wrappedElement);
-                } finally {
-                    lock.unlock();
-                }
+                PriorityBlockingQueueWithBudget.this.release(this);
             }
 
             boolean isClosed() {
-                return unreleasedBudgetPerElement.containsKey(wrappedElement) == false;
+                return PriorityBlockingQueueWithBudget.this.isReleased(this);
             }
 
             E element() {
-                return wrappedElement.element();
+                return element;
             }
         }
-
-        private record Wrap<E>(E element) {}
     }
 
     private static long newTargetIORateBytesPerSec(

server/src/test/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorServiceDiskSpaceTests.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 
@@ -36,9 +37,7 @@
 import java.nio.file.attribute.FileAttributeView;
 import java.nio.file.attribute.FileStoreAttributeView;
 import java.nio.file.spi.FileSystemProvider;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Deque;
 import java.util.IdentityHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -104,6 +103,11 @@ public static void removeMockUsableSpaceFS() {
         nodeEnvironment.close();
     }
 
+    @After
+    public void cleanupThreadPool() {
+        testThreadPool.scheduledTasks.clear();
+    }
+
     static class CapturingThreadPool extends TestThreadPool {
         final List<Tuple<TimeValue, Cancellable>> scheduledTasks = new ArrayList<>();
 
@@ -295,14 +299,23 @@ public void testDiskSpaceMonitorStartsAsDisabled() throws Exception {
             assertThat(availableDiskSpaceUpdates.size(), is(1));
             assertThat(availableDiskSpaceUpdates.getLast().getBytes(), is(Long.MAX_VALUE));
             // updating monitoring interval should enable the monitor
-            String intervalSettingValue = randomFrom("1s", "123ms", "5ns", "2h");
+            String intervalSettingValue = randomFrom("1s", "123ms", "5nanos", "2h");
             clusterSettings.applySettings(
                 Settings.builder()
                     .put(ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_CHECK_INTERVAL_SETTING.getKey(), intervalSettingValue)
                     .build()
             );
             assertThat(diskSpacePeriodicMonitor.isScheduled(), is(true));
             assertThat(testThreadPool.scheduledTasks.size(), is(1));
+            assertThat(
+                testThreadPool.scheduledTasks.getLast().v1(),
+                is(
+                    TimeValue.parseTimeValue(
+                        intervalSettingValue,
+                        ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_CHECK_INTERVAL_SETTING.getKey()
+                    )
+                )
+            );
         }
         aFileStore.throwIoException = false;
         bFileStore.throwIoException = false;