@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; int writeBufferSize = 4096; FsCheckpointStreamFactory checkpointStreamFactory = new FsCheckpointStreamFactory( fileSystem, checkpointPrivateDirectory, checkpointSharedDirectory, fileStateSizeThreshold, writeBufferSize); String localFolder = "local"; temporaryFolder.newFolder(localFolder); int sstFileCount = 6; Map<StateHandleID, Path> sstFilePaths = generateRandomSstFiles(localFolder, sstFileCount, fileStateSizeThreshold); try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) { Map<StateHandleID, StreamStateHandle> sstFiles = rocksDBStateUploader.uploadFilesToCheckpointFs(sstFilePaths, checkpointStreamFactory, new CloseableRegistry()); for (Map.Entry<StateHandleID, Path> entry : sstFilePaths.entrySet()) { assertStateContentEqual(entry.getValue(), sstFiles.get(entry.getKey()).openInputStream()); } } }
public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs( @Nonnull Map<StateHandleID, Path> files, CheckpointStreamFactory checkpointStreamFactory, CloseableRegistry closeableRegistry) throws Exception { Map<StateHandleID, StreamStateHandle> handles = new HashMap<>(); Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = createUploadFutures(files, checkpointStreamFactory, closeableRegistry); try { FutureUtils.waitForAll(futures.values()).get(); for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) { handles.put(entry.getKey(), entry.getValue().get()); } } catch (ExecutionException e) { Throwable throwable = ExceptionUtils.stripExecutionException(e); throwable = ExceptionUtils.stripException(throwable, RuntimeException.class); if (throwable instanceof IOException) { throw (IOException) throwable; } else { throw new FlinkRuntimeException("Failed to download data for state handles.", e); } } return handles; }
RocksDBStateUploader extends RocksDBStateDataTransfer { public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs( @Nonnull Map<StateHandleID, Path> files, CheckpointStreamFactory checkpointStreamFactory, CloseableRegistry closeableRegistry) throws Exception { Map<StateHandleID, StreamStateHandle> handles = new HashMap<>(); Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = createUploadFutures(files, checkpointStreamFactory, closeableRegistry); try { FutureUtils.waitForAll(futures.values()).get(); for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) { handles.put(entry.getKey(), entry.getValue().get()); } } catch (ExecutionException e) { Throwable throwable = ExceptionUtils.stripExecutionException(e); throwable = ExceptionUtils.stripException(throwable, RuntimeException.class); if (throwable instanceof IOException) { throw (IOException) throwable; } else { throw new FlinkRuntimeException("Failed to download data for state handles.", e); } } return handles; } }
RocksDBStateUploader extends RocksDBStateDataTransfer { public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs( @Nonnull Map<StateHandleID, Path> files, CheckpointStreamFactory checkpointStreamFactory, CloseableRegistry closeableRegistry) throws Exception { Map<StateHandleID, StreamStateHandle> handles = new HashMap<>(); Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = createUploadFutures(files, checkpointStreamFactory, closeableRegistry); try { FutureUtils.waitForAll(futures.values()).get(); for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) { handles.put(entry.getKey(), entry.getValue().get()); } } catch (ExecutionException e) { Throwable throwable = ExceptionUtils.stripExecutionException(e); throwable = ExceptionUtils.stripException(throwable, RuntimeException.class); if (throwable instanceof IOException) { throw (IOException) throwable; } else { throw new FlinkRuntimeException("Failed to download data for state handles.", e); } } return handles; } RocksDBStateUploader(int numberOfSnapshottingThreads); }
RocksDBStateUploader extends RocksDBStateDataTransfer { public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs( @Nonnull Map<StateHandleID, Path> files, CheckpointStreamFactory checkpointStreamFactory, CloseableRegistry closeableRegistry) throws Exception { Map<StateHandleID, StreamStateHandle> handles = new HashMap<>(); Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = createUploadFutures(files, checkpointStreamFactory, closeableRegistry); try { FutureUtils.waitForAll(futures.values()).get(); for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) { handles.put(entry.getKey(), entry.getValue().get()); } } catch (ExecutionException e) { Throwable throwable = ExceptionUtils.stripExecutionException(e); throwable = ExceptionUtils.stripException(throwable, RuntimeException.class); if (throwable instanceof IOException) { throw (IOException) throwable; } else { throw new FlinkRuntimeException("Failed to download data for state handles.", e); } } return handles; } RocksDBStateUploader(int numberOfSnapshottingThreads); Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs(
		@Nonnull Map<StateHandleID, Path> files,
		CheckpointStreamFactory checkpointStreamFactory,
		CloseableRegistry closeableRegistry); }
RocksDBStateUploader extends RocksDBStateDataTransfer { public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs( @Nonnull Map<StateHandleID, Path> files, CheckpointStreamFactory checkpointStreamFactory, CloseableRegistry closeableRegistry) throws Exception { Map<StateHandleID, StreamStateHandle> handles = new HashMap<>(); Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = createUploadFutures(files, checkpointStreamFactory, closeableRegistry); try { FutureUtils.waitForAll(futures.values()).get(); for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) { handles.put(entry.getKey(), entry.getValue().get()); } } catch (ExecutionException e) { Throwable throwable = ExceptionUtils.stripExecutionException(e); throwable = ExceptionUtils.stripException(throwable, RuntimeException.class); if (throwable instanceof IOException) { throw (IOException) throwable; } else { throw new FlinkRuntimeException("Failed to download data for state handles.", e); } } return handles; } RocksDBStateUploader(int numberOfSnapshottingThreads); Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs(
		@Nonnull Map<StateHandleID, Path> files,
		CheckpointStreamFactory checkpointStreamFactory,
		CloseableRegistry closeableRegistry);  }