org.apache.flink.runtime.state.StateHandleID Java Examples

/**
 * 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 = createFailingCheckpointStateOutputStream(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"), file.toPath());
	try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

private void createUploadFilePaths(
	FileStatus[] fileStatuses,
	Map<StateHandleID, StreamStateHandle> sstFiles,
	Map<StateHandleID, Path> sstFilePaths,
	Map<StateHandleID, Path> miscFilePaths) {
	for (FileStatus fileStatus : fileStatuses) {
		final Path filePath = fileStatus.getPath();
		final String fileName = filePath.getName();
		final StateHandleID stateHandleID = new StateHandleID(fileName);

		if (fileName.endsWith(SST_FILE_SUFFIX)) {
			final boolean existsAlready = baseSstFiles != null && baseSstFiles.contains(stateHandleID);

			if (existsAlready) {
				// we introduce a placeholder state handle, that is replaced with the
				// original from the shared state registry (created from a previous checkpoint)
				sstFiles.put(stateHandleID, new PlaceholderStreamStateHandle());
			} else {
				sstFilePaths.put(stateHandleID, filePath);
			}
		} else {
			miscFilePaths.put(stateHandleID, filePath);
		}
	}
}
 

private List<Runnable> createDownloadRunnables(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) {
	List<Runnable> runnables = new ArrayList<>(stateHandleMap.size());
	for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
		StateHandleID stateHandleID = entry.getKey();
		StreamStateHandle remoteFileHandle = entry.getValue();

		Path path = new Path(restoreInstancePath, stateHandleID.toString());

		runnables.add(ThrowingRunnable.unchecked(
			() -> downloadDataForStateHandle(path, remoteFileHandle, closeableRegistry)));
	}
	return runnables;
}
 

private void uploadSstFiles(
	@Nonnull Map<StateHandleID, StreamStateHandle> sstFiles,
	@Nonnull Map<StateHandleID, StreamStateHandle> miscFiles) throws Exception {

	// write state data
	Preconditions.checkState(localBackupDirectory.exists());

	Map<StateHandleID, Path> sstFilePaths = new HashMap<>();
	Map<StateHandleID, Path> miscFilePaths = new HashMap<>();

	FileStatus[] fileStatuses = localBackupDirectory.listStatus();
	if (fileStatuses != null) {
		createUploadFilePaths(fileStatuses, sstFiles, sstFilePaths, miscFilePaths);

		sstFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			sstFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
		miscFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			miscFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
	}
}
 

private Set<StateHandleID> snapshotMetaData(
	long checkpointId,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {

	final long lastCompletedCheckpoint;
	final Set<StateHandleID> baseSstFiles;

	// use the last completed checkpoint as the comparison base.
	synchronized (materializedSstFiles) {
		lastCompletedCheckpoint = lastCompletedCheckpointId;
		baseSstFiles = materializedSstFiles.get(lastCompletedCheckpoint);
	}
	LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
		"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);

	// snapshot meta data to save
	for (Map.Entry<String, RocksDbKvStateInfo> stateMetaInfoEntry : kvStateInformation.entrySet()) {
		stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().metaInfo.snapshot());
	}
	return baseSstFiles;
}
 

@Nonnull
@Override
protected RunnableFuture<SnapshotResult<KeyedStateHandle>> doSnapshot(
	long checkpointId,
	long checkpointTimestamp,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	final SnapshotDirectory snapshotDirectory = prepareLocalSnapshotDirectory(checkpointId);
	LOG.trace("Local RocksDB checkpoint goes to backup path {}.", snapshotDirectory);

	final List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = new ArrayList<>(kvStateInformation.size());
	final Set<StateHandleID> baseSstFiles = snapshotMetaData(checkpointId, stateMetaInfoSnapshots);

	takeDBNativeCheckpoint(snapshotDirectory);

	final RocksDBIncrementalSnapshotOperation snapshotOperation =
		new RocksDBIncrementalSnapshotOperation(
			checkpointId,
			checkpointStreamFactory,
			snapshotDirectory,
			baseSstFiles,
			stateMetaInfoSnapshots);

	return snapshotOperation.toAsyncSnapshotFutureTask(cancelStreamRegistry);
}
 

/**
 * 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 = createFailingCheckpointStateOutputStream(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 (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

/**
 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
 * {@link StateHandleID}.
 */
