SNAPSHOT - Bulk Ops

Author: lukewhiting

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

@@ -43,6 +43,8 @@
 import org.elasticsearch.cluster.routing.IndexRouting;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.common.streams.StreamType;
+import org.elasticsearch.common.streams.StreamsPermissionsUtils;
 import org.elasticsearch.common.util.concurrent.AtomicArray;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
@@ -59,6 +61,8 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -70,6 +74,7 @@
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.LongSupplier;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.action.bulk.TransportBulkAction.LAZY_ROLLOVER_ORIGIN;
 import static org.elasticsearch.cluster.metadata.IndexNameExpressionResolver.EXCLUDED_DATA_STREAMS_KEY;
@@ -104,6 +109,7 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
     private final FailureStoreMetrics failureStoreMetrics;
     private final DataStreamFailureStoreSettings dataStreamFailureStoreSettings;
     private final boolean clusterHasFailureStoreFeature;
+    private final StreamsPermissionsUtils streamsPermissionsUtils;
 
     BulkOperation(
         Task task,
@@ -139,7 +145,8 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
             new FailureStoreDocumentConverter(),
             failureStoreMetrics,
             dataStreamFailureStoreSettings,
-            clusterHasFailureStoreFeature
+            clusterHasFailureStoreFeature,
+            StreamsPermissionsUtils.getInstance()
         );
     }
 
@@ -160,7 +167,8 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
         FailureStoreDocumentConverter failureStoreDocumentConverter,
         FailureStoreMetrics failureStoreMetrics,
         DataStreamFailureStoreSettings dataStreamFailureStoreSettings,
-        boolean clusterHasFailureStoreFeature
+        boolean clusterHasFailureStoreFeature,
+        StreamsPermissionsUtils streamsPermissionsUtils
     ) {
         super(listener);
         this.task = task;
@@ -182,6 +190,7 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
         this.failureStoreMetrics = failureStoreMetrics;
         this.dataStreamFailureStoreSettings = dataStreamFailureStoreSettings;
         this.clusterHasFailureStoreFeature = clusterHasFailureStoreFeature;
+        this.streamsPermissionsUtils = streamsPermissionsUtils;
     }
 
     @Override
