[FLINK-37900][rocksdb] support configurable jitter in rocksdb state uploader

HuangZhenQiu · HuangZhenQiu · commit 38abdb460186 · 2025-06-12T05:16:47.000-07:00
diff --git a/docs/layouts/shortcodes/generated/expert_rocksdb_section.html b/docs/layouts/shortcodes/generated/expert_rocksdb_section.html
@@ -14,6 +14,12 @@
             <td>Integer</td>
             <td>The number of threads (per stateful operator) used to transfer (download and upload) files in RocksDBStateBackend.If negative, the common (TM) IO thread pool is used (see cluster.io-pool.size)</td>
         </tr>
+        <tr>
+            <td><h5>state.backend.rocksdb.checkpoint.upload-jitter</h5></td>
+            <td style="word-wrap: break-word;">0 ms</td>
+            <td>Duration</td>
+            <td>The time interval that will be used to create jitter for each checkpoint file upload.</td>
+        </tr>
         <tr>
             <td><h5>state.backend.rocksdb.localdir</h5></td>
             <td style="word-wrap: break-word;">(none)</td>
diff --git a/docs/layouts/shortcodes/generated/rocksdb_configuration.html b/docs/layouts/shortcodes/generated/rocksdb_configuration.html
@@ -14,6 +14,12 @@
             <td>Integer</td>
             <td>The number of threads (per stateful operator) used to transfer (download and upload) files in RocksDBStateBackend.If negative, the common (TM) IO thread pool is used (see cluster.io-pool.size)</td>
         </tr>
+        <tr>
+            <td><h5>state.backend.rocksdb.checkpoint.upload-jitter</h5></td>
+            <td style="word-wrap: break-word;">0 ms</td>
+            <td>Duration</td>
+            <td>The time interval that will be used to create jitter for each checkpoint file upload.</td>
+        </tr>
         <tr>
             <td><h5>state.backend.rocksdb.localdir</h5></td>
             <td style="word-wrap: break-word;">(none)</td>
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackendBuilder.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackendBuilder.java
@@ -72,6 +72,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashMap;
@@ -132,6 +133,7 @@ public class RocksDBKeyedStateBackendBuilder<K> extends AbstractKeyedStateBacken
     private RocksDBNativeMetricOptions nativeMetricOptions;
 
     private int numberOfTransferingThreads;
+    private Duration uploadJitter;
     private long writeBatchSize =
             RocksDBConfigurableOptions.WRITE_BATCH_SIZE.defaultValue().getBytes();
 
@@ -200,6 +202,7 @@ public RocksDBKeyedStateBackendBuilder(
         this.nativeMetricOptions = new RocksDBNativeMetricOptions();
         this.numberOfTransferingThreads =
                 RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue();
+        this.uploadJitter = RocksDBOptions.CHECKPOINT_UPLOAD_JITTER.defaultValue();
     }
 
     @VisibleForTesting
