From 3f2262ed8895b281b837596c9d33d2df2e5e5ac1 Mon Sep 17 00:00:00 2001
From: Ashish Singh <ssashish@amazon.com>
Date: Tue, 19 Mar 2024 13:20:55 +0530
Subject: [PATCH] Use cluster default remote store path type as fallback during
 snapshot restore

Signed-off-by: Ashish Singh <ssashish@amazon.com>
---
 .../metadata/MetadataCreateIndexService.java  | 22 ++++++++++++++-----
 .../opensearch/snapshots/RestoreService.java  |  8 ++++++-
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java b/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
index f117d1a4a11a2..f6a14d8ec9d63 100644
--- a/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
+++ b/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
@@ -553,11 +553,7 @@ IndexMetadata buildAndValidateTemporaryIndexMetadata(
         tmpImdBuilder.setRoutingNumShards(routingNumShards);
         tmpImdBuilder.settings(indexSettings);
         tmpImdBuilder.system(isSystem);
-
-        if (remoteStorePathResolver != null) {
-            String pathType = remoteStorePathResolver.resolveType().toString();
-            tmpImdBuilder.putCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY, Map.of(RemoteStorePathType.NAME, pathType));
-        }
+        addRemoteCustomData(tmpImdBuilder);
 
         // Set up everything, now locally create the index to see that things are ok, and apply
         IndexMetadata tempMetadata = tmpImdBuilder.build();
@@ -566,6 +562,22 @@ IndexMetadata buildAndValidateTemporaryIndexMetadata(
         return tempMetadata;
     }
 
+    public void addRemoteCustomData(IndexMetadata.Builder tmpImdBuilder) {
+        if (remoteStorePathResolver != null) {
+            // It is possible that remote custom data exists already. In such cases, we need to only update the path type
+            // in the remote store custom data map.
+            Map<String, String> existingRemoteCustomData = tmpImdBuilder.removeCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY);
+            Map<String, String> remoteCustomData = existingRemoteCustomData == null
+                ? new HashMap<>()
+                : new HashMap<>(existingRemoteCustomData);
+            // Determine the path type for use using the remoteStorePathResolver.
+            String newPathType = remoteStorePathResolver.resolveType().toString();
+            String oldPathType = remoteCustomData.put(RemoteStorePathType.NAME, newPathType);
+            logger.trace(() -> new ParameterizedMessage("Added new path type {}, replaced old path type {}", newPathType, oldPathType));
+            tmpImdBuilder.putCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY, remoteCustomData);
+        }
+    }
+
     private ClusterState applyCreateIndexRequestWithV1Templates(
         final ClusterState currentState,
         final CreateIndexClusterStateUpdateRequest request,
diff --git a/server/src/main/java/org/opensearch/snapshots/RestoreService.java b/server/src/main/java/org/opensearch/snapshots/RestoreService.java
index bf2c7fc74be92..a90bd71633581 100644
--- a/server/src/main/java/org/opensearch/snapshots/RestoreService.java
+++ b/server/src/main/java/org/opensearch/snapshots/RestoreService.java
@@ -72,6 +72,7 @@
 import org.opensearch.cluster.service.ClusterManagerTaskKeys;
 import org.opensearch.cluster.service.ClusterManagerTaskThrottler;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.Priority;
 import org.opensearch.common.UUIDs;
 import org.opensearch.common.lucene.Lucene;
@@ -85,6 +86,7 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.IndexModule;
 import org.opensearch.index.IndexSettings;
+import org.opensearch.index.remote.RemoteStorePathResolver;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.snapshots.IndexShardSnapshotStatus;
 import org.opensearch.index.store.remote.filecache.FileCacheStats;
@@ -191,6 +193,9 @@ public class RestoreService implements ClusterStateApplier {
 
     private static final CleanRestoreStateTaskExecutor cleanRestoreStateTaskExecutor = new CleanRestoreStateTaskExecutor();
 
+    @Nullable
+    private final RemoteStorePathResolver remoteStorePathResolver;
+
     public RestoreService(
         ClusterService clusterService,
         RepositoriesService repositoriesService,
@@ -216,7 +221,7 @@ public RestoreService(
 
         // Task is onboarded for throttling, it will get retried from associated TransportClusterManagerNodeAction.
         restoreSnapshotTaskKey = clusterService.registerClusterManagerTask(ClusterManagerTaskKeys.RESTORE_SNAPSHOT_KEY, true);
-
+        remoteStorePathResolver = new RemoteStorePathResolver(clusterService.getClusterSettings());
     }
 
     /**
@@ -452,6 +457,7 @@ public ClusterState execute(ClusterState currentState) {
                                             .put(snapshotIndexMetadata.getSettings())
                                             .put(IndexMetadata.SETTING_INDEX_UUID, UUIDs.randomBase64UUID())
                                     );
+                                    createIndexService.addRemoteCustomData(indexMdBuilder);
                                     shardLimitValidator.validateShardLimit(
                                         renamedIndexName,
                                         snapshotIndexMetadata.getSettings(),