Changes

Author: Tim-Brooks

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

@@ -97,7 +97,7 @@ public static class Handler implements Releasable {
 
         private final ArrayList<Releasable> releasables = new ArrayList<>(4);
         private final ArrayList<BulkResponse> responses = new ArrayList<>(2);
-        private final IndexingPressure.Coordinating coordinatingOperation;
+        private final IndexingPressure.Incremental incrementalOperation;
         private boolean closed = false;
         private boolean globalFailure = false;
         private boolean incrementalRequestSubmitted = false;
@@ -116,10 +116,14 @@ protected Handler(
             this.waitForActiveShards = waitForActiveShards != null ? ActiveShardCount.parseString(waitForActiveShards) : null;
             this.timeout = timeout;
             this.refresh = refresh;
-            this.coordinatingOperation = indexingPressure.markCoordinatingOperationStarted(0, 0, false);
+            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;
@@ -129,7 +133,8 @@ public void addItems(List<DocWriteRequest<?>> items, Releasable releasable, Runn
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
-                    if (coordinatingOperation.shouldSplit()) {
+                    if (incrementalOperation.shouldSplit()) {
+                        IndexingPressure.Coordinating coordinating = incrementalOperation.split();
                         final boolean isFirstRequest = incrementalRequestSubmitted == false;
                         incrementalRequestSubmitted = true;
                         final ArrayList<Releasable> toRelease = new ArrayList<>(releasables);
@@ -151,7 +156,7 @@ public void onFailure(Exception e) {
                             }
                         }, () -> {
                             bulkInProgress = false;
-                            coordinatingOperation.releaseCurrent();
+                            coordinating.close();
                             toRelease.forEach(Releasable::close);
                             nextItems.run();
                         }));
@@ -172,6 +177,7 @@ public void lastItems(List<DocWriteRequest<?>> items, Releasable releasable, Act
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
+                    IndexingPressure.Coordinating 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.
@@ -192,7 +198,7 @@ public void onFailure(Exception e) {
                             errorResponse(listener);
                         }
                     }, () -> {
-                        coordinatingOperation.releaseCurrent();
+                        coordinating.close();
                         toRelease.forEach(Releasable::close);
                     }));
                 } else {
@@ -205,15 +211,15 @@ public void onFailure(Exception e) {
         public void close() {
             if (closed == false) {
                 closed = true;
-                coordinatingOperation.close();
+                incrementalOperation.close();
                 releasables.forEach(Releasable::close);
                 releasables.clear();
             }
         }
 
         private void shortCircuitDueToTopLevelFailure(List<DocWriteRequest<?>> items, Releasable releasable) {
             assert releasables.isEmpty();
-            assert coordinatingOperation.currentSize() == 0;
+            assert incrementalOperation.currentOperationsSize() == 0;
             assert bulkRequest == null;
             if (globalFailure == false) {
                 addItemLevelFailures(items);
@@ -258,11 +264,11 @@ private boolean internalAddItems(List<DocWriteRequest<?>> items, Releasable rele
                 bulkRequest.add(items);
                 releasables.add(releasable);
                 long size = items.stream().mapToLong(Accountable::ramBytesUsed).sum();
-                coordinatingOperation.increment(items.size(), size);
+                incrementalOperation.increment(items.size(), size);
                 return true;
             } catch (EsRejectedExecutionException e) {
                 handleBulkFailure(incrementalRequestSubmitted == false, e);
-                coordinatingOperation.releaseCurrent();
+                incrementalOperation.split().close();
                 releasables.forEach(Releasable::close);
                 releasables.clear();
                 return false;

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

@@ -138,24 +138,111 @@ private static Releasable wrapReleasable(Releasable releasable) {
         };
     }
 
+    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 boolean shouldSplit() {
+            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 true;
+            }
+            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 true;
+            }
+            return false;
+        }
+
+        public Coordinating 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 currentSize = 0;
+        private long currentOperationsSize = 0;
 
         public Coordinating(boolean forceExecution) {
             this.forceExecution = forceExecution;
         }
 
-        public void increment(int operations, long bytes) {
+        private void increment(int operations, long bytes) {
             assert closed.get() == false;
             long combinedBytes = currentCombinedCoordinatingAndPrimaryBytes.addAndGet(bytes);
             long replicaWriteBytes = currentReplicaBytes.get();
@@ -186,7 +273,7 @@ public void increment(int operations, long bytes) {
                 );
             }
             currentOperations += operations;
-            currentSize += bytes;
+            currentOperationsSize += bytes;
             logger.trace(() -> Strings.format("adding [%d] coordinating operations and [%d] bytes", operations, bytes));
             currentCoordinatingBytes.getAndAdd(bytes);
             currentCoordinatingOps.getAndAdd(operations);
@@ -196,43 +283,17 @@ public void increment(int operations, long bytes) {
             totalCoordinatingRequests.getAndIncrement();
         }
 
-        public long currentSize() {
-            return currentSize;
-        }
-
-        public void releaseCurrent() {
-            currentCombinedCoordinatingAndPrimaryBytes.getAndAdd(-currentSize);
-            currentCoordinatingBytes.getAndAdd(-currentSize);
-            currentCoordinatingOps.getAndAdd(-currentOperations);
-            currentSize = 0;
-            currentOperations = 0;
-        }
-
-        public boolean shouldSplit() {
-            long currentUsage = (currentCombinedCoordinatingAndPrimaryBytes.get() + currentReplicaBytes.get());
-            if (currentUsage >= highWatermark && currentSize >= highWatermarkSize) {
-                highWaterMarkSplits.getAndIncrement();
-                logger.trace(
-                    () -> Strings.format("Split bulk due to high watermark: current bytes [%d] and size [%d]", currentUsage, currentSize)
-                );
-                return true;
-            }
-            if (currentUsage >= lowWatermark && currentSize >= lowWatermarkSize) {
-                lowWaterMarkSplits.getAndIncrement();
-                logger.trace(
-                    () -> Strings.format("Split bulk due to low watermark: current bytes [%d] and size [%d]", currentUsage, currentSize)
-                );
-                return true;
-            }
-            return false;
-
-        }
-
         @Override
         public void close() {
             if (closed.compareAndSet(false, true)) {
-                logger.trace(() -> Strings.format("removing [%d] coordinating operations and [%d] bytes", currentOperations, currentSize));
-                releaseCurrent();
+                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";

server/src/main/java/org/elasticsearch/rest/action/document/RestBulkAction.java

@@ -155,7 +155,6 @@ static class ChunkHandler implements BaseRestHandler.RequestBodyChunkConsumer {
 
         private volatile RestChannel restChannel;
         private boolean shortCircuited;
-        private int bytesParsed = 0;
         private final ArrayDeque<ReleasableBytesReference> unParsedChunks = new ArrayDeque<>(4);
         private final ArrayList<DocWriteRequest<?>> items = new ArrayList<>(4);
 
@@ -202,6 +201,7 @@ public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boo
                 bytesConsumed = 0;
             } else {
                 try {
+                    handler.getIncrementalOperation().incrementUnparsedBytes(chunk.length());
                     unParsedChunks.add(chunk);
 
                     if (unParsedChunks.size() > 1) {
@@ -210,10 +210,8 @@ public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boo
                         data = chunk;
                     }
 
-                    // TODO: Check that the behavior here vs. globalRouting, globalPipeline, globalRequireAlias, globalRequireDatsStream in
-                    // BulkRequest#add is fine
                     bytesConsumed = parser.parse(data, isLast);
-                    bytesParsed += bytesConsumed;
+                    handler.getIncrementalOperation().transferUnparsedBytesToParsed(bytesConsumed);
 
                 } catch (Exception e) {
                     shortCircuit();
@@ -225,7 +223,7 @@ public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boo
             final ArrayList<Releasable> releasables = accountParsing(bytesConsumed);
             if (isLast) {
                 assert unParsedChunks.isEmpty();
-                if (bytesParsed == 0) {
+                if (handler.getIncrementalOperation().totalParsedBytes() == 0) {
                     shortCircuit();
                     new RestToXContentListener<>(channel).onFailure(new ElasticsearchParseException("request body is required"));
                 } else {

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

@@ -41,22 +41,18 @@ public void testHighAndLowWatermarkSplits() {
         IndexingPressure indexingPressure = new IndexingPressure(settings);
 
         try (
-            IndexingPressure.Coordinating coordinating1 = indexingPressure.markCoordinatingOperationStarted(
-                10,
-                randomIntBetween(1, 127),
-                false
-            );
-            IndexingPressure.Coordinating coordinating2 = indexingPressure.markCoordinatingOperationStarted(
+            IndexingPressure.Incremental coordinating1 = indexingPressure.startIncrementalCoordinating(10, randomIntBetween(1, 127), false);
+            IndexingPressure.Incremental coordinating2 = indexingPressure.startIncrementalCoordinating(
                 10,
                 randomIntBetween(128, 1023),
                 false
             );
-            IndexingPressure.Coordinating coordinating3 = indexingPressure.markCoordinatingOperationStarted(
+            IndexingPressure.Incremental coordinating3 = indexingPressure.startIncrementalCoordinating(
                 10,
                 randomIntBetween(1024, 6000),
                 false
             );
-            Releasable ignored1 = indexingPressure.markCoordinatingOperationStarted(
+            Releasable ignored1 = indexingPressure.startIncrementalCoordinating(
                 10,
                 1 + (8 * 1024) - indexingPressure.stats().getCurrentCoordinatingBytes(),
                 false

server/src/test/java/org/elasticsearch/rest/action/document/RestBulkActionTests.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.http.HttpBody;
 import org.elasticsearch.index.IndexVersion;
@@ -42,11 +43,7 @@
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
-import static org.mockito.ArgumentMatchers.anyBoolean;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 /**
  * Tests for {@link RestBulkAction}.
@@ -228,10 +225,7 @@ public void next() {
             .build();
         FakeRestChannel channel = new FakeRestChannel(request, randomBoolean(), 1);
 
-        IndexingPressure indexingPressure = mock(IndexingPressure.class);
-        when(indexingPressure.markCoordinatingOperationStarted(anyInt(), anyLong(), anyBoolean())).thenReturn(
-            mock(IndexingPressure.Coordinating.class)
-        );
+        IndexingPressure indexingPressure = new IndexingPressure(Settings.EMPTY);
         RestBulkAction.ChunkHandler chunkHandler = new RestBulkAction.ChunkHandler(true, request, () -> {
             return new IncrementalBulkService.Handler(null, indexingPressure, null, null, null) {