Make BulkItemRequest immutable

I was talking with @itschrispeck about some JIT optimization issues in
BulkItemRequest's serialization. While looking at the code, the
`volatile` keyword on the `primaryResponse` field made me cringe. Why
is a `BulkItemRequest` mutable at all?

It turns out that we modify the existing `BulkItemRequest` instances
on the primary shard. These modified requests are send to the replicas.

This change makes `BulkItemRequest` immutable. The primary execution
context collects all of the primary responses, then produces a new
`BulkShardRequest` that gets forwarded to replicas.

Signed-off-by: Michael Froh <msfroh@apache.org>

Author: msfroh

server/src/main/java/org/opensearch/action/bulk/BulkItemRequest.java

@@ -36,15 +36,12 @@
 import org.apache.lucene.util.RamUsageEstimator;
 import org.opensearch.action.DocWriteRequest;
 import org.opensearch.common.Nullable;
-import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.index.shard.ShardId;
-import org.opensearch.core.xcontent.MediaTypeRegistry;
 
 import java.io.IOException;
-import java.util.Objects;
 
 /**
  * Transport request for a Single bulk item
@@ -55,9 +52,9 @@ public class BulkItemRequest implements Writeable, Accountable {
 
     private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(BulkItemRequest.class);
 
-    private int id;
-    private DocWriteRequest<?> request;
-    private volatile BulkItemResponse primaryResponse;
+    private final int id;
+    private final DocWriteRequest<?> request;
+    private final BulkItemResponse primaryResponse;
 
     /**
      * @param shardId the shard id
@@ -71,13 +68,20 @@ public class BulkItemRequest implements Writeable, Accountable {
             } else {
                 primaryResponse = new BulkItemResponse(shardId, in);
             }
+        } else {
+            primaryResponse = null;
         }
     }
 
     // NOTE: public for testing only
     public BulkItemRequest(int id, DocWriteRequest<?> request) {
+        this(id, request, null);
+    }
+
+    BulkItemRequest(int id, DocWriteRequest<?> request, BulkItemResponse primaryResponse) {
         this.id = id;
         this.request = request;
+        this.primaryResponse = primaryResponse;
     }
 
     public int id() {
@@ -97,38 +101,6 @@ BulkItemResponse getPrimaryResponse() {
         return primaryResponse;
     }
 
-    void setPrimaryResponse(BulkItemResponse primaryResponse) {
-        this.primaryResponse = primaryResponse;
-    }
-
-    /**
-     * Abort this request, and store a {@link org.opensearch.action.bulk.BulkItemResponse.Failure} response.
-     *
-     * @param index The concrete index that was resolved for this request
-     * @param cause The cause of the rejection (may not be null)
-     * @throws IllegalStateException If a response already exists for this request
-     */
-    public void abort(String index, Exception cause) {
-        if (primaryResponse == null) {
-            final BulkItemResponse.Failure failure = new BulkItemResponse.Failure(index, request.id(), Objects.requireNonNull(cause), true);
-            setPrimaryResponse(new BulkItemResponse(id, request.opType(), failure));
-        } else {
-            assert primaryResponse.isFailed() && primaryResponse.getFailure().isAborted() : "response ["
-                + Strings.toString(MediaTypeRegistry.JSON, primaryResponse)
-                + "]; cause ["
-                + cause
-                + "]";
-            if (primaryResponse.isFailed() && primaryResponse.getFailure().isAborted()) {
-                primaryResponse.getFailure().getCause().addSuppressed(cause);
-            } else {
-                throw new IllegalStateException(
-                    "aborting item that with response [" + primaryResponse + "] that was previously processed",
-                    cause
-                );
-            }
-        }
-    }
-
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVInt(id);

server/src/main/java/org/opensearch/action/bulk/BulkItemResponse.java