private void downloadDataForAllStateHandles(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) throws Exception {

	try {
		List<Runnable> runnables = createDownloadRunnables(stateHandleMap, restoreInstancePath, closeableRegistry);
		List<CompletableFuture<Void>> futures = new ArrayList<>(runnables.size());
		for (Runnable runnable : runnables) {
			futures.add(CompletableFuture.runAsync(runnable, executorService));
		}
		FutureUtils.waitForAll(futures).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);
		}
	}
}
 

private List<Runnable> createDownloadRunnables(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) {
	List<Runnable> runnables = new ArrayList<>(stateHandleMap.size());
	for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
		StateHandleID stateHandleID = entry.getKey();
		StreamStateHandle remoteFileHandle = entry.getValue();

		Path path = new Path(restoreInstancePath, stateHandleID.toString());

		runnables.add(ThrowingRunnable.unchecked(
			() -> downloadDataForStateHandle(path, remoteFileHandle, closeableRegistry)));
	}
	return runnables;
}
 

@Nonnull
@Override
protected RunnableFuture<SnapshotResult<KeyedStateHandle>> doSnapshot(
	long checkpointId,
	long checkpointTimestamp,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	final SnapshotDirectory snapshotDirectory = prepareLocalSnapshotDirectory(checkpointId);
	LOG.trace("Local RocksDB checkpoint goes to backup path {}.", snapshotDirectory);

	final List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = new ArrayList<>(kvStateInformation.size());
	final Set<StateHandleID> baseSstFiles = snapshotMetaData(checkpointId, stateMetaInfoSnapshots);

	takeDBNativeCheckpoint(snapshotDirectory);

	final RocksDBIncrementalSnapshotOperation snapshotOperation =
		new RocksDBIncrementalSnapshotOperation(
			checkpointId,
			checkpointStreamFactory,
			snapshotDirectory,
			baseSstFiles,
			stateMetaInfoSnapshots);

	return snapshotOperation.toAsyncSnapshotFutureTask(cancelStreamRegistry);
}
 

private Set<StateHandleID> snapshotMetaData(
	long checkpointId,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {

	final long lastCompletedCheckpoint;
	final Set<StateHandleID> baseSstFiles;

	// use the last completed checkpoint as the comparison base.
	synchronized (materializedSstFiles) {
		lastCompletedCheckpoint = lastCompletedCheckpointId;
		baseSstFiles = materializedSstFiles.get(lastCompletedCheckpoint);
	}
	LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
		"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);

	// snapshot meta data to save
	for (Map.Entry<String, RocksDbKvStateInfo> stateMetaInfoEntry : kvStateInformation.entrySet()) {
		stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().metaInfo.snapshot());
	}
	return baseSstFiles;
}
 

private void createUploadFilePaths(
	FileStatus[] fileStatuses,
	Map<StateHandleID, StreamStateHandle> sstFiles,
	Map<StateHandleID, Path> sstFilePaths,
	Map<StateHandleID, Path> miscFilePaths) {
	for (FileStatus fileStatus : fileStatuses) {
		final Path filePath = fileStatus.getPath();
		final String fileName = filePath.getName();
		final StateHandleID stateHandleID = new StateHandleID(fileName);

		if (fileName.endsWith(SST_FILE_SUFFIX)) {
			final boolean existsAlready = baseSstFiles != null && baseSstFiles.contains(stateHandleID);

			if (existsAlready) {
				// we introduce a placeholder state handle, that is replaced with the
				// original from the shared state registry (created from a previous checkpoint)
				sstFiles.put(stateHandleID, new PlaceholderStreamStateHandle());
			} else {
				sstFilePaths.put(stateHandleID, filePath);
			}
		} else {
			miscFilePaths.put(stateHandleID, filePath);
		}
	}
}
 

