Add wait for restore

Signed-off-by: Lakshya Taragi <[email protected]>
ltaragi · Sep 6, 2024 · 926067a · 926067a
1 parent 3c6019d
commit 926067a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 28 deletions.
diff --git a/...ClusterTest/java/org/opensearch/remotemigration/RemoteStoreMigrationSettingsUpdateIT.java b/...ClusterTest/java/org/opensearch/remotemigration/RemoteStoreMigrationSettingsUpdateIT.java
@@ -8,16 +8,22 @@
 
 package org.opensearch.remotemigration;
 
+import org.opensearch.action.admin.cluster.snapshots.create.CreateSnapshotResponse;
+import org.opensearch.action.admin.cluster.snapshots.status.SnapshotStatus;
 import org.opensearch.client.Client;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.settings.SettingsException;
+import org.opensearch.snapshots.SnapshotInfo;
+import org.opensearch.snapshots.SnapshotState;
 import org.opensearch.test.InternalTestCluster;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.List;
 import java.util.Optional;
 
+import static org.hamcrest.Matchers.is;
 import static org.opensearch.node.remotestore.RemoteStoreNodeService.CompatibilityMode.MIXED;
 import static org.opensearch.node.remotestore.RemoteStoreNodeService.CompatibilityMode.STRICT;
 import static org.opensearch.node.remotestore.RemoteStoreNodeService.Direction.REMOTE_STORE;
@@ -73,9 +79,9 @@ public void testNewRestoredIndexIsRemoteStoreBackedForRemoteStoreDirectionAndMix
 
         logger.info("Add remote and non-remote nodes");
         setClusterMode(MIXED.mode);
-        addRemote = false;
+        setAddRemote(false);
         String nonRemoteNodeName = internalCluster().startNode();
-        addRemote = true;
+        setAddRemote(true);
         String remoteNodeName = internalCluster().startNode();
         internalCluster().validateClusterFormed();
         assertNodeInCluster(nonRemoteNodeName);
@@ -84,7 +90,7 @@ public void testNewRestoredIndexIsRemoteStoreBackedForRemoteStoreDirectionAndMix
         logger.info("Create a non remote-backed index");
         createIndex(TEST_INDEX, 0);
 
-        logger.info("Verify that non remote stored backed index is created");
+        logger.info("Verify that non remote store backed index is created");
         assertNonRemoteStoreBackedIndex(TEST_INDEX);
 
         logger.info("Create repository");
@@ -94,8 +100,7 @@ public void testNewRestoredIndexIsRemoteStoreBackedForRemoteStoreDirectionAndMix
         createRepository(snapshotRepoName, "fs", Settings.builder().put("location", snapshotRepoNameAbsolutePath));
 
         logger.info("Create snapshot of non remote stored backed index");
-
-        createSnapshot(snapshotRepoName, snapshotName, TEST_INDEX);
+        createSnapshot(snapshotRepoName, snapshotName);
 
         logger.info("Restore index from snapshot under NONE direction");
         String restoredIndexName1 = TEST_INDEX + "-restored1";
@@ -164,7 +169,7 @@ private void createMixedModeCluster() {
 
     // bootstrap a cluster
     private void initializeCluster(boolean remoteClusterManager) {
-        addRemote = remoteClusterManager;
+        setAddRemote(remoteClusterManager);
         internalCluster().startClusterManagerOnlyNode();
         client = internalCluster().client();
     }

diff --git a/.../java/org/opensearch/remotemigration/RemoteStoreMigrationShardAllocationBaseTestCase.java b/.../java/org/opensearch/remotemigration/RemoteStoreMigrationShardAllocationBaseTestCase.java
@@ -10,6 +10,7 @@
 
 import org.opensearch.action.admin.cluster.allocation.ClusterAllocationExplanation;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
+import org.opensearch.action.admin.cluster.snapshots.create.CreateSnapshotResponse;
 import org.opensearch.action.admin.cluster.snapshots.restore.RestoreSnapshotResponse;
 import org.opensearch.action.support.ActiveShardCount;
 import org.opensearch.cluster.metadata.IndexMetadata;
@@ -33,7 +34,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
+import static org.hamcrest.Matchers.is;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_REMOTE_SEGMENT_STORE_REPOSITORY;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_REMOTE_STORE_ENABLED;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_REMOTE_TRANSLOG_STORE_REPOSITORY;
@@ -276,19 +280,15 @@ protected void assertAllocation(boolean isPrimary, @Nullable DiscoveryNode targe
     }
 
     // create a snapshot
-    public static SnapshotInfo createSnapshot(String snapshotRepoName, String snapshotName, String... indices) {
-        SnapshotInfo snapshotInfo = internalCluster().client()
-            .admin()
+    public static SnapshotInfo createSnapshot(String snapshotRepoName, String snapshotName) {
+        final CreateSnapshotResponse response = client().admin()
             .cluster()
             .prepareCreateSnapshot(snapshotRepoName, snapshotName)
-            .setIndices(indices)
             .setWaitForCompletion(true)
-            .get()
-            .getSnapshotInfo();
-
-        assertEquals(SnapshotState.SUCCESS, snapshotInfo.state());
-        assertTrue(snapshotInfo.successfulShards() > 0);
-        assertEquals(0, snapshotInfo.failedShards());
+            .get();
+        final SnapshotInfo snapshotInfo = response.getSnapshotInfo();
+        assertThat(snapshotInfo.successfulShards(), is(snapshotInfo.totalShards()));
+        assertThat(snapshotInfo.state(), is(SnapshotState.SUCCESS));
         return snapshotInfo;
     }
 
@@ -310,18 +310,27 @@ public static void createIndex(String indexName, int replicaCount) {
     }
 
     // restore indices from a snapshot
-    public static RestoreSnapshotResponse restoreSnapshot(String snapshotRepoName, String snapshotName, String restoredIndexName) {
-        RestoreSnapshotResponse restoreSnapshotResponse = internalCluster().client()
-            .admin()
-            .cluster()
-            .prepareRestoreSnapshot(snapshotRepoName, snapshotName)
-            .setWaitForCompletion(false)
-            .setIndices(TEST_INDEX)
-            .setRenamePattern(TEST_INDEX)
-            .setRenameReplacement(restoredIndexName)
-            .get();
-        assertEquals(restoreSnapshotResponse.status(), RestStatus.ACCEPTED);
-        return restoreSnapshotResponse;
+    public static RestoreSnapshotResponse restoreSnapshot(String snapshotRepoName, String snapshotName, String restoredIndexName) throws Exception {
+        AtomicReference<RestoreSnapshotResponse> responseRef = new AtomicReference<>();
+        assertBusy(() -> {
+            RestoreSnapshotResponse response = client()
+                .admin()
+                .cluster()
+                .prepareRestoreSnapshot(snapshotRepoName, snapshotName)
+                .setIndices(TEST_INDEX)
+                .setRenamePattern(TEST_INDEX)
+                .setRenameReplacement(restoredIndexName)
+                .setWaitForCompletion(true)
+                .execute()
+                .actionGet();
+
+            assertTrue(response.getRestoreInfo().totalShards() > 0);
+            assertEquals(0, response.getRestoreInfo().failedShards());
+
+            responseRef.set(response);
+        }, 1, TimeUnit.MINUTES);
+
+        return responseRef.get();
     }
 
     // verify that the created index is not remote store backed