Add ingestion management APIs

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

Diff for: plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/KafkaIngestionBaseIT.java

@@ -14,11 +14,15 @@
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.serialization.StringSerializer;
+import org.opensearch.action.admin.indices.streamingingestion.pause.PauseIngestionResponse;
+import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionResponse;
+import org.opensearch.action.admin.indices.streamingingestion.state.GetIngestionStateResponse;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.transport.client.Requests;
 import org.junit.After;
 import org.junit.Before;
 
@@ -28,6 +32,7 @@
 import java.util.Locale;
 import java.util.Properties;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
 import org.testcontainers.containers.KafkaContainer;
@@ -112,6 +117,11 @@ protected void waitForSearchableDocs(long docCount, List<String> nodes) throws E
         }, 1, TimeUnit.MINUTES);
     }
 
+    protected long getSearchableDocCount(String node) throws Exception {
+        final SearchResponse response = client(node).prepareSearch(indexName).setSize(0).setPreference("_only_local").get();
+        return response.getHits().getTotalHits().value();
+    }
+
     protected void waitForState(Callable<Boolean> checkState) throws Exception {
         assertBusy(() -> {
             if (checkState.call() == false) {
@@ -124,6 +134,18 @@ protected String getSettings(String indexName, String setting) {
         return client().admin().indices().prepareGetSettings(indexName).get().getSetting(indexName, setting);
     }
 
+    protected GetIngestionStateResponse getIngestionState(String indexName) throws ExecutionException, InterruptedException {
+        return client().admin().indices().getIngestionState(Requests.getIngestionStateRequest(indexName)).get();
+    }
+
+    protected PauseIngestionResponse pauseIngestion(String indexName) throws ExecutionException, InterruptedException {
+        return client().admin().indices().pauseIngestion(Requests.pauseIngestionRequest(indexName)).get();
+    }
+
+    protected ResumeIngestionResponse resumeIngestion(String indexName) throws ExecutionException, InterruptedException {
+        return client().admin().indices().resumeIngestion(Requests.resumeIngestionRequest(indexName)).get();
+    }
+
     protected void createIndexWithDefaultSettings(int numShards, int numReplicas) {
         createIndex(
             indexName,

Diff for: plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/RemoteStoreKafkaIT.java

@@ -9,6 +9,9 @@
 package org.opensearch.plugin.kafka;
 
 import org.opensearch.action.admin.indices.settings.get.GetSettingsResponse;
+import org.opensearch.action.admin.indices.streamingingestion.pause.PauseIngestionResponse;
+import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionResponse;
+import org.opensearch.action.admin.indices.streamingingestion.state.GetIngestionStateResponse;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.routing.allocation.command.AllocateReplicaAllocationCommand;
@@ -154,6 +157,56 @@ public void testErrorStrategy() throws Exception {
         waitForSearchableDocs(2, Arrays.asList(node));
     }
 
+    public void testPauseAndResumeIngestion() throws Exception {
+        // setup nodes and index
+        produceData("1", "name1", "24");
+        produceData("2", "name2", "20");
+        internalCluster().startClusterManagerOnlyNode();
+        final String nodeA = internalCluster().startDataOnlyNode();
+        final String nodeB = internalCluster().startDataOnlyNode();
+
+        createIndexWithDefaultSettings(1, 1);
+        ensureGreen(indexName);
+        waitForSearchableDocs(2, Arrays.asList(nodeA, nodeB));
+
+        // pause ingestion
+        PauseIngestionResponse pauseResponse = pauseIngestion(indexName);
+        assertTrue(pauseResponse.isAcknowledged());
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.getPollerState().equalsIgnoreCase("paused"));
+        });
+
+        // verify ingestion state is persisted
+        produceData("3", "name3", "30");
+        produceData("4", "name4", "31");
+        internalCluster().stopRandomNode(InternalTestCluster.nameFilter(nodeA));
+        ensureYellowAndNoInitializingShards(indexName);
+        assertTrue(nodeB.equals(primaryNodeName(indexName)));
+
+        final String nodeC = internalCluster().startDataOnlyNode();
+        client().admin().cluster().prepareReroute().add(new AllocateReplicaAllocationCommand(indexName, 0, nodeC)).get();
+        ensureGreen(indexName);
+        assertTrue(nodeC.equals(replicaNodeName(indexName)));
+        assertEquals(2, getSearchableDocCount(nodeB));
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.getPollerState().equalsIgnoreCase("paused"));
+        });
+
+        // resume ingestion
+        ResumeIngestionResponse resumeResponse = resumeIngestion(indexName);
+        assertTrue(resumeResponse.isAcknowledged());
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return Arrays.stream(ingestionState.getShardStates())
+                .allMatch(
+                    state -> state.getPollerState().equalsIgnoreCase("polling") || state.getPollerState().equalsIgnoreCase("processing")
+                );
+        });
+        waitForSearchableDocs(4, Arrays.asList(nodeB, nodeC));
+    }
+
     private void verifyRemoteStoreEnabled(String node) {
         GetSettingsResponse settingsResponse = client(node).admin().indices().prepareGetSettings(indexName).get();
         String remoteStoreEnabled = settingsResponse.getIndexToSettings().get(indexName).get("index.remote_store.enabled");

Diff for: server/src/main/java/org/opensearch/action/ActionModule.java

@@ -198,6 +198,12 @@
 import org.opensearch.action.admin.indices.shrink.TransportResizeAction;
 import org.opensearch.action.admin.indices.stats.IndicesStatsAction;
 import org.opensearch.action.admin.indices.stats.TransportIndicesStatsAction;
+import org.opensearch.action.admin.indices.streamingingestion.pause.PauseIngestionAction;
+import org.opensearch.action.admin.indices.streamingingestion.pause.TransportPauseIngestionAction;
+import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionAction;
+import org.opensearch.action.admin.indices.streamingingestion.resume.TransportResumeIngestionAction;
+import org.opensearch.action.admin.indices.streamingingestion.state.GetIngestionStateAction;
+import org.opensearch.action.admin.indices.streamingingestion.state.TransportGetIngestionStateAction;
 import org.opensearch.action.admin.indices.template.delete.DeleteComponentTemplateAction;
 import org.opensearch.action.admin.indices.template.delete.DeleteComposableIndexTemplateAction;
 import org.opensearch.action.admin.indices.template.delete.DeleteIndexTemplateAction;
@@ -401,6 +407,7 @@
 import org.opensearch.rest.action.admin.indices.RestGetFieldMappingAction;
 import org.opensearch.rest.action.admin.indices.RestGetIndexTemplateAction;
 import org.opensearch.rest.action.admin.indices.RestGetIndicesAction;
+import org.opensearch.rest.action.admin.indices.RestGetIngestionStateAction;
 import org.opensearch.rest.action.admin.indices.RestGetMappingAction;
 import org.opensearch.rest.action.admin.indices.RestGetSettingsAction;
 import org.opensearch.rest.action.admin.indices.RestIndexDeleteAliasesAction;
@@ -410,6 +417,7 @@
 import org.opensearch.rest.action.admin.indices.RestIndicesShardStoresAction;
 import org.opensearch.rest.action.admin.indices.RestIndicesStatsAction;
 import org.opensearch.rest.action.admin.indices.RestOpenIndexAction;
+import org.opensearch.rest.action.admin.indices.RestPauseIngestionAction;
 import org.opensearch.rest.action.admin.indices.RestPutComponentTemplateAction;
 import org.opensearch.rest.action.admin.indices.RestPutComposableIndexTemplateAction;
 import org.opensearch.rest.action.admin.indices.RestPutIndexTemplateAction;
@@ -418,6 +426,7 @@
 import org.opensearch.rest.action.admin.indices.RestRefreshAction;
 import org.opensearch.rest.action.admin.indices.RestResizeHandler;
 import org.opensearch.rest.action.admin.indices.RestResolveIndexAction;
+import org.opensearch.rest.action.admin.indices.RestResumeIngestionAction;
 import org.opensearch.rest.action.admin.indices.RestRolloverIndexAction;
 import org.opensearch.rest.action.admin.indices.RestSimulateIndexTemplateAction;
 import org.opensearch.rest.action.admin.indices.RestSimulateTemplateAction;
@@ -806,6 +815,11 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(GetSearchPipelineAction.INSTANCE, GetSearchPipelineTransportAction.class);
         actions.register(DeleteSearchPipelineAction.INSTANCE, DeleteSearchPipelineTransportAction.class);
 
+        // Pull-based ingestion actions
+        actions.register(PauseIngestionAction.INSTANCE, TransportPauseIngestionAction.class);
+        actions.register(ResumeIngestionAction.INSTANCE, TransportResumeIngestionAction.class);
+        actions.register(GetIngestionStateAction.INSTANCE, TransportGetIngestionStateAction.class);
+
         return unmodifiableMap(actions.getRegistry());
     }
 