/**
 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
 * {@link StateHandleID}.
 */
private void downloadDataForAllStateHandles(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) throws Exception {

	try {
		List<Runnable> runnables = createDownloadRunnables(stateHandleMap, restoreInstancePath, closeableRegistry);
		List<CompletableFuture<Void>> futures = new ArrayList<>(runnables.size());
		for (Runnable runnable : runnables) {
			futures.add(CompletableFuture.runAsync(runnable, executorService));
		}
		FutureUtils.waitForAll(futures).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);
		}
	}
}
 

private void createUploadFilePaths(
	Path[] files,
	Map<StateHandleID, StreamStateHandle> sstFiles,
	Map<StateHandleID, Path> sstFilePaths,
	Map<StateHandleID, Path> miscFilePaths) {
	for (Path filePath : files) {
		final String fileName = filePath.getFileName().toString();
		final StateHandleID stateHandleID = new StateHandleID(fileName);

		if (fileName.endsWith(SST_FILE_SUFFIX)) {
			final boolean existsAlready = baseSstFiles != null && baseSstFiles.contains(stateHandleID);

			if (existsAlready) {
				// we introduce a placeholder state handle, that is replaced with the
				// original from the shared state registry (created from a previous checkpoint)
				sstFiles.put(stateHandleID, new PlaceholderStreamStateHandle());
			} else {
				sstFilePaths.put(stateHandleID, filePath);
			}
		} else {
			miscFilePaths.put(stateHandleID, filePath);
		}
	}
}
 

private void uploadSstFiles(
	@Nonnull Map<StateHandleID, StreamStateHandle> sstFiles,
	@Nonnull Map<StateHandleID, StreamStateHandle> miscFiles) throws Exception {

	// write state data
	Preconditions.checkState(localBackupDirectory.exists());

	Map<StateHandleID, Path> sstFilePaths = new HashMap<>();
	Map<StateHandleID, Path> miscFilePaths = new HashMap<>();

	Path[] files = localBackupDirectory.listDirectory();
	if (files != null) {
		createUploadFilePaths(files, sstFiles, sstFilePaths, miscFilePaths);

		sstFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			sstFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
		miscFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			miscFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
	}
}
 

private Set<StateHandleID> snapshotMetaData(
	long checkpointId,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {

	final long lastCompletedCheckpoint;
	final Set<StateHandleID> baseSstFiles;

	// use the last completed checkpoint as the comparison base.
	synchronized (materializedSstFiles) {
		lastCompletedCheckpoint = lastCompletedCheckpointId;
		baseSstFiles = materializedSstFiles.get(lastCompletedCheckpoint);
	}
	LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
		"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);

	// snapshot meta data to save
	for (Map.Entry<String, RocksDbKvStateInfo> stateMetaInfoEntry : kvStateInformation.entrySet()) {
		stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().metaInfo.snapshot());
	}
	return baseSstFiles;
}
 

@Nonnull
@Override
protected RunnableFuture<SnapshotResult<KeyedStateHandle>> doSnapshot(
	long checkpointId,
	long checkpointTimestamp,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	final SnapshotDirectory snapshotDirectory = prepareLocalSnapshotDirectory(checkpointId);
	LOG.trace("Local RocksDB checkpoint goes to backup path {}.", snapshotDirectory);

	final List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = new ArrayList<>(kvStateInformation.size());
	final Set<StateHandleID> baseSstFiles = snapshotMetaData(checkpointId, stateMetaInfoSnapshots);

	takeDBNativeCheckpoint(snapshotDirectory);

	final RocksDBIncrementalSnapshotOperation snapshotOperation =
		new RocksDBIncrementalSnapshotOperation(
			checkpointId,
			checkpointStreamFactory,
			snapshotDirectory,
			baseSstFiles,
			stateMetaInfoSnapshots);

	return snapshotOperation.toAsyncSnapshotFutureTask(cancelStreamRegistry);
}
 

