From 7bccf910df8cb5556743223810f2d6e89a6e550a Mon Sep 17 00:00:00 2001
From: Lakshya Taragi <lakshya.taragi@gmail.com>
Date: Thu, 22 Aug 2024 17:38:56 +0530
Subject: [PATCH] Populate RecoveryState details for shallow snapshot restore

Signed-off-by: Lakshya Taragi <lakshya.taragi@gmail.com>
---
 .../remotestore/RemoteRestoreSnapshotIT.java  | 20 +++++++++++++++++++
 .../opensearch/index/shard/IndexShard.java    | 19 +++++++++++++++---
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
index f8e5079b01a36..e8720b3fff059 100644
--- a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
@@ -12,6 +12,7 @@
 import org.opensearch.action.admin.cluster.remotestore.restore.RestoreRemoteStoreRequest;
 import org.opensearch.action.admin.cluster.snapshots.restore.RestoreSnapshotResponse;
 import org.opensearch.action.admin.indices.delete.DeleteIndexRequest;
+import org.opensearch.action.admin.indices.recovery.RecoveryResponse;
 import org.opensearch.action.delete.DeleteResponse;
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.client.Client;
@@ -31,6 +32,7 @@
 import org.opensearch.index.remote.RemoteStoreEnums.PathType;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.indices.IndicesService;
+import org.opensearch.indices.recovery.RecoveryState;
 import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.snapshots.AbstractSnapshotIntegTestCase;
 import org.opensearch.snapshots.SnapshotInfo;
@@ -579,6 +581,24 @@ public void testRestoreShallowSnapshotRepository() throws ExecutionException, In
         ensureGreen(restoredIndexName1);
         assertDocsPresentInIndex(client, restoredIndexName1, numDocsInIndex1);
 
+        // ensure recovery details are non-zero
+        RecoveryResponse recoveryResponse = client().admin().indices().prepareRecoveries(restoredIndexName1).execute().actionGet();
+        for (Map.Entry<String, List<RecoveryState>> entry : recoveryResponse.shardRecoveryStates().entrySet()) {
+            for (RecoveryState recoveryState : entry.getValue()) {
+                // ensure populated file details
+                assertThat(recoveryState.getIndex().totalFileCount(), greaterThan(0));
+                assertThat(recoveryState.getIndex().totalRecoverFiles(), greaterThan(0));
+                assertThat(recoveryState.getIndex().recoveredFileCount(), greaterThan(0));
+                assertThat(recoveryState.getIndex().recoveredFilesPercent(), greaterThan(0f));
+
+                // ensure populated bytes details
+                assertThat(recoveryState.getIndex().recoveredBytes(), greaterThan(0L));
+                assertThat(recoveryState.getIndex().totalBytes(), greaterThan(0L));
+                assertThat(recoveryState.getIndex().totalRecoverBytes(), greaterThan(0L));
+                assertThat(recoveryState.getIndex().recoveredBytesPercent(), greaterThan(0f));
+            }
+        }
+
         // indexing some new docs and validating
         indexDocuments(client, restoredIndexName1, numDocsInIndex1, numDocsInIndex1 + 2);
         ensureGreen(restoredIndexName1);
diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java
index 82b68b32f3bf8..e65d1b0754dc5 100644
--- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java
+++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java
@@ -5117,10 +5117,23 @@ public void syncSegmentsFromGivenRemoteSegmentStore(
         }
         Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegments = sourceRemoteDirectory
             .getSegmentsUploadedToRemoteStore();
-        final Directory storeDirectory = store.directory();
-        store.incRef();
-
         try {
+            final Directory storeDirectory;
+            if (recoveryState.getStage() == RecoveryState.Stage.INDEX) {
+                storeDirectory = new StoreRecovery.StatsDirectoryWrapper(store.directory(), recoveryState.getIndex());
+                for (String file : uploadedSegments.keySet()) {
+                    long checksum = Long.parseLong(uploadedSegments.get(file).getChecksum());
+                    if (overrideLocal || localDirectoryContains(storeDirectory, file, checksum) == false) {
+                        recoveryState.getIndex().addFileDetail(file, uploadedSegments.get(file).getLength(), false);
+                    } else {
+                        recoveryState.getIndex().addFileDetail(file, uploadedSegments.get(file).getLength(), true);
+                    }
+                }
+            } else {
+                storeDirectory = store.directory();
+            }
+            store.incRef();
+
             String segmentsNFile = copySegmentFiles(
                 storeDirectory,
                 sourceRemoteDirectory,