Backport Add coordinating object to track bytes (elastic#122460) and Simplify check to split bulk request (elastic#124035)

Author: Mikep86

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;
 
@@ -91,19 +92,18 @@ public static class Handler implements Releasable {
         public static final BulkRequest.IncrementalState EMPTY_STATE = new BulkRequest.IncrementalState(Collections.emptyMap(), true);
 
         private final Client client;
-        private final IndexingPressure indexingPressure;
         private final ActiveShardCount waitForActiveShards;
         private final TimeValue timeout;
         private final String refresh;
 
         private final ArrayList<Releasable> releasables = new ArrayList<>(4);
         private final ArrayList<BulkResponse> responses = new ArrayList<>(2);
+        private final IndexingPressure.Incremental incrementalOperation;
         private boolean closed = false;
         private boolean globalFailure = false;
         private boolean incrementalRequestSubmitted = false;
         private boolean bulkInProgress = false;
         private Exception bulkActionLevelFailure = null;
-        private long currentBulkSize = 0L;
         private BulkRequest bulkRequest = null;
 
         protected Handler(
@@ -114,13 +114,17 @@ protected Handler(
             @Nullable String refresh
         ) {
             this.client = client;
-            this.indexingPressure = indexingPressure;
             this.waitForActiveShards = waitForActiveShards != null ? ActiveShardCount.parseString(waitForActiveShards) : null;
             this.timeout = timeout;
             this.refresh = refresh;
+            this.incrementalOperation = indexingPressure.startIncrementalCoordinating(0, 0, false);
             createNewBulkRequest(EMPTY_STATE);
         }
 
+        public IndexingPressure.Incremental getIncrementalOperation() {
+            return incrementalOperation;
+        }
+
         public void addItems(List<DocWriteRequest<?>> items, Releasable releasable, Runnable nextItems) {
             assert closed == false;
             assert bulkInProgress == false;
@@ -130,7 +134,9 @@ public void addItems(List<DocWriteRequest<?>> items, Releasable releasable, Runn
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
-                    if (shouldBackOff()) {
+                    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);
@@ -153,6 +159,7 @@ public void onFailure(Exception e) {
                         }, () -> {
                             bulkInProgress = false;
                             toRelease.forEach(Releasable::close);
+                            coordinating.close();
                             nextItems.run();
                         }));
                     } else {
@@ -164,10 +171,6 @@ public void onFailure(Exception e) {
             }
         }
 
-        private boolean shouldBackOff() {
-            return indexingPressure.shouldSplitBulk(currentBulkSize);
-        }
-
         public void lastItems(List<DocWriteRequest<?>> items, Releasable releasable, ActionListener<BulkResponse> listener) {
             assert bulkInProgress == false;
             if (bulkActionLevelFailure != null) {
@@ -176,6 +179,7 @@ public void lastItems(List<DocWriteRequest<?>> items, Releasable releasable, Act
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
+                    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.
@@ -195,7 +199,10 @@ public void onFailure(Exception e) {
                             handleBulkFailure(isFirstRequest, e);
                             errorResponse(listener);
                         }
-                    }, () -> toRelease.forEach(Releasable::close)));
+                    }, () -> {
+                        toRelease.forEach(Releasable::close);
+                        coordinating.close();
+                    }));
                 } else {
                     errorResponse(listener);
                 }