private List<Runnable> createDownloadRunnables(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) {
	List<Runnable> runnables = new ArrayList<>(stateHandleMap.size());
	for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
		StateHandleID stateHandleID = entry.getKey();
		StreamStateHandle remoteFileHandle = entry.getValue();

		Path path = restoreInstancePath.resolve(stateHandleID.toString());

		runnables.add(ThrowingRunnable.unchecked(
			() -> downloadDataForStateHandle(path, remoteFileHandle, closeableRegistry)));
	}
	return runnables;
}
 

/**
 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
 * {@link StateHandleID}.
 */
private void downloadDataForAllStateHandles(
	Map<StateHandleID, StreamStateHandle> stateHandleMap,
	Path restoreInstancePath,
	CloseableRegistry closeableRegistry) throws Exception {

	try {
		List<Runnable> runnables = createDownloadRunnables(stateHandleMap, restoreInstancePath, closeableRegistry);
		List<CompletableFuture<Void>> futures = new ArrayList<>(runnables.size());
		for (Runnable runnable : runnables) {
			futures.add(CompletableFuture.runAsync(runnable, executorService));
		}
		FutureUtils.waitForAll(futures).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);
		}
	}
}
 

/**
 * 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 = createFailingCheckpointStateOutputStream(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 (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

private void uploadSstFiles(
	@Nonnull Map<StateHandleID, StreamStateHandle> sstFiles,
	@Nonnull Map<StateHandleID, StreamStateHandle> miscFiles) throws Exception {

	// write state data
	Preconditions.checkState(localBackupDirectory.exists());

	Map<StateHandleID, Path> sstFilePaths = new HashMap<>();
	Map<StateHandleID, Path> miscFilePaths = new HashMap<>();

	FileStatus[] fileStatuses = localBackupDirectory.listStatus();
	if (fileStatuses != null) {
		createUploadFilePaths(fileStatuses, sstFiles, sstFilePaths, miscFilePaths);

		sstFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			sstFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
		miscFiles.putAll(stateUploader.uploadFilesToCheckpointFs(
			miscFilePaths,
			checkpointStreamFactory,
			snapshotCloseableRegistry));
	}
}
 

public RocksIncrementalSnapshotStrategy(
	@Nonnull RocksDB db,
	@Nonnull ResourceGuard rocksDBResourceGuard,
	@Nonnull TypeSerializer<K> keySerializer,
	@Nonnull LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull CloseableRegistry cancelStreamRegistry,
	@Nonnull File instanceBasePath,
	@Nonnull UUID backendUID,
	@Nonnull SortedMap<Long, Set<StateHandleID>> materializedSstFiles,
	long lastCompletedCheckpointId,
	int numberOfTransferingThreads) {

	super(
		DESCRIPTION,
		db,
		rocksDBResourceGuard,
		keySerializer,
		kvStateInformation,
		keyGroupRange,
		keyGroupPrefixBytes,
		localRecoveryConfig,
		cancelStreamRegistry);

	this.instanceBasePath = instanceBasePath;
	this.backendUID = backendUID;
	this.materializedSstFiles = materializedSstFiles;
	this.lastCompletedCheckpointId = lastCompletedCheckpointId;
	this.stateUploader = new RocksDBStateUploader(numberOfTransferingThreads);
	this.localDirectoryName = backendUID.toString().replaceAll("[\\-]", "");
}
 

private Map<StateHandleID, CompletableFuture<StreamStateHandle>> createUploadFutures(
	Map<StateHandleID, Path> files,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) {
	Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = new HashMap<>(files.size());

	for (Map.Entry<StateHandleID, Path> entry : files.entrySet()) {
		final Supplier<StreamStateHandle> supplier =
			CheckedSupplier.unchecked(() -> uploadLocalFileToCheckpointFs(entry.getValue(), checkpointStreamFactory, closeableRegistry));
		futures.put(entry.getKey(), CompletableFuture.supplyAsync(supplier, executorService));
	}

	return futures;
}
 

/**
 * Transfer all state data to the target directory using specified number of threads.
 *
 * @param restoreStateHandle Handles used to retrieve the state data.
 * @param dest The target directory which the state data will be stored.
 *
 * @throws Exception Thrown if can not transfer all the state data.
 */
