Fix currentSize calculation in BulkIngester

CHANGELOG.md

@@ -30,6 +30,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Fixed
 - Fix `unitTest` task not running the tests in the `test` source set ([#2074](https://github.com/opensearch-project/opensearch-java/pull/2074))
 - Run model tests against both JSON mappers instead of picking one at random ([#2085](https://github.com/opensearch-project/opensearch-java/pull/2085))
+- Fix currentSize calculation in BulkIngester ([#2113](https://github.com/opensearch-project/opensearch-java/pull/2113))
 
 ## [Unreleased 3.x]
 ### Added

java-client/src/main/java/org/opensearch/client/opensearch/_helpers/bulk/BulkIngester.java

@@ -125,7 +125,7 @@ public class BulkIngester<Context> implements AutoCloseable {
     private BackoffPolicy backoffPolicy;
 
     // Current state
-    private List<RetryableBulkOperation<Context>> operations = new ArrayList<>();
+    private List<IngesterOperation<Context>> operations = new ArrayList<>();
     private long currentSize;
     private int requestsInFlightCount;
     private volatile boolean isClosed = false;
@@ -237,7 +237,7 @@ public Duration flushInterval() {
      * The number of operations that have been buffered, waiting to be sent.
      */
     public int pendingOperations() {
-        List<RetryableBulkOperation<Context>> operations = this.operations;
+        List<IngesterOperation<Context>> operations = this.operations;
         return operations == null ? 0 : operations.size();
     }
 
@@ -353,32 +353,31 @@ private void failsafeFlush() {
      * automatic flush triggers (maxOperations, maxSize, or flushInterval).
      */
     public void flush() {
-        List<RetryableBulkOperation<Context>> sentRequests = new ArrayList<>();
+        List<IngesterOperation<Context>> sentRequests = new ArrayList<>();
         RequestExecution<Context> exec = sendRequestCondition.whenReadyIf(() -> {
             // May happen on manual and periodic flushes
-            return !operations.isEmpty() && operations.stream().anyMatch(RetryableBulkOperation::isSendable);
+            return !operations.isEmpty() && operations.stream().anyMatch(IngesterOperation::isSendable);
         }, () -> {
             // Selecting operations that can be sent immediately,
             // Dividing actual operations from contexts
             List<BulkOperation> immediateOps = new ArrayList<>();
             List<Context> contexts = new ArrayList<>();
 
-            for (Iterator<RetryableBulkOperation<Context>> it = operations.iterator(); it.hasNext();) {
-                RetryableBulkOperation<Context> op = it.next();
+            for (Iterator<IngesterOperation<Context>> it = operations.iterator(); it.hasNext();) {
+                IngesterOperation<Context> op = it.next();
                 if (op.isSendable()) {
                     immediateOps.add(op.operation());
                     contexts.add(op.context());
 
                     sentRequests.add(op);
+                    currentSize -= op.size();
                     it.remove();
                 }
             }
 
             // Build the request
             BulkRequest request = newRequest().operations(immediateOps).build();
 
-            // Prepare for next round
-            currentSize = operations.size();
             addCondition.signalIfReady();
 
             long id = sendRequestCondition.invocations();
@@ -423,7 +422,7 @@ public void flush() {
                         // Partial success, retrying failed requests if policy allows it
                         // Keeping list of retryable requests/responses, to exclude them for calling
                         // listener later
-                        List<RetryableBulkOperation<Context>> retryableReq = new ArrayList<>();
+                        List<IngesterOperation<Context>> retryableReq = new ArrayList<>();
                         List<RetryableBulkOperation<Context>> refires = new ArrayList<>();
                         List<BulkResponseItem> retryableResp = new ArrayList<>();
 
@@ -442,7 +441,7 @@ public void flush() {
                                 // Creating partial BulkRequest
                                 List<BulkOperation> partialOps = new ArrayList<>();
                                 List<Context> partialCtx = new ArrayList<>();
-                                for (RetryableBulkOperation<Context> op : sentRequests) {
+                                for (IngesterOperation<Context> op : sentRequests) {
                                     partialOps.add(op.operation());
                                     partialCtx.add(op.context());
                                 }
@@ -489,14 +488,13 @@ public void flush() {
     private void selectingRetries(
         int index,
         BulkResponseItem bulkItemResponse,
-        List<RetryableBulkOperation<Context>> sentRequests,
+        List<IngesterOperation<Context>> sentRequests,
         List<BulkResponseItem> retryableResp,
-        List<RetryableBulkOperation<Context>> retryableReq,
+        List<IngesterOperation<Context>> retryableReq,
         List<RetryableBulkOperation<Context>> refires
     ) {
-
         // Getting original failed, requests and keeping successful ones to send to the listener
-        RetryableBulkOperation<Context> original = sentRequests.get(index);
+        IngesterOperation<Context> original = sentRequests.get(index);
         if (original.canRetry()) {
             retryableResp.add(bulkItemResponse);
             Iterator<Long> retryTimes = Optional.ofNullable(original.retries()).orElse(backoffPolicy.iterator());
@@ -592,10 +590,10 @@ private void addRetry(RetryableBulkOperation<Context> repeatableOp) {
     }
 
     private void innerAdd(RetryableBulkOperation<Context> repeatableOp) {
-        IngesterOperation ingestOp = IngesterOperation.of(repeatableOp, client._transport().jsonpMapper());
+        IngesterOperation<Context> ingestOp = IngesterOperation.of(repeatableOp, client._transport().jsonpMapper());
 
         addCondition.whenReady(() -> {
-            operations.add(ingestOp.repeatableOperation());
+            operations.add(ingestOp);
             currentSize += ingestOp.size();
 
             if (!canAddOperation()) {

java-client/src/main/java/org/opensearch/client/opensearch/_helpers/bulk/IngesterOperation.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.client.opensearch._helpers.bulk;
 
+import java.util.Iterator;
 import javax.annotation.Nullable;
 import org.opensearch.client.json.JsonEnum;
 import org.opensearch.client.json.JsonpMapper;
@@ -53,11 +54,11 @@
  * <p>
  * This is an internal utility class used by {@link BulkIngester} to track buffered operation sizes.
  */
-class IngesterOperation {
-    private final RetryableBulkOperation repeatableOp;
+class IngesterOperation<Context> {
+    private final RetryableBulkOperation<Context> repeatableOp;
     private final long size;
 
-    IngesterOperation(RetryableBulkOperation repeatableOp, long size) {
+    IngesterOperation(RetryableBulkOperation<Context> repeatableOp, long size) {
         this.repeatableOp = repeatableOp;
         this.size = size;
     }
@@ -69,7 +70,7 @@ class IngesterOperation {
      * @param mapper       the JSON mapper for serialization
      * @return an IngesterOperation with calculated size
      */
-    public static IngesterOperation of(RetryableBulkOperation repeatableOp, JsonpMapper mapper) {
+    public static <Context> IngesterOperation<Context> of(RetryableBulkOperation<Context> repeatableOp, JsonpMapper mapper) {
         switch (repeatableOp.operation()._kind()) {
             case Create:
                 return createOperation(repeatableOp, mapper);
@@ -89,7 +90,7 @@ public static IngesterOperation of(RetryableBulkOperation repeatableOp, JsonpMap
      *
      * @return the retryable bulk operation
      */
-    public RetryableBulkOperation repeatableOperation() {
+    public RetryableBulkOperation<Context> repeatableOperation() {
         return this.repeatableOp;
     }
 
@@ -102,9 +103,29 @@ public long size() {
         return this.size;
     }
 
-    private static IngesterOperation createOperation(RetryableBulkOperation repeatableOp, JsonpMapper mapper) {
+    public BulkOperation operation() {
+        return repeatableOp.operation();
+    }
+
+    public Context context() {
+        return repeatableOp.context();
+    }
+
+    public boolean isSendable() {
+        return repeatableOp.isSendable();
+    }
+
+    public boolean canRetry() {
+        return repeatableOp.canRetry();
+    }
+
+    public Iterator<Long> retries() {
+        return repeatableOp.retries();
+    }
+
+    private static <Context> IngesterOperation<Context> createOperation(RetryableBulkOperation<Context> repeatableOp, JsonpMapper mapper) {
         CreateOperation<?> create = repeatableOp.operation().create();
-        RetryableBulkOperation newOperation;
+        RetryableBulkOperation<Context> newOperation;
 
         long size = basePropertiesSize(create);
 
@@ -115,18 +136,18 @@ private static IngesterOperation createOperation(RetryableBulkOperation repeatab
         } else {
             BinaryData binaryDoc = BinaryData.of(create.document(), mapper);
             size += binaryDoc.size();
-            newOperation = new RetryableBulkOperation(BulkOperation.of(bo -> bo.create(idx -> {
+            newOperation = new RetryableBulkOperation<>(BulkOperation.of(bo -> bo.create(idx -> {
                 copyCreateProperties(create, idx);
                 return idx.document(binaryDoc);
             })), repeatableOp.context(), repeatableOp.retries());
         }
 
-        return new IngesterOperation(newOperation, size);
+        return new IngesterOperation<>(newOperation, size);
     }
 
-    private static IngesterOperation indexOperation(RetryableBulkOperation repeatableOp, JsonpMapper mapper) {
+    private static <Context> IngesterOperation<Context> indexOperation(RetryableBulkOperation<Context> repeatableOp, JsonpMapper mapper) {
         IndexOperation<?> index = repeatableOp.operation().index();
-        RetryableBulkOperation newOperation;
+        RetryableBulkOperation<Context> newOperation;
 
         long size = basePropertiesSize(index);
 
@@ -137,16 +158,16 @@ private static IngesterOperation indexOperation(RetryableBulkOperation repeatabl
         } else {
             BinaryData binaryDoc = BinaryData.of(index.document(), mapper);
             size += binaryDoc.size();
-            newOperation = new RetryableBulkOperation(BulkOperation.of(bo -> bo.index(idx -> {
+            newOperation = new RetryableBulkOperation<>(BulkOperation.of(bo -> bo.index(idx -> {
                 copyIndexProperties(index, idx);
                 return idx.document(binaryDoc);
             })), repeatableOp.context(), repeatableOp.retries());
         }
 
-        return new IngesterOperation(newOperation, size);
+        return new IngesterOperation<>(newOperation, size);
     }
 
-    private static IngesterOperation updateOperation(RetryableBulkOperation repeatableOp, JsonpMapper mapper) {
+    private static <Context> IngesterOperation<Context> updateOperation(RetryableBulkOperation<Context> repeatableOp, JsonpMapper mapper) {
         UpdateOperation<?> update = repeatableOp.operation().update();
 
         // UpdateOperation implements NdJsonpSerializable, which means it serializes as two separate JSON objects:
@@ -185,12 +206,12 @@ private static IngesterOperation updateOperation(RetryableBulkOperation repeatab
             ) + 300; // Fallback estimate for data
         }
 
-        return new IngesterOperation(repeatableOp, size);
+        return new IngesterOperation<>(repeatableOp, size);
     }
 
-    private static IngesterOperation deleteOperation(RetryableBulkOperation repeatableOp) {
+    private static <Context> IngesterOperation<Context> deleteOperation(RetryableBulkOperation<Context> repeatableOp) {
         DeleteOperation delete = repeatableOp.operation().delete();
-        return new IngesterOperation(repeatableOp, basePropertiesSize(delete));
+        return new IngesterOperation<>(repeatableOp, basePropertiesSize(delete));
     }
 
     private static void copyBaseProperties(BulkOperationBase op, BulkOperationBase.AbstractBuilder<?> builder) {