Simplify check to split bulk request (#124035)

Use an optional to check if bulk operation should be split.

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/action/bulk/IncrementalBulkService.java

@@ -26,6 +26,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 
@@ -133,8 +134,9 @@ public void addItems(List<DocWriteRequest<?>> items, Releasable releasable, Runn
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
-                    if (incrementalOperation.shouldSplit()) {
-                        IndexingPressure.Coordinating coordinating = incrementalOperation.split();
+                    Optional<Releasable> maybeSplit = incrementalOperation.maybeSplit();
+                    if (maybeSplit.isPresent()) {
+                        Releasable coordinating = maybeSplit.get();
                         final boolean isFirstRequest = incrementalRequestSubmitted == false;
                         incrementalRequestSubmitted = true;
                         final ArrayList<Releasable> toRelease = new ArrayList<>(releasables);
@@ -156,8 +158,8 @@ public void onFailure(Exception e) {
                             }
                         }, () -> {
                             bulkInProgress = false;
-                            coordinating.close();
                             toRelease.forEach(Releasable::close);
+                            coordinating.close();
                             nextItems.run();
                         }));
                     } else {
@@ -177,7 +179,7 @@ public void lastItems(List<DocWriteRequest<?>> items, Releasable releasable, Act
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
-                    IndexingPressure.Coordinating coordinating = incrementalOperation.split();
+                    Releasable coordinating = incrementalOperation.split();
                     final ArrayList<Releasable> toRelease = new ArrayList<>(releasables);
                     releasables.clear();
                     // We do not need to set this back to false as this will be the last request.
@@ -198,8 +200,8 @@ public void onFailure(Exception e) {
                             errorResponse(listener);
                         }
                     }, () -> {
-                        coordinating.close();
                         toRelease.forEach(Releasable::close);
+                        coordinating.close();
                     }));
                 } else {
                     errorResponse(listener);

server/src/main/java/org/elasticsearch/index/IndexingPressure.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.index.stats.IndexingPressureStats;
 
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -189,7 +190,7 @@ public long currentOperationsSize() {
             return coordinating.currentOperationsSize;
         }
 
-        public boolean shouldSplit() {
+        public Optional<Releasable> maybeSplit() {
             long currentUsage = (currentCombinedCoordinatingAndPrimaryBytes.get() + currentReplicaBytes.get());
             long currentOperationsSize = coordinating.currentOperationsSize;
             if (currentUsage >= highWatermark && currentOperationsSize >= highWatermarkSize) {
@@ -201,7 +202,7 @@ public boolean shouldSplit() {
                         currentOperationsSize
                     )
                 );
-                return true;
+                return Optional.of(split());
             }
             if (currentUsage >= lowWatermark && currentOperationsSize >= lowWatermarkSize) {
                 lowWaterMarkSplits.getAndIncrement();
@@ -212,12 +213,12 @@ public boolean shouldSplit() {
                         currentOperationsSize
                     )
                 );
-                return true;
+                return Optional.of(split());
             }
-            return false;
+            return Optional.empty();
         }
 
-        public Coordinating split() {
+        public Releasable split() {
             Coordinating toReturn = coordinating;
             coordinating = new Coordinating(forceExecution);
             return toReturn;

server/src/test/java/org/elasticsearch/index/IndexingPressureTests.java

@@ -17,6 +17,8 @@
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matchers;
 
+import java.util.Optional;
+
 public class IndexingPressureTests extends ESTestCase {
 
     private final Settings settings = Settings.builder()
@@ -58,28 +60,31 @@ public void testHighAndLowWatermarkSplits() {
                 false
             )
         ) {
-            assertFalse(coordinating1.shouldSplit());
-            assertFalse(coordinating2.shouldSplit());
+            assertFalse(coordinating1.maybeSplit().isPresent());
+            assertFalse(coordinating2.maybeSplit().isPresent());
             assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 0L);
             assertEquals(indexingPressure.stats().getLowWaterMarkSplits(), 0L);
-            assertTrue(coordinating3.shouldSplit());
-            assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 0L);
-            assertEquals(indexingPressure.stats().getLowWaterMarkSplits(), 1L);
-
-            try (
-                Releasable ignored2 = indexingPressure.markCoordinatingOperationStarted(
-                    10,
-                    1 + (9 * 1024) - indexingPressure.stats().getCurrentCoordinatingBytes(),
-                    false
-                )
-            ) {
-                assertFalse(coordinating1.shouldSplit());
-                assertTrue(coordinating2.shouldSplit());
-                assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 1L);
+            Optional<Releasable> split1 = coordinating3.maybeSplit();
+            assertTrue(split1.isPresent());
+            try (Releasable ignored2 = split1.get()) {
+                assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 0L);
                 assertEquals(indexingPressure.stats().getLowWaterMarkSplits(), 1L);
-                assertTrue(coordinating3.shouldSplit());
-                assertEquals(indexingPressure.stats().getLowWaterMarkSplits(), 1L);
-                assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 2L);
+
+                try (
+                    Releasable ignored3 = indexingPressure.markCoordinatingOperationStarted(
+                        10,
+                        1 + (9 * 1024) - indexingPressure.stats().getCurrentCoordinatingBytes(),
+                        false
+                    )
+                ) {
+                    assertFalse(coordinating1.maybeSplit().isPresent());
+                    Optional<Releasable> split2 = coordinating2.maybeSplit();
+                    assertTrue(split2.isPresent());
+                    try (Releasable ignored4 = split2.get()) {
+                        assertEquals(indexingPressure.stats().getHighWaterMarkSplits(), 1L);
+                        assertEquals(indexingPressure.stats().getLowWaterMarkSplits(), 1L);
+                    }
+                }
             }
         }
     }