public void transferAllStateDataToDirectory(
	IncrementalRemoteKeyedStateHandle restoreStateHandle,
	Path dest,
	CloseableRegistry closeableRegistry) throws Exception {

	final Map<StateHandleID, StreamStateHandle> sstFiles =
		restoreStateHandle.getSharedState();
	final Map<StateHandleID, StreamStateHandle> miscFiles =
		restoreStateHandle.getPrivateState();

	downloadDataForAllStateHandles(sstFiles, dest, closeableRegistry);
	downloadDataForAllStateHandles(miscFiles, dest, closeableRegistry);
}
 

public RocksDBRestoreResult(
	RocksDB db,
	ColumnFamilyHandle defaultColumnFamilyHandle,
	RocksDBNativeMetricMonitor nativeMetricMonitor,
	long lastCompletedCheckpointId,
	UUID backendUID,
	SortedMap<Long, Set<StateHandleID>> restoredSstFiles) {
	this.db = db;
	this.defaultColumnFamilyHandle = defaultColumnFamilyHandle;
	this.nativeMetricMonitor = nativeMetricMonitor;
	this.lastCompletedCheckpointId = lastCompletedCheckpointId;
	this.backendUID = backendUID;
	this.restoredSstFiles = restoredSstFiles;
}
 

private static Map<StateHandleID, StreamStateHandle> deserializeStreamStateHandleMap(
	DataInputStream dis) throws IOException {

	final int size = dis.readInt();
	Map<StateHandleID, StreamStateHandle> result = new HashMap<>(size);

	for (int i = 0; i < size; ++i) {
		StateHandleID stateHandleID = new StateHandleID(dis.readUTF());
		StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
		result.put(stateHandleID, stateHandle);
	}

	return result;
}
 

/**
 * 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;
	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());
		}
	}
}
 

private Map<StateHandleID, Path> generateRandomSstFiles(
	String localFolder,
	int sstFileCount,
	int fileStateSizeThreshold) throws IOException {
	ThreadLocalRandom random = ThreadLocalRandom.current();

	Map<StateHandleID, Path> sstFilePaths = new HashMap<>(sstFileCount);
	for (int i = 0; i < sstFileCount; ++i) {
		File file = temporaryFolder.newFile(String.format("%s/%d.sst", localFolder, i));
		generateRandomFileContent(file.getPath(), random.nextInt(1_000_000) + fileStateSizeThreshold);
		sstFilePaths.put(new StateHandleID(String.valueOf(i)), Path.fromLocalFile(file));
	}
	return sstFilePaths;
}
 

private void transferAllDataFromStateHandles(
		Map<StateHandleID, StreamStateHandle> stateHandleMap,
		Path restoreInstancePath) throws IOException {

	for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
		StateHandleID stateHandleID = entry.getKey();
		StreamStateHandle remoteFileHandle = entry.getValue();
		copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
	}
}
 

public static Map<StateHandleID, StreamStateHandle> createRandomStateHandleMap(Random rnd) {
	final int size = rnd.nextInt(4);
	Map<StateHandleID, StreamStateHandle> result = new HashMap<>(size);
	for (int i = 0; i < size; ++i) {
		StateHandleID randomId = new StateHandleID(createRandomUUID(rnd).toString());
		StreamStateHandle stateHandle = createDummyStreamStateHandle(rnd);
		result.put(randomId, stateHandle);
	}

	return result;
}