[server] Added Retry Handler to Address Snapshotting Test Asynchrony

Author: rionmonster

fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTest.java

@@ -498,34 +498,34 @@ void testBrokenSnapshotRecovery(@TempDir File snapshotKvTabletDir) throws Except
         // put initial data and create first snapshot
         KvRecordBatch kvRecords =
                 genKvRecordBatch(
-                        Tuple2.of("k1", new Object[] {1, "a"}),
-                        Tuple2.of("k2", new Object[] {2, "b"}));
+                        Tuple2.of("k1", new Object[]{1, "a"}),
+                        Tuple2.of("k2", new Object[]{2, "b"}));
         putRecordsToLeader(kvReplica, kvRecords);
 
         // trigger first snapshot
-        scheduledExecutorService.triggerNonPeriodicScheduledTask();
+        triggerSnapshotTaskWithRetry(scheduledExecutorService, 5);
         kvSnapshotStore.waitUntilSnapshotComplete(tableBucket, 0);
 
         // put more data and create second snapshot
         kvRecords =
                 genKvRecordBatch(
-                        Tuple2.of("k1", new Object[] {3, "c"}),
-                        Tuple2.of("k3", new Object[] {4, "d"}));
+                        Tuple2.of("k1", new Object[]{3, "c"}),
+                        Tuple2.of("k3", new Object[]{4, "d"}));
         putRecordsToLeader(kvReplica, kvRecords);
 
         // trigger second snapshot
-        scheduledExecutorService.triggerNonPeriodicScheduledTask();
+        triggerSnapshotTaskWithRetry(scheduledExecutorService, 5);
         kvSnapshotStore.waitUntilSnapshotComplete(tableBucket, 1);
 
         // put more data and create third snapshot (this will be the broken one)
         kvRecords =
                 genKvRecordBatch(
-                        Tuple2.of("k4", new Object[] {5, "e"}),
-                        Tuple2.of("k5", new Object[] {6, "f"}));
+                        Tuple2.of("k4", new Object[]{5, "e"}),
+                        Tuple2.of("k5", new Object[]{6, "f"}));
         putRecordsToLeader(kvReplica, kvRecords);
 
         // trigger third snapshot
-        scheduledExecutorService.triggerNonPeriodicScheduledTask();
+        triggerSnapshotTaskWithRetry(scheduledExecutorService, 5);
         CompletedSnapshot snapshot2 = kvSnapshotStore.waitUntilSnapshotComplete(tableBucket, 2);
 
         // verify that snapshot2 is the latest one before we break it
@@ -566,8 +566,8 @@ void testBrokenSnapshotRecovery(@TempDir File snapshotKvTabletDir) throws Except
         List<Tuple2<byte[], byte[]>> expectedKeyValues =
                 getKeyValuePairs(
                         genKvRecords(
-                                Tuple2.of("k1", new Object[] {3, "c"}),
-                                Tuple2.of("k3", new Object[] {4, "d"})));
+                                Tuple2.of("k1", new Object[]{3, "c"}),
+                                Tuple2.of("k3", new Object[]{4, "d"})));
         verifyGetKeyValues(kvTablet, expectedKeyValues);
 
         // Verify the core functionality: KvTablet successfully initialized despite broken snapshot
@@ -766,4 +766,22 @@ public void reset() {
             isScheduled = false;
         }
     }
+
+    /** A helper function with support for retries for flaky triggering operations. */
+    private static void triggerSnapshotTaskWithRetry(
+            ManuallyTriggeredScheduledExecutorService scheduledExecutorService,
+            int maxRetries) throws Exception {
+        for (int i = 0; i < maxRetries; i++) {
+            try {
+                scheduledExecutorService.triggerNonPeriodicScheduledTask();
+                return;
+            } catch (java.util.NoSuchElementException e) {
+                if (i == maxRetries - 1) {
+                    throw e;
+                }
+
+                Thread.sleep(50);
+            }
+        }
+    }
 }