Address PR Comments

Signed-off-by: Gaurav Bafna <[email protected]>
gbbafna · Oct 10, 2024 · 5d14a37 · 5d14a37
1 parent 0089d64
commit 5d14a37
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 18 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteRestoreSnapshotIT.java
@@ -963,15 +963,17 @@ public void testConcurrentSnapshotV2CreateOperation() throws InterruptedExceptio
             thread.join();
         }
 
+        client().admin().cluster().prepareCreateSnapshot(snapshotRepoName, "snapshot-cleanup-timestamp").setWaitForCompletion(true).get();
         // Validate that only one snapshot has been created
         Repository repository = internalCluster().getInstance(RepositoriesService.class).repository(snapshotRepoName);
         PlainActionFuture<RepositoryData> repositoryDataPlainActionFuture = new PlainActionFuture<>();
         repository.getRepositoryData(repositoryDataPlainActionFuture);
-
         RepositoryData repositoryData = repositoryDataPlainActionFuture.get();
-        assertThat(repositoryData.getSnapshotIds().size(), greaterThanOrEqualTo(1));
-        forceSyncPinnedTimestamps();
-        assertEquals(RemoteStorePinnedTimestampService.getPinnedEntities().size(), repositoryData.getSnapshotIds().size());
+        assertThat(repositoryData.getSnapshotIds().size(), greaterThanOrEqualTo(2));
+        waitUntil(() -> {
+            forceSyncPinnedTimestamps();
+            return RemoteStorePinnedTimestampService.getPinnedEntities().size() == repositoryData.getSnapshotIds().size();
+        });
     }
 
     public void testConcurrentSnapshotV2CreateOperation_MasterChange() throws Exception {

diff --git a/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java b/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java
@@ -724,23 +724,22 @@ private void createSnapshotPreValidations(
     private void updateSnapshotPinnedTimestamp(Snapshot snapshot, long timestampToPin) throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
         SetOnce<Exception> ex = new SetOnce<>();
+        ActionListener<Void> listener = new ActionListener<>() {
+            @Override
+            public void onResponse(Void unused) {
+                logger.debug("Timestamp pinned successfully for snapshot {}", snapshot.getSnapshotId().getName());
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Failed to pin timestamp for snapshot {} with exception {}", snapshot.getSnapshotId().getName(), e);
+                ex.set(e);
+            }
+        };
         remoteStorePinnedTimestampService.pinTimestamp(
             timestampToPin,
             getPinningEntity(snapshot.getRepository(), snapshot.getSnapshotId().getUUID()),
-            new ActionListener<>() {
-                @Override
-                public void onResponse(Void unused) {
-                    logger.debug("Timestamp pinned successfully for snapshot {}", snapshot.getSnapshotId().getName());
-                    latch.countDown();
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    logger.error("Failed to pin timestamp for snapshot {} with exception {}", snapshot.getSnapshotId().getName(), e);
-                    ex.set(e);
-                    latch.countDown();
-                }
-            }
+            new LatchedActionListener<>(listener, latch)
         );
         latch.await();
         if (ex.get() != null) {