Address PR comments

Signed-off-by: Gaurav Bafna <[email protected]>
opensearch-project · Oct 3, 2024 · 14f0323 · 14f0323
1 parent 48f147f
commit 14f0323
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
@@ -893,8 +893,8 @@ public void testMixedSnapshotCreationWithV2RepositorySetting() throws Exception
         assertThat(snapshotInfo.successfulShards(), equalTo(snapshotInfo.totalShards()));
         assertThat(snapshotInfo.snapshotId().getName(), equalTo(snapshotName2));
         assertThat(snapshotInfo.getPinnedTimestamp(), greaterThan(0L));
+        forceSyncPinnedTimestamps();
         assertEquals(RemoteStorePinnedTimestampService.getPinnedEntities().size(), 1);
-
     }
 
     public void testConcurrentSnapshotV2CreateOperation() throws InterruptedException, ExecutionException {
@@ -970,6 +970,7 @@ public void testConcurrentSnapshotV2CreateOperation() throws InterruptedExceptio
 
         RepositoryData repositoryData = repositoryDataPlainActionFuture.get();
         assertThat(repositoryData.getSnapshotIds().size(), greaterThanOrEqualTo(1));
+        forceSyncPinnedTimestamps();
         assertEquals(RemoteStorePinnedTimestampService.getPinnedEntities().size(), repositoryData.getSnapshotIds().size());
     }
 
@@ -1044,12 +1045,11 @@ public void testConcurrentSnapshotV2CreateOperation_MasterChange() throws Except
 
         }, 30, TimeUnit.SECONDS);
         thread.join();
-        RemoteStorePinnedTimestampService remoteStorePinnedTimestampService = internalCluster().getInstance(
-            RemoteStorePinnedTimestampService.class,
-            internalCluster().getClusterManagerName()
-        );
         forceSyncPinnedTimestamps();
-        assertEquals(RemoteStorePinnedTimestampService.getPinnedEntities().size(), totalSnaps.intValue());
+        waitUntil(() -> {
+            this.forceSyncPinnedTimestamps();
+            return RemoteStorePinnedTimestampService.getPinnedEntities().size() == totalSnaps.intValue();
+        });
     }
 
     public void testCreateSnapshotV2() throws Exception {

diff --git a/...rc/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java b/...rc/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java
@@ -19,6 +19,9 @@
 import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
@@ -76,6 +79,11 @@ public void testTimestampPinUnpin() throws Exception {
             Tuple<Long, Set<Long>> pinnedTimestampWithFetchTimestamp_2 = RemoteStorePinnedTimestampService.getPinnedTimestamps();
             long lastFetchTimestamp_2 = pinnedTimestampWithFetchTimestamp_2.v1();
             assertTrue(lastFetchTimestamp_2 != -1);
+            Map<String, List<Long>> pinnedEntities = RemoteStorePinnedTimestampService.getPinnedEntities();
+            assertEquals(3, pinnedEntities.size());
+            assertEquals(Set.of("ss2", "ss3", "ss4"), pinnedEntities.keySet());
+            assertEquals(pinnedEntities.get("ss2").size(), 1);
+            assertEquals(Optional.of(timestamp1).get(), pinnedEntities.get("ss2").get(0));
             assertEquals(Set.of(timestamp1, timestamp2, timestamp3), pinnedTimestampWithFetchTimestamp_2.v2());
         });
 
@@ -103,10 +111,14 @@ public void onFailure(Exception e) {
         // Adding different entity to already pinned timestamp
         remoteStorePinnedTimestampService.pinTimestamp(timestamp3, "ss5", noOpActionListener);
 
-        remoteStorePinnedTimestampService.rescheduleAsyncUpdatePinnedTimestampTask(TimeValue.timeValueSeconds(1));
+        remoteStorePinnedTimestampService.forceSyncPinnedTimestamps();
 
         assertBusy(() -> {
             Tuple<Long, Set<Long>> pinnedTimestampWithFetchTimestamp_3 = RemoteStorePinnedTimestampService.getPinnedTimestamps();
+            Map<String, List<Long>> pinnedEntities = RemoteStorePinnedTimestampService.getPinnedEntities();
+            assertEquals(3, pinnedEntities.size());
+            assertEquals(pinnedEntities.get("ss5").size(), 1);
+            assertEquals(Optional.of(timestamp3).get(), pinnedEntities.get("ss5").get(0));
             long lastFetchTimestamp_3 = pinnedTimestampWithFetchTimestamp_3.v1();
             assertTrue(lastFetchTimestamp_3 != -1);
             assertEquals(Set.of(timestamp1, timestamp3), pinnedTimestampWithFetchTimestamp_3.v2());

diff --git a/server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java b/server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java
@@ -301,6 +301,7 @@ protected void runInternal() {
                 Map<String, BlobMetadata> pinnedTimestampList = blobContainer.listBlobs();
                 if (pinnedTimestampList.isEmpty()) {
                     pinnedTimestampsSet = new Tuple<>(triggerTimestamp, Set.of());
+                    pinnedEntityToTimestampsMap = new HashMap<>();
                     return;
                 }
                 Set<Long> pinnedTimestamps = pinnedTimestampList.keySet()