azagrebin commented on a change in pull request #7351: [FLINK-11008][State
Backends, Checkpointing]SpeedUp upload state files using multithread
URL: https://github.com/apache/flink/pull/7351#discussion_r246341548
##########
File path:
flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateDataTransferTest.java
##########
@@ -138,6 +148,157 @@ public void testMultiThreadRestoreCorrectly() throws
Exception {
}
}
+ /**
+ * Test that the exception arose in the thread pool will rethrow to the
main thread.
+ */
+ @Test
+ public void testMultiThreadUploadThreadPoolExceptionRethrow() throws
IOException {
+ SpecifiedException expectedException = new
SpecifiedException("throw exception while multi thread upload states.");
+
+ CheckpointStreamFactory.CheckpointStateOutputStream
outputStream = new CheckpointStreamFactory.CheckpointStateOutputStream() {
+ @Nullable
+ @Override
+ public StreamStateHandle closeAndGetHandle() {
+ return new ByteStreamStateHandle("testHandle",
"testHandle".getBytes());
+ }
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public long getPos() {
+ return 0;
+ }
+
+ @Override
+ public void flush() {
+ }
+
+ @Override
+ public void sync() {
+ }
+
+ @Override
+ public void write(int b) throws IOException {
+ throw expectedException;
+ }
+ };
+ CheckpointStreamFactory checkpointStreamFactory =
(CheckpointedStateScope scope) -> outputStream;
+
+ File file =
temporaryFolder.newFile(String.valueOf(UUID.randomUUID()));
+ generateRandomFileContent(file.getPath(), 20);
+
+ Map<StateHandleID, Path> filePaths = new HashMap<>(1);
+ filePaths.put(new StateHandleID("mockHandleID"), new
Path(file.getPath()));
+ try {
+ RocksDbStateDataTransfer.uploadFilesToCheckpointFs(
+ filePaths,
+ 5,
+ checkpointStreamFactory,
+ new CloseableRegistry(),
+ new HashMap<>());
+ fail();
+ } catch (Exception e) {
+ assertEquals(expectedException, e);
+ }
+ }
+
+ /**
+ * Test that upload files with multi-thread correctly.
+ */
+ @Test
+ public void testMultiThreadUploadCorrectly() throws Exception {
+
+ File checkpointPrivateFolder =
temporaryFolder.newFolder("private");
+ Path checkpointPrivateDirectory = new
Path(checkpointPrivateFolder.getPath());
+
+ File checkpointSharedFolder =
temporaryFolder.newFolder("shared");
+ Path checkpointSharedDirectory = new
Path(checkpointSharedFolder.getPath());
+
+ FileSystem fileSystem =
checkpointPrivateDirectory.getFileSystem();
+ int fileStateSizeThreshold = 1024;
+ FsCheckpointStreamFactory checkpointStreamFactory =
+ new FsCheckpointStreamFactory(fileSystem,
checkpointPrivateDirectory, checkpointSharedDirectory, fileStateSizeThreshold);
+
+ String localFolder = "local";
+ temporaryFolder.newFolder(localFolder);
+
+ int sstFileCount = 6;
+ Map<StateHandleID, Path> sstFilePaths = new
HashMap<>(sstFileCount);
+ generateRandomSstFiles(localFolder, sstFileCount,
fileStateSizeThreshold, sstFilePaths);
+
+ int miscFileCount = 3;
+ Map<StateHandleID, Path> miscFilePaths = new
HashMap<>(miscFileCount);
+ ThreadLocalRandom random = ThreadLocalRandom.current();
+ File currentFile =
temporaryFolder.newFile(String.format("%s/CURRENT", localFolder));
+ generateRandomFileContent(currentFile.getPath(),
random.nextInt(fileStateSizeThreshold) + 1);
+ miscFilePaths.put(new StateHandleID("CURRENT"),
Path.fromLocalFile(currentFile));
+
+ File manifest =
temporaryFolder.newFile(String.format("%s/MANIFEST", localFolder));
+ generateRandomFileContent(manifest.getPath(),
random.nextInt(fileStateSizeThreshold) + 1);
+ miscFilePaths.put(new StateHandleID("MANIFEST"),
Path.fromLocalFile(manifest));
+
+ File options =
temporaryFolder.newFile(String.format("%s/OPTIONS", localFolder));
+ generateRandomFileContent(options.getPath(),
random.nextInt(fileStateSizeThreshold) + 1);
+ miscFilePaths.put(new StateHandleID("OPTIONS"),
Path.fromLocalFile(options));
+
+ Map<StateHandleID, StreamStateHandle> sstFiles = new
HashMap<>(sstFileCount);
+ Map<StateHandleID, StreamStateHandle> miscFiles = new
HashMap<>(miscFileCount);
+
+ RocksDbStateDataTransfer.uploadFilesToCheckpointFs(
+ sstFilePaths,
+ 5,
+ checkpointStreamFactory,
+ new CloseableRegistry(),
+ sstFiles);
+
+ RocksDbStateDataTransfer.uploadFilesToCheckpointFs(
+ miscFilePaths,
+ 5,
+ checkpointStreamFactory,
+ new CloseableRegistry(),
+ miscFiles);
+
+ for (Map.Entry<StateHandleID, Path> entry :
sstFilePaths.entrySet()) {
+ assertStateContentEqual(entry.getValue(),
sstFiles.get(entry.getKey()).openInputStream());
+ }
+
+ for (Map.Entry<StateHandleID, Path> entry :
miscFilePaths.entrySet()) {
+ assertStateContentEqual(entry.getValue(),
miscFiles.get(entry.getKey()).openInputStream());
+ }
+ }
+
+ private void generateRandomSstFiles(
+ String localFolder,
+ int sstFileCount,
+ int fileStateSizeThreshold,
+ Map<StateHandleID, Path> stateFilePaths) throws IOException {
Review comment:
also here, I would rather return `stateFilePaths`.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
