Add more unit tests for RemoteStoreUtils.getPinnedTimestampLockedFiles

Signed-off-by: Sachin Kale <[email protected]>
sachinpkale · Sep 30, 2024 · d5deaf9 · d5deaf9
1 parent 8ddb3ee
commit d5deaf9
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/index/remote/RemoteStoreUtils.java b/server/src/main/java/org/opensearch/index/remote/RemoteStoreUtils.java
@@ -476,7 +476,7 @@ private static Set<String> getPinnedTimestampLockedFiles(
         Set<Long> newPinnedTimestamps = new TreeSet<>(Collections.reverseOrder());
         for (Long pinnedTimestamp : pinnedTimestampSet) {
             String cachedFile = metadataFilePinnedTimestampMap.get(pinnedTimestamp);
-            if (cachedFile != null) {
+            if (cachedFile != null && metadataFiles.contains(cachedFile)) {
                 implicitLockedFiles.add(cachedFile);
             } else {
                 newPinnedTimestamps.add(pinnedTimestamp);

diff --git a/server/src/test/java/org/opensearch/index/remote/RemoteStoreUtilsTests.java b/server/src/test/java/org/opensearch/index/remote/RemoteStoreUtilsTests.java
@@ -1083,4 +1083,92 @@ public void testIsPinnedTimestampStateStaleFeatureEnabled() {
         setupRemotePinnedTimestampFeature(true);
         assertTrue(RemoteStoreUtils.isPinnedTimestampStateStale());
     }
+
+    public void testGetPinnedTimestampLockedFilesWithCache() {
+        setupRemotePinnedTimestampFeature(true);
+
+        Map<Long, String> metadataFilePinnedTimestampCache = new HashMap<>();
+
+        // Pinned timestamps 800, 900, 1000, 2000
+        // Metadata with timestamp 990, 995, 1000, 1001
+        // Metadata timestamp 1000 <= Pinned Timestamp 1000
+        // Metadata timestamp 1001 <= Pinned Timestamp 2000
+        Tuple<Map<Long, String>, Set<String>> metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(990L, 995L, 1000L, 1001L),
+            Set.of(800L, 900L, 1000L, 2000L),
+            metadataFilePinnedTimestampCache
+        );
+        Map<Long, String> metadataFiles = metadataAndLocks.v1();
+        Set<String> implicitLockedFiles = metadataAndLocks.v2();
+
+        assertEquals(2, implicitLockedFiles.size());
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(1000L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(1001L)));
+        // Now we cache all the matches except the last one.
+        assertEquals(1, metadataFilePinnedTimestampCache.size());
+        assertEquals(metadataFiles.get(1000L), metadataFilePinnedTimestampCache.get(1000L));
+
+        metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(990L, 995L, 1000L, 1001L, 2000L, 2200L),
+            Set.of(800L, 900L, 1000L, 2000L, 3000L),
+            metadataFilePinnedTimestampCache
+        );
+        metadataFiles = metadataAndLocks.v1();
+        implicitLockedFiles = metadataAndLocks.v2();
+        assertEquals(3, implicitLockedFiles.size());
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(1000L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(2000L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(2200L)));
+        assertEquals(2, metadataFilePinnedTimestampCache.size());
+        assertEquals(metadataFiles.get(1000L), metadataFilePinnedTimestampCache.get(1000L));
+        assertEquals(metadataFiles.get(2000L), metadataFilePinnedTimestampCache.get(2000L));
+
+        metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(990L, 995L, 1000L, 1001L, 2000L, 2200L, 2500L),
+            Set.of(2000L, 3000L),
+            metadataFilePinnedTimestampCache
+        );
+        metadataFiles = metadataAndLocks.v1();
+        implicitLockedFiles = metadataAndLocks.v2();
+        assertEquals(2, implicitLockedFiles.size());
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(2000L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(2500L)));
+        assertEquals(1, metadataFilePinnedTimestampCache.size());
+        assertEquals(metadataFiles.get(2000L), metadataFilePinnedTimestampCache.get(2000L));
+
+        metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(2000L, 2200L, 2500L, 3001L, 4200L, 4600L, 5010L),
+            Set.of(3000L, 4000L, 5000L, 6000L),
+            metadataFilePinnedTimestampCache
+        );
+        metadataFiles = metadataAndLocks.v1();
+        implicitLockedFiles = metadataAndLocks.v2();
+        assertEquals(4, implicitLockedFiles.size());
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(2500L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(3001L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(4600L)));
+        assertTrue(implicitLockedFiles.contains(metadataFiles.get(5010L)));
+        assertEquals(3, metadataFilePinnedTimestampCache.size());
+        assertEquals(metadataFiles.get(2500L), metadataFilePinnedTimestampCache.get(3000L));
+        assertEquals(metadataFiles.get(3001L), metadataFilePinnedTimestampCache.get(4000L));
+        assertEquals(metadataFiles.get(4600L), metadataFilePinnedTimestampCache.get(5000L));
+
+        metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(),
+            Set.of(3000L, 4000L, 5000L, 6000L),
+            metadataFilePinnedTimestampCache
+        );
+        implicitLockedFiles = metadataAndLocks.v2();
+        assertEquals(0, implicitLockedFiles.size());
+        assertEquals(3, metadataFilePinnedTimestampCache.size());
+
+        metadataAndLocks = testGetPinnedTimestampLockedFilesWithPinnedTimestamps(
+            List.of(2000L, 2200L, 2500L, 3001L, 4200L, 4600L, 5010L),
+            Set.of(),
+            metadataFilePinnedTimestampCache
+        );
+        implicitLockedFiles = metadataAndLocks.v2();
+        assertEquals(0, implicitLockedFiles.size());
+        assertEquals(0, metadataFilePinnedTimestampCache.size());
+    }
 }