@@ -1041,6 +1055,11 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestGetDecommissionStateAction());
         registerHandler.accept(new RestRemoteStoreStatsAction());
         registerHandler.accept(new RestRestoreRemoteStoreAction());
+
+        // pull-based ingestion API
+        registerHandler.accept(new RestPauseIngestionAction());
+        registerHandler.accept(new RestResumeIngestionAction());
+        registerHandler.accept(new RestGetIngestionStateAction());
     }
 
     @Override

Diff for: server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/package-info.java

@@ -0,0 +1,10 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/** Streaming ingestion transport handlers. */
+package org.opensearch.action.admin.indices.streamingingestion;

Diff for: server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/pause/PauseIngestionAction.java

@@ -0,0 +1,26 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.pause;
+
+import org.opensearch.action.ActionType;
+
+/**
+ * Transport action for pausing ingestion.
+ *
+ * @opensearch.experimental
+ */
+public class PauseIngestionAction extends ActionType<PauseIngestionResponse> {
+
+    public static final PauseIngestionAction INSTANCE = new PauseIngestionAction();
+    public static final String NAME = "indices:admin/ingestion/pause";
+
+    private PauseIngestionAction() {
+        super(NAME, PauseIngestionResponse::new);
+    }
+}

Diff for: server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/pause/PauseIngestionClusterStateUpdateRequest.java