@@ -274,6 +283,30 @@ private long buildTookInMillis(long startTimeNanos) {
     }
 
     private Map<ShardId, List<BulkItemRequest>> groupBulkRequestsByShards(ClusterState clusterState) {
+        ProjectMetadata projectMetadata = projectResolver.getProjectMetadata(clusterState);
+
+        Set<StreamType> enabledStreamTypes = Arrays.stream(StreamType.values())
+            .filter(t -> streamsPermissionsUtils.streamTypeIsEnabled(t, projectMetadata))
+            .collect(Collectors.toCollection(() -> EnumSet.noneOf(StreamType.class)));
+
+        for (StreamType streamType : enabledStreamTypes) {
+            for (int i = 0; i < bulkRequest.requests.size(); i++) {
+                DocWriteRequest<?> req = bulkRequest.requests.get(i);
+                String prefix = streamType.getStreamName() + ".";
+                if (req != null && req.index().startsWith(prefix)) {
+                    IllegalArgumentException exception = new IllegalArgumentException(
+                        "Bulk requests for streams with type ["
+                            + streamType.getStreamName()
+                            + "] are not supported, use the ["
+                            + streamType.getStreamName()
+                            + "] API instead."
+                    );
+                    IndexDocFailureStoreStatus failureStoreStatus = processFailure(new BulkItemRequest(i, req), projectMetadata, exception);
+                    addFailureAndDiscardRequest(req, i, req.index(), exception, failureStoreStatus);
+                }
+            }
+        }
+
         return groupRequestsByShards(
             clusterState,
             Iterators.enumerate(bulkRequest.requests.iterator(), BulkItemRequest::new),

server/src/main/java/org/elasticsearch/common/streams/StreamTypes.java renamed to server/src/main/java/org/elasticsearch/common/streams/StreamType.java

@@ -9,13 +9,13 @@
 
 package org.elasticsearch.common.streams;
 
-public enum StreamTypes {
+public enum StreamType {
 
     LOGS("logs");
 
     private final String streamName;
 
-    StreamTypes(String streamName) {
+    StreamType(String streamName) {
         this.streamName = streamName;
     }
 

server/src/main/java/org/elasticsearch/common/streams/StreamsPermissionsUtils.java

@@ -34,7 +34,7 @@ public static StreamsPermissionsUtils getInstance() {
 
     public void throwIfRetrouteToSubstreamNotAllowed(ProjectMetadata projectMetadata, Set<String> indexHistory, String destination)
         throws IllegalArgumentException {
-        for (StreamTypes streamType : StreamTypes.values()) {
+        for (StreamType streamType : StreamType.values()) {
             String streamName = streamType.getStreamName();
             if (streamTypeIsEnabled(streamType, projectMetadata)
                 && destination.startsWith(streamName + ".")
@@ -51,7 +51,7 @@ public void throwIfRetrouteToSubstreamNotAllowed(ProjectMetadata projectMetadata
         }
     }
 
-    public boolean streamTypeIsEnabled(StreamTypes streamType, ProjectMetadata projectMetadata) {
+    public boolean streamTypeIsEnabled(StreamType streamType, ProjectMetadata projectMetadata) {
         StreamsMetadata metadata = projectMetadata.custom(StreamsMetadata.TYPE, StreamsMetadata.EMPTY);
         return switch (streamType) {
             case LOGS -> metadata.isLogsEnabled();

server/src/test/java/org/elasticsearch/action/bulk/BulkOperationTests.java

@@ -42,10 +42,12 @@
 import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.cluster.metadata.Template;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.project.TestProjectResolvers;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.streams.StreamsPermissionsUtils;
 import org.elasticsearch.common.util.concurrent.AtomicArray;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
@@ -215,10 +217,18 @@ public class BulkOperationTests extends ESTestCase {
 
     private TestThreadPool threadPool;
 
+    private StreamsPermissionsUtils streamsPermissionsUtilsMock;
+    private ProjectResolver projectResolverMock;
+    private IndexNameExpressionResolver indexNameExpressionResolverMock;
+
     @Before
-    public void setupThreadpool() {
+    public void setupTest() {
         threadPool = new TestThreadPool(getClass().getName());
         threadPool.getThreadContext().putHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER, projectId.id());
+        streamsPermissionsUtilsMock = mock(StreamsPermissionsUtils.class);
+        when(streamsPermissionsUtilsMock.streamTypeIsEnabled(any(), any())).thenReturn(false);
+        projectResolverMock = mock(ProjectResolver.class);
+        indexNameExpressionResolverMock = mock(IndexNameExpressionResolver.class);
     }
 
     @After
@@ -688,8 +698,8 @@ public void testRetryableBlockAcceptsFailureStoreDocument() throws Exception {
     }
 
     /**
-     * A bulk operation to a data stream with a failure store enabled may still partially fail if the cluster is experiencing a
-     * non-retryable block when the redirected documents would be sent to the shard-level action.
+     * A bulk operation to a data stream with a failure store enabled may still partially fail if the redirected documents experience
+     * a shard-level failure while writing to the failure store indices.
      */
     public void testBlockedClusterRejectsFailureStoreDocument() throws Exception {
         // Requests that go to two separate shards
@@ -1239,7 +1249,8 @@ private BulkOperation newBulkOperation(
             failureStoreDocumentConverter,
             FailureStoreMetrics.NOOP,
             dataStreamFailureStoreSettings,
-            failureStoreNodeFeatureEnabled
+            failureStoreNodeFeatureEnabled,
+            streamsPermissionsUtilsMock
         );
     }
 

server/src/test/java/org/elasticsearch/common/streams/StreamsPermissionsUtilsTests.java

@@ -47,22 +47,22 @@ public void testGetInstanceReturnsSingleton() {
     public void testStreamTypeIsEnabledReturnsTrueWhenLogsEnabled() {
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(true);
 
-        boolean result = utils.streamTypeIsEnabled(StreamTypes.LOGS, projectMetadataMock);
+        boolean result = utils.streamTypeIsEnabled(StreamType.LOGS, projectMetadataMock);
         assertTrue(result);
     }
 
     public void testStreamTypeIsEnabledReturnsFalseWhenLogsDisabled() {
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(false);
 
-        boolean result = utils.streamTypeIsEnabled(StreamTypes.LOGS, projectMetadataMock);
+        boolean result = utils.streamTypeIsEnabled(StreamType.LOGS, projectMetadataMock);
         assertFalse(result);
     }
 
     public void testIfRetrouteToSubstreamNotAllowedThrows() {
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(true);
 
         Set<String> indexHistory = new HashSet<>(); // empty, so reroute not allowed
-        String destination = StreamTypes.LOGS.getStreamName() + ".substream";
+        String destination = StreamType.LOGS.getStreamName() + ".substream";
 
         IllegalArgumentException ex = expectThrows(
             IllegalArgumentException.class,
@@ -77,7 +77,7 @@ public void testThrowIfRetrouteToSubstreamNotAllowedDoesNotThrowWhenStreamTypeDi
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(false);
 
         Set<String> indexHistory = Collections.emptySet();
-        String destination = StreamTypes.LOGS.getStreamName() + ".substream";
+        String destination = StreamType.LOGS.getStreamName() + ".substream";
 
         // Should not throw since stream type is disabled
         utils.throwIfRetrouteToSubstreamNotAllowed(projectMetadataMock, indexHistory, destination);
@@ -87,7 +87,7 @@ public void testThrowIfRetrouteToSubstreamNotAllowedDoesNotThrowWhenDestinationN
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(true);
 
         Set<String> indexHistory = Collections.emptySet();
-        String destination = StreamTypes.LOGS.getStreamName(); // not a substream
+        String destination = StreamType.LOGS.getStreamName(); // not a substream
 
         // Should not throw since destination is not a substream
         utils.throwIfRetrouteToSubstreamNotAllowed(projectMetadataMock, indexHistory, destination);
@@ -97,8 +97,8 @@ public void testThrowIfRetrouteToSubstreamNotAllowedDoesNotThrowWhenIndexHistory
         when(streamsMetadataMock.isLogsEnabled()).thenReturn(true);
 
         Set<String> indexHistory = new HashSet<>();
-        indexHistory.add(StreamTypes.LOGS.getStreamName());
-        String destination = StreamTypes.LOGS.getStreamName() + ".substream";
+        indexHistory.add(StreamType.LOGS.getStreamName());
+        String destination = StreamType.LOGS.getStreamName() + ".substream";
 
         // Should not throw since indexHistory contains the stream name
         utils.throwIfRetrouteToSubstreamNotAllowed(projectMetadataMock, indexHistory, destination);