Fix remotestore/restore API to honour wait_for_completion

Signed-off-by: Sachin Kale <[email protected]>
sachinpkale · Sep 28, 2023 · f594192 · f594192
1 parent d656e3d
commit f594192
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java
@@ -13,6 +13,7 @@
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.cluster.health.ClusterHealthStatus;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.test.InternalTestCluster;
 import org.opensearch.test.OpenSearchIntegTestCase;
@@ -30,6 +31,30 @@
 @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 0)
 public class RemoteStoreRestoreIT extends BaseRemoteStoreRestoreIT {
 
+    public void testWaitForCompletion() throws IOException, ExecutionException, InterruptedException {
+        prepareCluster(1, 3, INDEX_NAME, 0, 1);
+        Map<String, Long> indexStats = indexData(2, true, INDEX_NAME);
+
+        assertHitCount(client().prepareSearch(INDEX_NAME).setSize(0).get(), indexStats.get(REFRESHED_OR_FLUSHED_OPERATIONS));
+
+        internalCluster().stopRandomNode(InternalTestCluster.nameFilter(primaryNodeName(INDEX_NAME)));
+        ensureRed(INDEX_NAME);
+
+        PlainActionFuture<RestoreRemoteStoreResponse> listener = PlainActionFuture.newFuture();
+        client().admin()
+            .cluster()
+            .restoreRemoteStore(
+                new RestoreRemoteStoreRequest().indices(INDEX_NAME).waitForCompletion(true),
+                listener
+            );
+        RestoreRemoteStoreResponse restoreRemoteStoreResponse = listener.get();
+
+        assertEquals(1, restoreRemoteStoreResponse.getRestoreInfo().successfulShards());
+        assertEquals(0, restoreRemoteStoreResponse.getRestoreInfo().failedShards());
+        assertHitCount(client().prepareSearch(INDEX_NAME).setSize(0).get(), indexStats.get(TOTAL_OPERATIONS));
+        logger.info("--> Successful");
+    }
+
     /**
      * Simulates all data restored using Remote Translog Store.
      * @throws IOException IO Exception.

diff --git a/...pensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java b/...pensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java
@@ -85,7 +85,7 @@ protected void clusterManagerOperation(
         final ActionListener<RestoreRemoteStoreResponse> listener
     ) {
         restoreService.restore(request, ActionListener.delegateFailure(listener, (delegatedListener, restoreCompletionResponse) -> {
-            if (restoreCompletionResponse.getRestoreInfo() == null && request.waitForCompletion()) {
+            if (request.waitForCompletion()) {
                 RestoreClusterStateListener.createAndRegisterListener(
                     clusterService,
                     restoreCompletionResponse,