@@ -204,13 +211,17 @@ public void onFailure(Exception e) {
 
         @Override
         public void close() {
-            closed = true;
-            releasables.forEach(Releasable::close);
-            releasables.clear();
+            if (closed == false) {
+                closed = true;
+                incrementalOperation.close();
+                releasables.forEach(Releasable::close);
+                releasables.clear();
+            }
         }
 
         private void shortCircuitDueToTopLevelFailure(List<DocWriteRequest<?>> items, Releasable releasable) {
             assert releasables.isEmpty();
+            assert incrementalOperation.currentOperationsSize() == 0;
             assert bulkRequest == null;
             if (globalFailure == false) {
                 addItemLevelFailures(items);
@@ -228,7 +239,6 @@ private void errorResponse(ActionListener<BulkResponse> listener) {
 
         private void handleBulkSuccess(BulkResponse bulkResponse) {
             responses.add(bulkResponse);
-            currentBulkSize = 0L;
             bulkRequest = null;
         }
 
@@ -237,7 +247,6 @@ private void handleBulkFailure(boolean isFirstRequest, Exception e) {
             globalFailure = isFirstRequest;
             bulkActionLevelFailure = e;
             addItemLevelFailures(bulkRequest.requests());
-            currentBulkSize = 0;
             bulkRequest = null;
         }
 
@@ -257,19 +266,18 @@ private boolean internalAddItems(List<DocWriteRequest<?>> items, Releasable rele
                 bulkRequest.add(items);
                 releasables.add(releasable);
                 long size = items.stream().mapToLong(Accountable::ramBytesUsed).sum();
-                releasables.add(indexingPressure.markCoordinatingOperationStarted(items.size(), size, false));
-                currentBulkSize += size;
+                incrementalOperation.increment(items.size(), size);
                 return true;
             } catch (EsRejectedExecutionException e) {
                 handleBulkFailure(incrementalRequestSubmitted == false, e);
+                incrementalOperation.split().close();
                 releasables.forEach(Releasable::close);
                 releasables.clear();
                 return false;
             }
         }
 
         private void createNewBulkRequest(BulkRequest.IncrementalState incrementalState) {
-            assert currentBulkSize == 0L;
             assert bulkRequest == null;
             bulkRequest = new BulkRequest();
             bulkRequest.incrementalState(incrementalState);

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;
 
@@ -138,48 +139,167 @@ private static Releasable wrapReleasable(Releasable releasable) {
         };
     }
 
-    public Releasable markCoordinatingOperationStarted(int operations, long bytes, boolean forceExecution) {
-        long combinedBytes = this.currentCombinedCoordinatingAndPrimaryBytes.addAndGet(bytes);
-        long replicaWriteBytes = this.currentReplicaBytes.get();
-        long totalBytes = combinedBytes + replicaWriteBytes;
-        if (forceExecution == false && totalBytes > coordinatingLimit) {
-            long bytesWithoutOperation = combinedBytes - bytes;
-            long totalBytesWithoutOperation = totalBytes - bytes;
-            this.currentCombinedCoordinatingAndPrimaryBytes.getAndAdd(-bytes);
-            this.coordinatingRejections.getAndIncrement();
-            throw new EsRejectedExecutionException(
-                "rejected execution of coordinating operation ["
-                    + "coordinating_and_primary_bytes="
-                    + bytesWithoutOperation
-                    + ", "
-                    + "replica_bytes="
-                    + replicaWriteBytes
-                    + ", "
-                    + "all_bytes="
-                    + totalBytesWithoutOperation
-                    + ", "
-                    + "coordinating_operation_bytes="
-                    + bytes
-                    + ", "
-                    + "max_coordinating_bytes="
-                    + coordinatingLimit
-                    + "]",
-                false
-            );
+    public Incremental startIncrementalCoordinating(int operations, long bytes, boolean forceExecution) {
+        Incremental coordinating = new Incremental(forceExecution);
+        coordinating.coordinating.increment(operations, bytes);
+        return coordinating;
+    }
+
+    public Coordinating markCoordinatingOperationStarted(int operations, long bytes, boolean forceExecution) {
+        Coordinating coordinating = new Coordinating(forceExecution);
+        coordinating.increment(operations, bytes);
+        return coordinating;
+    }
+
+    public class Incremental implements Releasable {
+
+        private final AtomicBoolean closed = new AtomicBoolean();
+        private final boolean forceExecution;
+        private long currentUnparsedSize = 0;
+        private long totalParsedBytes = 0;
+        private Coordinating coordinating;
+
+        public Incremental(boolean forceExecution) {
+            this.forceExecution = forceExecution;
+            this.coordinating = new Coordinating(forceExecution);
+        }
+
+        public long totalParsedBytes() {
+            return totalParsedBytes;
+        }
+
+        public void incrementUnparsedBytes(long bytes) {
+            assert closed.get() == false;
+            // TODO: Implement integration with IndexingPressure for unparsed bytes
+            currentUnparsedSize += bytes;
+        }
+
+        public void transferUnparsedBytesToParsed(long bytes) {
+            assert closed.get() == false;
+            assert currentUnparsedSize >= bytes;
+            currentUnparsedSize -= bytes;
+            totalParsedBytes += bytes;
+        }
+
+        public void increment(int operations, long bytes) {
+            // TODO: Eventually most of the memory will already be accounted for in unparsed.
+            coordinating.increment(operations, bytes);
+        }
+
+        public long currentOperationsSize() {
+            return coordinating.currentOperationsSize;
+        }
+
+        public Optional<Releasable> maybeSplit() {
+            long currentUsage = (currentCombinedCoordinatingAndPrimaryBytes.get() + currentReplicaBytes.get());
+            long currentOperationsSize = coordinating.currentOperationsSize;
+            if (currentUsage >= highWatermark && currentOperationsSize >= highWatermarkSize) {
+                highWaterMarkSplits.getAndIncrement();
+                logger.trace(
+                    () -> Strings.format(
+                        "Split bulk due to high watermark: current bytes [%d] and size [%d]",
+                        currentUsage,
+                        currentOperationsSize
+                    )
+                );
+                return Optional.of(split());
+            }
+            if (currentUsage >= lowWatermark && currentOperationsSize >= lowWatermarkSize) {
+                lowWaterMarkSplits.getAndIncrement();
+                logger.trace(
+                    () -> Strings.format(
+                        "Split bulk due to low watermark: current bytes [%d] and size [%d]",
+                        currentUsage,
+                        currentOperationsSize
+                    )
+                );
+                return Optional.of(split());
+            }
+            return Optional.empty();
+        }
+
+        public Releasable split() {
+            Coordinating toReturn = coordinating;
+            coordinating = new Coordinating(forceExecution);
+            return toReturn;
+        }
+
+        @Override
+        public void close() {
+            coordinating.close();
+        }
+    }
+
+    // TODO: Maybe this should be re-named and used for primary operations too. Eventually we will need to account for: ingest pipeline
+    // expansions, reading updates, etc. This could just be a generic OP that could be expanded as appropriate
+    public class Coordinating implements Releasable {
+
+        private final AtomicBoolean closed = new AtomicBoolean();
+        private final boolean forceExecution;
+        private int currentOperations = 0;
+        private long currentOperationsSize = 0;
+
+        public Coordinating(boolean forceExecution) {
+            this.forceExecution = forceExecution;
+        }
+
+        private void increment(int operations, long bytes) {
+            assert closed.get() == false;
+            long combinedBytes = currentCombinedCoordinatingAndPrimaryBytes.addAndGet(bytes);
+            long replicaWriteBytes = currentReplicaBytes.get();
+            long totalBytes = combinedBytes + replicaWriteBytes;
+            if (forceExecution == false && totalBytes > coordinatingLimit) {
+                long bytesWithoutOperation = combinedBytes - bytes;
+                long totalBytesWithoutOperation = totalBytes - bytes;
+                currentCombinedCoordinatingAndPrimaryBytes.getAndAdd(-bytes);
+                coordinatingRejections.getAndIncrement();
+                throw new EsRejectedExecutionException(
+                    "rejected execution of coordinating operation ["
+                        + "coordinating_and_primary_bytes="
+                        + bytesWithoutOperation
+                        + ", "
+                        + "replica_bytes="
+                        + replicaWriteBytes
+                        + ", "
+                        + "all_bytes="
+                        + totalBytesWithoutOperation
+                        + ", "
+                        + "coordinating_operation_bytes="
+                        + bytes
+                        + ", "
+                        + "max_coordinating_bytes="
+                        + coordinatingLimit
+                        + "]",
+                    false
+                );
+            }
+            currentOperations += operations;
+            currentOperationsSize += bytes;
+            logger.trace(() -> Strings.format("adding [%d] coordinating operations and [%d] bytes", operations, bytes));
+            currentCoordinatingBytes.getAndAdd(bytes);
+            currentCoordinatingOps.getAndAdd(operations);
+            totalCombinedCoordinatingAndPrimaryBytes.getAndAdd(bytes);
+            totalCoordinatingBytes.getAndAdd(bytes);
+            totalCoordinatingOps.getAndAdd(operations);
+            totalCoordinatingRequests.getAndIncrement();
+        }
+
+        @Override
+        public void close() {
+            if (closed.compareAndSet(false, true)) {
+                logger.trace(
+                    () -> Strings.format("removing [%d] coordinating operations and [%d] bytes", currentOperations, currentOperationsSize)
+                );
+                currentCombinedCoordinatingAndPrimaryBytes.getAndAdd(-currentOperationsSize);
+                currentCoordinatingBytes.getAndAdd(-currentOperationsSize);
+                currentCoordinatingOps.getAndAdd(-currentOperations);
+                currentOperationsSize = 0;
+                currentOperations = 0;
+            } else {
+                logger.error("IndexingPressure memory is adjusted twice", new IllegalStateException("Releasable is called twice"));
+                assert false : "IndexingPressure is adjusted twice";
+            }
         }
-        logger.trace(() -> Strings.format("adding [%d] coordinating operations and [%d] bytes", operations, bytes));
-        currentCoordinatingBytes.getAndAdd(bytes);
-        currentCoordinatingOps.getAndAdd(operations);
-        totalCombinedCoordinatingAndPrimaryBytes.getAndAdd(bytes);
-        totalCoordinatingBytes.getAndAdd(bytes);
-        totalCoordinatingOps.getAndAdd(operations);
-        totalCoordinatingRequests.getAndIncrement();
-        return wrapReleasable(() -> {
-            logger.trace(() -> Strings.format("removing [%d] coordinating operations and [%d] bytes", operations, bytes));
-            this.currentCombinedCoordinatingAndPrimaryBytes.getAndAdd(-bytes);
-            this.currentCoordinatingBytes.getAndAdd(-bytes);
-            this.currentCoordinatingOps.getAndAdd(-operations);
-        });
     }
 
     public Releasable markPrimaryOperationLocalToCoordinatingNodeStarted(int operations, long bytes) {
@@ -266,21 +386,6 @@ public Releasable markReplicaOperationStarted(int operations, long bytes, boolea
         });
     }
 
-    public boolean shouldSplitBulk(long size) {
-        long currentUsage = (currentCombinedCoordinatingAndPrimaryBytes.get() + currentReplicaBytes.get());
-        if (currentUsage >= highWatermark && size >= highWatermarkSize) {
-            highWaterMarkSplits.getAndIncrement();
-            logger.trace(() -> Strings.format("Split bulk due to high watermark: current bytes [%d] and size [%d]", currentUsage, size));
-            return (true);
-        }
-        if (currentUsage >= lowWatermark && size >= lowWatermarkSize) {
-            lowWaterMarkSplits.getAndIncrement();
-            logger.trace(() -> Strings.format("Split bulk due to low watermark: current bytes [%d] and size [%d]", currentUsage, size));
-            return (true);
-        }
-        return (false);
-    }
-
     public IndexingPressureStats stats() {
         return new IndexingPressureStats(
             totalCombinedCoordinatingAndPrimaryBytes.get(),