Make BulkItemRequest immutable

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

@@ -36,99 +36,50 @@
 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
  *
  * @opensearch.internal
  */
-public class BulkItemRequest implements Writeable, Accountable {
+public record BulkItemRequest(int id, DocWriteRequest<?> request, BulkItemResponse primaryResponse) implements Writeable, Accountable {
 
     private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(BulkItemRequest.class);
 
-    private int id;
-    private DocWriteRequest<?> request;
-    private volatile BulkItemResponse primaryResponse;
-
     /**
      * @param shardId the shard id
      */
     BulkItemRequest(@Nullable ShardId shardId, StreamInput in) throws IOException {
-        id = in.readVInt();
-        request = DocWriteRequest.readDocumentRequest(shardId, in);
+        this(in.readVInt(), DocWriteRequest.readDocumentRequest(shardId, in), readPrimaryResponse(shardId, in));
+    }
+
+    private static BulkItemResponse readPrimaryResponse(ShardId shardId, StreamInput in) throws IOException {
         if (in.readBoolean()) {
             if (shardId == null) {
-                primaryResponse = new BulkItemResponse(in);
+                return new BulkItemResponse(in);
             } else {
-                primaryResponse = new BulkItemResponse(shardId, in);
+                return new BulkItemResponse(shardId, in);
             }
         }
+        return null;
     }
 
     // NOTE: public for testing only
     public BulkItemRequest(int id, DocWriteRequest<?> request) {
-        this.id = id;
-        this.request = request;
-    }
-
-    public int id() {
-        return id;
-    }
-
-    public DocWriteRequest<?> request() {
-        return request;
+        this(id, request, null);
     }
 
     public String index() {
         assert request.indices().length == 1;
         return request.indices()[0];
     }
 
-    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);
@@ -139,8 +90,6 @@ public void writeTo(StreamOutput out) throws IOException {
     public void writeThin(StreamOutput out) throws IOException {
         out.writeVInt(id);
         DocWriteRequest.writeDocumentRequestThin(out, request);
-
-        BulkItemResponse primaryResponse = this.primaryResponse; // Read volatile once
         out.writeOptionalWriteable((o, resp) -> resp.writeThin(o), primaryResponse);
     }
 

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,19 +88,21 @@ 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();
     }
 
     private int findNextNonAborted(int startIndex) {
         final int length = request.items().length;
-        while (startIndex < length && isAborted(request.items()[startIndex].getPrimaryResponse())) {
+        while (startIndex < length && isAborted(request.items()[startIndex].primaryResponse())) {
             startIndex++;
         }
         return startIndex;
@@ -131,7 +130,7 @@ public DocWriteRequest<?> getCurrent() {
     }
 
     public BulkShardRequest getBulkShardRequest() {
-        return request;
+        return request.setPrimaryResponses(primaryResponses);
     }
 
     /** returns the result of the request that has been executed on the shard */
@@ -145,21 +144,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 +195,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 +242,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 +257,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 +337,16 @@ 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;
+        assertInvariants(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,26 +361,22 @@ 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;
                 break;
             case COMPLETED:
-                assert requestToExecute != null;
+                // requestToExecute can be null if the update ended up as NOOP
                 assert executionResult != null;
-                assert getCurrentItem().getPrimaryResponse() != null;
+                assert primaryResponses[currentIndex] != null;
                 break;
         }
         return true;

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

@@ -68,6 +68,24 @@ public BulkShardRequest(ShardId shardId, RefreshPolicy refreshPolicy, BulkItemRe
         setRefreshPolicy(refreshPolicy);
     }
 
+    BulkShardRequest setPrimaryResponses(BulkItemResponse[] primaryResponses) {
+        if (primaryResponses == null || primaryResponses.length != items.length) {
+            throw new IllegalArgumentException("Primary responses must have same length as BulkItemRequests");
+        }
+        BulkItemRequest[] newRequests = new BulkItemRequest[items.length];
+        for (int i = 0; i < items.length; i++) {
+            BulkItemRequest request = items[i];
+            if (request == null) {
+                newRequests[i] = null;
+            } else {
+                newRequests[i] = new BulkItemRequest(request.id(), request.request(), primaryResponses[i]);
+            }
+        }
+        BulkShardRequest bulkShardRequest = new BulkShardRequest(shardId, getRefreshPolicy(), newRequests);
+        cloneProperties(bulkShardRequest);
+        return bulkShardRequest;
+    }
+
     public BulkItemRequest[] items() {
         return items;
     }

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

@@ -860,9 +860,10 @@ public static Translog.Location performOnReplica(BulkShardRequest request, Index
         Translog.Location location = null;
         for (int i = 0; i < request.items().length; i++) {
             final BulkItemRequest item = request.items()[i];
-            final BulkItemResponse response = item.getPrimaryResponse();
+            final BulkItemResponse response = item.primaryResponse();
             final Engine.Result operationResult;
-            if (item.getPrimaryResponse().isFailed()) {
+            assert response != null;
+            if (response.isFailed()) {
                 if (response.getFailure().getSeqNo() == SequenceNumbers.UNASSIGNED_SEQ_NO) {
                     continue; // ignore replication as we didn't generate a sequence number for this request.
                 }
@@ -886,7 +887,7 @@ public static Translog.Location performOnReplica(BulkShardRequest request, Index
                 assert response.getResponse().getSeqNo() != SequenceNumbers.UNASSIGNED_SEQ_NO;
                 operationResult = performOpOnReplica(response.getResponse(), item.request(), replica);
             }
-            assert operationResult != null : "operation result must never be null when primary response has no failure";
+            assert operationResult != null : "operation result must never be null";
             location = syncOperationResultOrThrow(operationResult, location);
         }
         return location;

server/src/main/java/org/opensearch/action/support/replication/ReplicationRequest.java

@@ -262,4 +262,11 @@ public String getDescription() {
     public void onRetry() {
         // nothing by default
     }
+
+    protected void cloneProperties(ReplicationRequest<?> target) {
+        target.waitForActiveShards(waitForActiveShards());
+        target.timeout(timeout());
+        target.routedBasedOnClusterVersion(routedBasedOnClusterVersion());
+        target.setParentTask(getParentTask());
+    }
 }