@@ -421,7 +421,6 @@ public long getTerm() {
         /**
          * Whether this failure is the result of an <em>abort</em>.
          * If {@code true}, the request to which this failure relates should never be retried, regardless of the {@link #getCause() cause}.
-         * @see BulkItemRequest#abort(String, Exception)
          */
         public boolean isAborted() {
             return aborted;

server/src/main/java/org/opensearch/action/bulk/BulkPrimaryExecutionContext.java

@@ -45,9 +45,6 @@
 
 import java.util.Arrays;
 
-import static org.opensearch.action.bulk.BulkShardResponse.DEFAULT_QUEUE_SIZE;
-import static org.opensearch.action.bulk.BulkShardResponse.DEFAULT_SERVICE_TIME_IN_NANOS;
-
 /**
  * This is a utility class that holds the per request state needed to perform bulk operations on the primary.
  * More specifically, it maintains an index to the current executing bulk item, which allows execution
@@ -91,13 +88,15 @@ enum ItemProcessingState {
     private int currentIndex = -1;
 
     private ItemProcessingState currentItemState;
-    private DocWriteRequest requestToExecute;
+    private DocWriteRequest<?> requestToExecute;
     private BulkItemResponse executionResult;
     private int retryCounter;
+    private final BulkItemResponse[] primaryResponses;
 
     BulkPrimaryExecutionContext(BulkShardRequest request, IndexShard primary) {
         this.request = request;
         this.primary = primary;
+        this.primaryResponses = new BulkItemResponse[request.items().length];
         advance();
     }
 
@@ -131,7 +130,12 @@ public DocWriteRequest<?> getCurrent() {
     }
 
     public BulkShardRequest getBulkShardRequest() {
-        return request;
+        BulkItemRequest[] newRequests = new BulkItemRequest[request.items().length];
+        for (int i = 0; i < newRequests.length; i++) {
+            BulkItemRequest oldRequest = request.items()[i];
+            newRequests[i] = new BulkItemRequest(oldRequest.id(), oldRequest.request(), primaryResponses[i]);
+        }
+        return new BulkShardRequest(request.shardId(), request.getRefreshPolicy(), newRequests);
     }
 
     /** returns the result of the request that has been executed on the shard */
@@ -145,21 +149,11 @@ public int getRetryCounter() {
         return retryCounter;
     }
 
-    /** returns true if the current request has been executed on the primary */
-    public boolean isOperationExecuted() {
-        return currentItemState == ItemProcessingState.EXECUTED;
-    }
-
     /** returns true if the request needs to wait for a mapping update to arrive from the cluster-manager */
     public boolean requiresWaitingForMappingUpdate() {
         return currentItemState == ItemProcessingState.WAIT_FOR_MAPPING_UPDATE;
     }
 
-    /** returns true if the current request should be retried without waiting for an external event */
-    public boolean requiresImmediateRetry() {
-        return currentItemState == ItemProcessingState.IMMEDIATE_RETRY;
-    }
-
     /**
      * returns true if the current request has been completed and it's result translated to a user
      * facing response
@@ -206,17 +200,18 @@ public IndexShard getPrimary() {
     }
 
     /**
-     * sets the request that should actually be executed on the primary. This can be different then the request
+     * sets the request that should actually be executed on the primary. This can be different from the request
      * received from the user (specifically, an update request is translated to an indexing or delete request).
      */
-    public void setRequestToExecute(DocWriteRequest writeRequest) {
+    public void setRequestToExecute(DocWriteRequest<?> writeRequest) {
         assert assertInvariants(ItemProcessingState.INITIAL);
         requestToExecute = writeRequest;
         currentItemState = ItemProcessingState.TRANSLATED;
         assert assertInvariants(ItemProcessingState.TRANSLATED);
     }
 
     /** returns the request that should be executed on the shard. */
+    @SuppressWarnings("unchecked")
     public <T extends DocWriteRequest<T>> T getRequestToExecute() {
         assert assertInvariants(ItemProcessingState.TRANSLATED);
         return (T) requestToExecute;
@@ -252,7 +247,7 @@ public void markOperationAsNoOp(DocWriteResponse response) {
     public void failOnMappingUpdate(Exception cause) {
         assert assertInvariants(ItemProcessingState.WAIT_FOR_MAPPING_UPDATE);
         currentItemState = ItemProcessingState.EXECUTED;
-        final DocWriteRequest docWriteRequest = getCurrentItem().request();
+        final DocWriteRequest<?> docWriteRequest = getCurrentItem().request();
         executionResult = new BulkItemResponse(
             getCurrentItem().id(),
             docWriteRequest.opType(),
@@ -267,7 +262,7 @@ public void failOnMappingUpdate(Exception cause) {
     public void markOperationAsExecuted(Engine.Result result) {
         assertInvariants(ItemProcessingState.TRANSLATED);
         final BulkItemRequest current = getCurrentItem();
-        DocWriteRequest docWriteRequest = getRequestToExecute();
+        DocWriteRequest<?> docWriteRequest = getRequestToExecute();
         switch (result.getResultType()) {
             case SUCCESS:
                 final DocWriteResponse response;
@@ -347,24 +342,15 @@ public void markAsCompleted(BulkItemResponse translatedResponse) {
         if (translatedResponse.isFailed() == false && requestToExecute != null && requestToExecute != getCurrent()) {
             request.items()[currentIndex] = new BulkItemRequest(request.items()[currentIndex].id(), requestToExecute);
         }
-        getCurrentItem().setPrimaryResponse(translatedResponse);
+        primaryResponses[currentIndex] = translatedResponse;
         currentItemState = ItemProcessingState.COMPLETED;
         advance();
     }
 
     /** builds the bulk shard response to return to the user */
     public BulkShardResponse buildShardResponse(long serviceTimeEWMAInNanos, int nodeQueueSize) {
         assert hasMoreOperationsToExecute() == false;
-        return new BulkShardResponse(
-            request.shardId(),
-            Arrays.stream(request.items()).map(BulkItemRequest::getPrimaryResponse).toArray(BulkItemResponse[]::new),
-            serviceTimeEWMAInNanos,
-            nodeQueueSize
-        );
-    }
-
-    public BulkShardResponse buildShardResponse() {
-        return buildShardResponse(DEFAULT_SERVICE_TIME_IN_NANOS, DEFAULT_QUEUE_SIZE);
+        return new BulkShardResponse(request.shardId(), primaryResponses, serviceTimeEWMAInNanos, nodeQueueSize);
     }
 
     private boolean assertInvariants(ItemProcessingState... expectedCurrentState) {
@@ -379,18 +365,14 @@ private boolean assertInvariants(ItemProcessingState... expectedCurrentState) {
                 assert requestToExecute == null : requestToExecute;
                 assert executionResult == null : executionResult;
                 break;
-            case TRANSLATED:
+            case TRANSLATED, IMMEDIATE_RETRY:
                 assert requestToExecute != null;
                 assert executionResult == null : executionResult;
                 break;
             case WAIT_FOR_MAPPING_UPDATE:
                 assert requestToExecute == null;
                 assert executionResult == null : executionResult;
                 break;
-            case IMMEDIATE_RETRY:
-                assert requestToExecute != null;
-                assert executionResult == null : executionResult;
-                break;
             case EXECUTED:
                 // requestToExecute can be null if the update ended up as NOOP
                 assert executionResult != null;

server/src/test/java/org/opensearch/action/bulk/BulkPrimaryExecutionContextTests.java

@@ -40,75 +40,18 @@
 import org.opensearch.action.index.IndexRequest;
 import org.opensearch.action.support.WriteRequest;
 import org.opensearch.action.update.UpdateRequest;
-import org.opensearch.core.index.AppendOnlyIndexOperationRetryException;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.engine.Engine;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.translog.Translog;
 import org.opensearch.test.OpenSearchTestCase;
 
-import java.util.ArrayList;
-
 import static org.hamcrest.Matchers.equalTo;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class BulkPrimaryExecutionContextTests extends OpenSearchTestCase {
 
-    public void testAbortedSkipped() {
-        BulkShardRequest shardRequest = generateRandomRequest();
-
-        ArrayList<DocWriteRequest<?>> nonAbortedRequests = new ArrayList<>();
-        for (BulkItemRequest request : shardRequest.items()) {
-            if (randomBoolean()) {
-                request.abort("index", new OpenSearchException("bla"));
-            } else {
-                nonAbortedRequests.add(request.request());
-            }
-        }
-
-        ArrayList<DocWriteRequest<?>> visitedRequests = new ArrayList<>();
-        for (BulkPrimaryExecutionContext context = new BulkPrimaryExecutionContext(shardRequest, null); context
-            .hasMoreOperationsToExecute();) {
-            visitedRequests.add(context.getCurrent());
-            context.setRequestToExecute(context.getCurrent());
-            // using failures prevents caring about types
-            context.markOperationAsExecuted(new Engine.IndexResult(new OpenSearchException("bla"), 1));
-            context.markAsCompleted(context.getExecutionResult());
-        }
-
-        assertThat(visitedRequests, equalTo(nonAbortedRequests));
-    }
-
-    public void testAppendOnlyIndexOperationRetryException() {
-        BulkShardRequest shardRequest = generateRandomRequest();
-
-        final IndexShard primary = mock(IndexShard.class);
-        when(primary.shardId()).thenReturn(shardRequest.shardId());
-        ArrayList<DocWriteRequest<?>> nonAbortedRequests = new ArrayList<>();
-        for (BulkItemRequest request : shardRequest.items()) {
-            if (randomBoolean()) {
-                request.abort("index", new AppendOnlyIndexOperationRetryException("Indexing operation retried for append only indices"));
-            } else {
-                nonAbortedRequests.add(request.request());
-            }
-        }
-
-        ArrayList<DocWriteRequest<?>> visitedRequests = new ArrayList<>();
-        for (BulkPrimaryExecutionContext context = new BulkPrimaryExecutionContext(shardRequest, primary); context
-            .hasMoreOperationsToExecute();) {
-            visitedRequests.add(context.getCurrent());
-            context.setRequestToExecute(context.getCurrent());
-            // using failures prevents caring about types
-            context.markOperationAsExecuted(
-                new Engine.IndexResult(new AppendOnlyIndexOperationRetryException("Indexing operation retried for append only indices"), 1)
-            );
-            context.markAsCompleted(context.getExecutionResult());
-        }
-
-        assertThat(visitedRequests, equalTo(nonAbortedRequests));
-    }
-
     private BulkShardRequest generateRandomRequest() {
         BulkItemRequest[] items = new BulkItemRequest[randomInt(20)];
         for (int i = 0; i < items.length; i++) {