@@ -626,7 +629,8 @@ private RocksDBRestoreOperation getRocksDBRestoreOperation(
                 injectRocksDBStateUploader == null
                         ? new RocksDBStateUploader(
                                 RocksDBStateDataTransferHelper.forThreadNumIfSpecified(
-                                        numberOfTransferingThreads, ioExecutor))
+                                        numberOfTransferingThreads, ioExecutor),
+                                uploadJitter)
                         : injectRocksDBStateUploader;
         if (enableIncrementalCheckpointing) {
             checkpointSnapshotStrategy =
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBOptions.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBOptions.java
@@ -27,6 +27,8 @@
 import org.apache.flink.configuration.description.Description;
 import org.apache.flink.configuration.description.TextElement;
 
+import java.time.Duration;
+
 import static org.apache.flink.configuration.ClusterOptions.CLUSTER_IO_EXECUTOR_POOL_SIZE;
 import static org.apache.flink.state.rocksdb.EmbeddedRocksDBStateBackend.PriorityQueueStateType.ROCKSDB;
 import static org.apache.flink.state.rocksdb.PredefinedOptions.DEFAULT;
@@ -91,6 +93,15 @@ public class RocksDBOptions {
                                     + CLUSTER_IO_EXECUTOR_POOL_SIZE.key()
                                     + ")");
 
+    /** The number of millisecond for RocksDBStateBackend checkpoint file upload jitter. */
+    @Documentation.Section(Documentation.Sections.EXPERT_ROCKSDB)
+    public static final ConfigOption<Duration> CHECKPOINT_UPLOAD_JITTER =
+            ConfigOptions.key("state.backend.rocksdb.checkpoint.upload-jitter")
+                    .durationType()
+                    .defaultValue(Duration.ZERO)
+                    .withDescription(
+                            "The time interval that will be used to create jitter for each checkpoint file upload.");
+
     /** The predefined settings for RocksDB DBOptions and ColumnFamilyOptions by Flink community. */
     @Documentation.Section(Documentation.Sections.EXPERT_ROCKSDB)
     public static final ConfigOption<String> PREDEFINED_OPTIONS =
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBStateUploader.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBStateUploader.java
@@ -39,25 +39,32 @@
 import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 
 /** Help class for uploading RocksDB state files. */
 public class RocksDBStateUploader implements Closeable {
     private static final int READ_BUFFER_SIZE = 16 * 1024;
-
+    private final Duration uploadJitter;
+    private final Random random;
     private final RocksDBStateDataTransferHelper transfer;
 
     @VisibleForTesting
-    public RocksDBStateUploader(int numberOfSnapshottingThreads) {
-        this(RocksDBStateDataTransferHelper.forThreadNum(numberOfSnapshottingThreads));
+    public RocksDBStateUploader(int numberOfSnapshottingThreads, Duration uploadJitter) {
+        this(
+                RocksDBStateDataTransferHelper.forThreadNum(numberOfSnapshottingThreads),
+                uploadJitter);
     }
 
-    public RocksDBStateUploader(RocksDBStateDataTransferHelper transfer) {
+    public RocksDBStateUploader(RocksDBStateDataTransferHelper transfer, Duration uploadJitter) {
         this.transfer = transfer;
+        this.uploadJitter = uploadJitter;
+        this.random = new Random(uploadJitter.toMillis());
     }
 
     /**
@@ -133,12 +140,14 @@ private HandleAndLocalPath uploadLocalFileToCheckpointFs(
             CheckpointedStateScope stateScope,
             CloseableRegistry closeableRegistry,
             CloseableRegistry tmpResourcesRegistry)
-            throws IOException {
+            throws IOException, InterruptedException {
 
         InputStream inputStream = null;
         CheckpointStateOutputStream outputStream = null;
 
         try {
+            // add a random jitter
+            applyJitter();
             final byte[] buffer = new byte[READ_BUFFER_SIZE];
 
             inputStream = Files.newInputStream(filePath);
@@ -180,6 +189,17 @@ private HandleAndLocalPath uploadLocalFileToCheckpointFs(
         }
     }
 
+    @VisibleForTesting
+    long applyJitter() throws InterruptedException {
+        if (uploadJitter.isZero()) {
+            return 0;
+        }
+
+        long sleepMilliseconds = random.nextLong(0, uploadJitter.toMillis() + 1);
+        Thread.sleep(sleepMilliseconds);
+        return sleepMilliseconds;
+    }
+
     @Override
     public void close() throws IOException {
         this.transfer.close();
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/EmbeddedRocksDBStateBackendTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/EmbeddedRocksDBStateBackendTest.java
@@ -282,7 +282,8 @@ public void setupRocksKeyedStateBackend() throws Exception {
             rocksDBStateUploader =
                     spy(
                             new RocksDBStateUploader(
-                                    RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue()));
+                                    RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue(),
+                                    RocksDBOptions.CHECKPOINT_UPLOAD_JITTER.defaultValue()));
             keyedStateBackendBuilder.setRocksDBStateUploader(rocksDBStateUploader);
         }
 
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/RocksDBStateUploaderTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/RocksDBStateUploaderTest.java
@@ -42,6 +42,7 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -91,7 +92,8 @@ public List<StreamStateHandle> duplicate(
 
         List<Path> filePaths = new ArrayList<>(1);
         filePaths.add(file.toPath());
-        try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
+        try (RocksDBStateUploader rocksDBStateUploader =
+                new RocksDBStateUploader(5, Duration.ZERO)) {
             assertThatThrownBy(
                             () ->
                                     rocksDBStateUploader.uploadFilesToCheckpointFs(
@@ -136,7 +138,8 @@ void testUploadedSstCanBeCleanedUp() throws Exception {
         List<Path> filePaths =
                 generateRandomSstFiles(localFolder, sstFileCount, fileStateSizeThreshold);
         CloseableRegistry tmpResourcesRegistry = new CloseableRegistry();
-        try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(1)) {
+        try (RocksDBStateUploader rocksDBStateUploader =
+                new RocksDBStateUploader(1, Duration.ZERO)) {
             rocksDBStateUploader.uploadFilesToCheckpointFs(
                     filePaths,
                     checkpointStreamFactory,
@@ -211,7 +214,8 @@ void testMultiThreadUploadCorrectly() throws Exception {
         List<Path> sstFilePaths =
                 generateRandomSstFiles(localFolder, sstFileCount, fileStateSizeThreshold);
 
-        try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
+        try (RocksDBStateUploader rocksDBStateUploader =
+                new RocksDBStateUploader(5, Duration.ZERO)) {
             List<HandleAndLocalPath> sstFiles =
                     rocksDBStateUploader.uploadFilesToCheckpointFs(
                             sstFilePaths,
@@ -233,6 +237,26 @@ void testMultiThreadUploadCorrectly() throws Exception {
         }
     }
 
+    @Test
+    void testApplyJitter() throws Exception {
+        try (RocksDBStateUploader rocksDBStateUploader =
+                new RocksDBStateUploader(1, Duration.ofMillis(1000L))) {
+            long startTime = System.currentTimeMillis();
+            long milliseconds = rocksDBStateUploader.applyJitter();
+            assertThat(milliseconds).isLessThanOrEqualTo(1000);
+            assertThat(System.currentTimeMillis() - startTime).isGreaterThanOrEqualTo(milliseconds);
+        }
+    }
+
+    @Test
+    void testApplyDefaultJitter() throws Exception {
+        try (RocksDBStateUploader rocksDBStateUploader =
+                new RocksDBStateUploader(1, Duration.ZERO)) {
+            long milliseconds = rocksDBStateUploader.applyJitter();
+            assertThat(milliseconds).isEqualTo(0L);
+        }
+    }
+
     private static CheckpointStateOutputStream createFailingCheckpointStateOutputStream(
             IOException failureException) {
         return new CheckpointStateOutputStream() {
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/snapshot/RocksIncrementalSnapshotStrategyTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/snapshot/RocksIncrementalSnapshotStrategyTest.java
@@ -114,7 +114,8 @@ public RocksIncrementalSnapshotStrategy<?> createSnapshotStrategy()
 
         RocksDBStateUploader rocksDBStateUploader =
                 new RocksDBStateUploader(
-                        RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue());
+                        RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue(),
+                        RocksDBOptions.CHECKPOINT_UPLOAD_JITTER.defaultValue());
 
         int keyGroupPrefixBytes =
                 CompositeKeySerializationUtils.computeRequiredBytesInKeyGroupPrefix(2);

Original file line number	Diff line number	Diff line change
`@@ -282,7 +282,8 @@ public void setupRocksKeyedStateBackend() throws Exception {`
`282`	`282`	`rocksDBStateUploader =`
`283`	`283`	`spy(`
`284`	`284`	`new RocksDBStateUploader(`
`285`		`- RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue()));`
	`285`	`+ RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue(),`
	`286`	`+ RocksDBOptions.CHECKPOINT_UPLOAD_JITTER.defaultValue()));`
`286`	`287`	`keyedStateBackendBuilder.setRocksDBStateUploader(rocksDBStateUploader);`
`287`	`288`	`}`
`288`	`289`