From e6bcc79b624cf21de30b9cdf4a9430dbd4dd79f0 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   | 18 ++++++++++++++++++
 .../org/opensearch/index/shard/IndexShard.java | 16 +++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
index f8e5079b01a36..11bf122cb444a 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,22 @@ 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()) {
+                assertThat(recoveryState.getIndex().totalFileCount(), greaterThan(0));
+                assertThat(recoveryState.getIndex().totalRecoverFiles(), greaterThan(0));
+                assertThat(recoveryState.getIndex().recoveredFileCount(), greaterThan(0));
+                assertThat(recoveryState.getIndex().recoveredFilesPercent(), greaterThan(0f));
+
+                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..b42d7f2afa3c0 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,24 @@ 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();
+            }
+
             String segmentsNFile = copySegmentFiles(
                 storeDirectory,
                 sourceRemoteDirectory,