@@ -0,0 +1,34 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.pause;
+
+import org.opensearch.cluster.ack.IndicesClusterStateUpdateRequest;
+
+/**
+ * Cluster state update request that allows to pause ingestion.
+ *
+ * @opensearch.experimental
+ */
+public class PauseIngestionClusterStateUpdateRequest extends IndicesClusterStateUpdateRequest<PauseIngestionClusterStateUpdateRequest> {
+
+    private long taskId;
+
+    public PauseIngestionClusterStateUpdateRequest(final long taskId) {
+        this.taskId = taskId;
+    }
+
+    public long taskId() {
+        return taskId;
+    }
+
+    public PauseIngestionClusterStateUpdateRequest taskId(final long taskId) {
+        this.taskId = taskId;
+        return this;
+    }
+}

Diff for: server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/pause/PauseIngestionRequest.java

@@ -0,0 +1,100 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.pause;
+
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.action.IndicesRequest;
+import org.opensearch.action.support.IndicesOptions;
+import org.opensearch.action.support.clustermanager.AcknowledgedRequest;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.util.CollectionUtils;
+
+import java.io.IOException;
+
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+/**
+ * A request to pause ingestion.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class PauseIngestionRequest extends AcknowledgedRequest<PauseIngestionRequest> implements IndicesRequest.Replaceable {
+
+    private String[] indices;
+    private IndicesOptions indicesOptions = IndicesOptions.strictExpandOpen();
+
+    public PauseIngestionRequest(StreamInput in) throws IOException {
+        super(in);
+        this.indices = in.readStringArray();
+        this.indicesOptions = IndicesOptions.readIndicesOptions(in);
+    }
+
+    public PauseIngestionRequest(String[] indices) {
+        this.indices = indices;
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException validationException = null;
+        if (CollectionUtils.isEmpty(indices)) {
+            validationException = addValidationError("index is missing", validationException);
+        }
+        return validationException;
+    }
+
+    /**
+     * Returns the indices to be paused.
+     */
+    @Override
+    public String[] indices() {
+        return indices;
+    }
+
+    /**
+     * Sets the indices to be paused.
+     */
+    @Override
+    public PauseIngestionRequest indices(String... indices) {
+        this.indices = indices;
+        return this;
+    }
+
+    /**
+     * Specifies what type of requested indices to ignore and how to deal with wildcard expressions.
+     * For example indices that don't exist.
+     *
+     * @return the desired behaviour regarding indices to ignore and wildcard indices expressions
+     */
+    @Override
+    public IndicesOptions indicesOptions() {
+        return indicesOptions;
+    }
+
+    /**
+     * Specifies what type of requested indices to ignore and how to deal wild wildcard expressions.
+     * For example indices that don't exist.
+     *
+     * @param indicesOptions the desired behaviour regarding indices to ignore and wildcard indices expressions
+     * @return the request itself
+     */
+    public PauseIngestionRequest indicesOptions(IndicesOptions indicesOptions) {
+        this.indicesOptions = indicesOptions;
+        return this;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeStringArray(indices);
+        indicesOptions.writeIndicesOptions(out);
+    }
+}