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

/**
 * Checks that the state snapshot context is closed after a successful snapshot operation.
 */
@Test
public void testSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(0L, 0L));

	whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);

	StreamTask<Void, AbstractStreamOperator<Void>> containingTask = mock(StreamTask.class);
	when(containingTask.getCancelables()).thenReturn(closeableRegistry);

	AbstractStreamOperator<Void> operator = mock(AbstractStreamOperator.class);
	when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod();
	doReturn(containingTask).when(operator).getContainingTask();

	operator.snapshotState(
			checkpointId,
			timestamp,
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			new MemCheckpointStreamFactory(Integer.MAX_VALUE));

}
 

@Nonnull
@Override
public final RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	long checkpointId,
	long timestamp,
	@Nonnull CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	if (kvStateInformation.isEmpty()) {
		if (LOG.isDebugEnabled()) {
			LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
				timestamp);
		}
		return DoneFuture.of(SnapshotResult.empty());
	} else {
		return doSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
	}
}
 

private SupplierWithException<CheckpointStreamWithResultProvider, Exception> createCheckpointStreamSupplier(
	long checkpointId,
	CheckpointStreamFactory primaryStreamFactory,
	CheckpointOptions checkpointOptions) {

	return localRecoveryConfig.isLocalRecoveryEnabled() &&
		(CheckpointType.SAVEPOINT != checkpointOptions.getCheckpointType()) ?

		() -> CheckpointStreamWithResultProvider.createDuplicatingStream(
			checkpointId,
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory,
			localRecoveryConfig.getLocalStateDirectoryProvider()) :

		() -> CheckpointStreamWithResultProvider.createSimpleStream(
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory);
}
 

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

/**
 * Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be canceled and
 * is also stopped when the backend is closed through {@link #dispose()}. For each backend, this method must always
 * be called by the same thread.
 *
 * @param checkpointId  The Id of the checkpoint.
 * @param timestamp     The timestamp of the checkpoint.
 * @param streamFactory The factory that we can use for writing our state to streams.
 * @param checkpointOptions Options for how to perform this checkpoint.
 * @return Future to the state handle of the snapshot data.
 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
 */
@Nonnull
@Override
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	long startTime = System.currentTimeMillis();

	// flush everything into db before taking a snapshot
	writeBatchWrapper.flush();

	RocksDBSnapshotStrategyBase<K> chosenSnapshotStrategy =
		CheckpointType.SAVEPOINT == checkpointOptions.getCheckpointType() ?
			savepointSnapshotStrategy : checkpointSnapshotStrategy;

	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		chosenSnapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	chosenSnapshotStrategy.logSyncCompleted(streamFactory, startTime);

	return snapshotRunner;
}
 

@Nonnull
@Override
public final RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	long checkpointId,
	long timestamp,
	@Nonnull CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	if (kvStateInformation.isEmpty()) {
		if (LOG.isDebugEnabled()) {
			LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
				timestamp);
		}
		return DoneFuture.of(SnapshotResult.empty());
	} else {
		return doSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
	}
}
 

/**
 * Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be canceled and
 * is also stopped when the backend is closed through {@link #dispose()}. For each backend, this method must always
 * be called by the same thread.
 *
 * @param checkpointId  The Id of the checkpoint.
 * @param timestamp     The timestamp of the checkpoint.
 * @param streamFactory The factory that we can use for writing our state to streams.
 * @param checkpointOptions Options for how to perform this checkpoint.
 * @return Future to the state handle of the snapshot data.
 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
 */
@Nonnull
@Override
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	long startTime = System.currentTimeMillis();

	// flush everything into db before taking a snapshot
	writeBatchWrapper.flush();

	RocksDBSnapshotStrategyBase<K> chosenSnapshotStrategy =
			checkpointOptions.getCheckpointType().isSavepoint() ? savepointSnapshotStrategy : checkpointSnapshotStrategy;

	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		chosenSnapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	chosenSnapshotStrategy.logSyncCompleted(streamFactory, startTime);

	return snapshotRunner;
}
 

@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);
}
 

@Nonnull
@Override
@SuppressWarnings("unchecked")
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws IOException {

	long startTime = System.currentTimeMillis();

	final RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		snapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	snapshotStrategy.logSyncCompleted(streamFactory, startTime);
	return snapshotRunner;
}
 

private void checkpointState(
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		CheckpointMetrics checkpointMetrics) throws Exception {

	CheckpointStreamFactory storage = checkpointStorage.resolveCheckpointStorageLocation(
			checkpointMetaData.getCheckpointId(),
			checkpointOptions.getTargetLocation());

	CheckpointingOperation checkpointingOperation = new CheckpointingOperation(
		this,
		checkpointMetaData,
		checkpointOptions,
		storage,
		checkpointMetrics);

	checkpointingOperation.executeCheckpointing();
}
 

private void checkpointState(
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		CheckpointMetrics checkpointMetrics) throws Exception {

	CheckpointStreamFactory storage = checkpointStorage.resolveCheckpointStorageLocation(
			checkpointMetaData.getCheckpointId(),
			checkpointOptions.getTargetLocation());

	CheckpointingOperation checkpointingOperation = new CheckpointingOperation(
		this,
		checkpointMetaData,
		checkpointOptions,
		storage,
		checkpointMetrics);

	checkpointingOperation.executeCheckpointing();
}
 

@Nullable
@Override
public StreamStateHandle closeAndGetHandle() throws IOException {

	if (!closed.compareAndSet(false, true)) {
		throw new IOException("Stream was already closed!");
	}

	if (delegate instanceof CheckpointStreamFactory.CheckpointStateOutputStream) {
		StreamStateHandle streamStateHandle =
			((CheckpointStreamFactory.CheckpointStateOutputStream) delegate).closeAndGetHandle();
		unblockAll();
		return streamStateHandle;
	} else {
		unblockAll();
		throw new IOException("Delegate is not a CheckpointStateOutputStream!");
	}
}
 

/**
 * Tests that the underlying stream file is deleted upon calling close.
 */
@Test
public void testCleanupWhenClosingStream() throws IOException {

	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	stream.close();

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

@Nonnull
@Override
@SuppressWarnings("unchecked")
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws IOException {

	long startTime = System.currentTimeMillis();

	final RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		snapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	snapshotStrategy.logSyncCompleted(streamFactory, startTime);
	return snapshotRunner;
}
 

private SupplierWithException<CheckpointStreamWithResultProvider, Exception> createCheckpointStreamSupplier(
	long checkpointId,
	CheckpointStreamFactory primaryStreamFactory,
	CheckpointOptions checkpointOptions) {

	return localRecoveryConfig.isLocalRecoveryEnabled() && !checkpointOptions.getCheckpointType().isSavepoint() ?

		() -> CheckpointStreamWithResultProvider.createDuplicatingStream(
			checkpointId,
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory,
			localRecoveryConfig.getLocalStateDirectoryProvider()) :

		() -> CheckpointStreamWithResultProvider.createSimpleStream(
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory);
}
 

private static OperatorSnapshotFutures checkpointStreamOperator(
		StreamOperator<?> op,
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		CheckpointStreamFactory storageLocation,
		Supplier<Boolean> isCanceled) throws Exception {
	try {
		return op.snapshotState(
			checkpointMetaData.getCheckpointId(),
			checkpointMetaData.getTimestamp(),
			checkpointOptions,
			storageLocation);
	}
	catch (Exception ex) {
		if (!isCanceled.get()) {
			LOG.info(ex.getMessage(), ex);
		}
		throw ex;
	}
}
 

/**
 * Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be canceled and
 * is also stopped when the backend is closed through {@link #dispose()}. For each backend, this method must always
 * be called by the same thread.
 *
 * @param checkpointId  The Id of the checkpoint.
 * @param timestamp     The timestamp of the checkpoint.
 * @param streamFactory The factory that we can use for writing our state to streams.
 * @param checkpointOptions Options for how to perform this checkpoint.
 * @return Future to the state handle of the snapshot data.
 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
 */
@Nonnull
@Override
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	long startTime = System.currentTimeMillis();

	// flush everything into db before taking a snapshot
	writeBatchWrapper.flush();

	RocksDBSnapshotStrategyBase<K> chosenSnapshotStrategy =
			checkpointOptions.getCheckpointType().isSavepoint() ? savepointSnapshotStrategy : checkpointSnapshotStrategy;

	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		chosenSnapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	chosenSnapshotStrategy.logSyncCompleted(streamFactory, startTime);

	return snapshotRunner;
}
 

@Before
public void setUp() throws Exception {
	CloseableRegistry closableRegistry = new CloseableRegistry();
	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
	this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}
 

@Override
public CheckpointStreamFactory resolveCheckpointStorageLocation(
        long checkpointId,
        CheckpointStorageLocationReference reference)
        throws IOException
{
    if (reference.isDefaultReference()) {
        // default reference, construct the default location for that particular checkpoint
        final Path checkpointDir = createCheckpointDirectory(checkpointsDirectory, checkpointId);

        return new FsCheckpointStorageLocation(
                fileSystem,
                checkpointDir,
                sharedStateDirectory,
                taskOwnedStateDirectory,
                reference,
                fileSizeThreshold,
                writeBufferSize);
    }
    else {
        // location encoded in the reference
        final Path path = decodePathFromReference(reference);

        return new FsCheckpointStorageLocation(
                path.getFileSystem(),
                path,
                path,
                path,
                reference,
                fileSizeThreshold,
                writeBufferSize);
    }
}
 

@Override
public CheckpointStreamFactory resolveCheckpointStorageLocation(
		long checkpointId,
		CheckpointStorageLocationReference reference) {

	// no matter where the checkpoint goes, we always return the storage location that stores
	// state inline with the state handles.
	return new MemCheckpointStreamFactory(maxStateSize);
}
 

@Override
public CheckpointStreamFactory resolveCheckpointStorageLocation(
		long checkpointId,
		CheckpointStorageLocationReference reference) throws IOException {

	if (reference.isDefaultReference()) {
		// default reference, construct the default location for that particular checkpoint
		final Path checkpointDir = createCheckpointDirectory(checkpointsDirectory, checkpointId);

		return new FsCheckpointStorageLocation(
				fileSystem,
				checkpointDir,
				sharedStateDirectory,
				taskOwnedStateDirectory,
				reference,
				fileSizeThreshold,
				writeBufferSize);
	}
	else {
		// location encoded in the reference
		final Path path = decodePathFromReference(reference);

		return new FsCheckpointStorageLocation(
				path.getFileSystem(),
				path,
				path,
				path,
				reference,
				fileSizeThreshold,
				writeBufferSize);
	}
}
 

private CheckpointStreamFactory.CheckpointStateOutputStream createFailingCheckpointStateOutputStream(
	IOException failureException) {
	return 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 failureException;
		}
	};
}
 

ChannelStateCheckpointWriter(
		CheckpointStartRequest startCheckpointItem,
		CheckpointStreamFactory streamFactory,
		ChannelStateSerializer serializer,
		RunnableWithException onComplete) throws Exception {
	this(
		startCheckpointItem.getCheckpointId(),
		startCheckpointItem.getTargetResult(),
		streamFactory.createCheckpointStateOutputStream(EXCLUSIVE),
		serializer,
		onComplete);
}
 

private RocksDBIncrementalSnapshotOperation(
	long checkpointId,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull SnapshotDirectory localBackupDirectory,
	@Nullable Set<StateHandleID> baseSstFiles,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {

	this.checkpointStreamFactory = checkpointStreamFactory;
	this.baseSstFiles = baseSstFiles;
	this.checkpointId = checkpointId;
	this.localBackupDirectory = localBackupDirectory;
	this.stateMetaInfoSnapshots = stateMetaInfoSnapshots;
}
 

private CheckpointStreamFactory createCheckpointStreamFactory() {
	try {
		return stateBackend
			.createCheckpointStorage(new JobID())
			.resolveCheckpointStorageLocation(2L, checkpointOptions.getTargetLocation());
	} catch (IOException e) {
		throw new RuntimeException("unexpected");
	}
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	OperatorStateBackend operatorStateBackend = mock(OperatorStateBackend.class);
	when(operatorStateBackend.snapshot(anyLong(), anyLong(), any(CheckpointStreamFactory.class), any(CheckpointOptions.class)))
		.thenReturn(new FutureTask<>(new BlockingCallable()));

	return operatorStateBackend;
}
 

private RocksDBIncrementalSnapshotOperation(
	long checkpointId,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull SnapshotDirectory localBackupDirectory,
	@Nullable Set<StateHandleID> baseSstFiles,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {

	this.checkpointStreamFactory = checkpointStreamFactory;
	this.baseSstFiles = baseSstFiles;
	this.checkpointId = checkpointId;
	this.localBackupDirectory = localBackupDirectory;
	this.stateMetaInfoSnapshots = stateMetaInfoSnapshots;
}
 

@Override
public RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot(
	long checkpointId, long timestamp, @Nonnull CheckpointStreamFactory streamFactory, @Nonnull CheckpointOptions checkpointOptions) {
	if (checkpointId == DECLINE_CHECKPOINT_ID) {
		return ExceptionallyDoneFuture.of(new ExpectedTestException());
	} else {
		return DoneFuture.of(SnapshotResult.empty());
	}
}
 

@Override
public final OperatorSnapshotFutures snapshotState(long checkpointId, long timestamp, CheckpointOptions checkpointOptions,
		CheckpointStreamFactory factory) throws Exception {
	return stateHandler.snapshotState(
		this,
		Optional.ofNullable(timeServiceManager),
		getOperatorName(),
		checkpointId,
		timestamp,
		checkpointOptions,
		factory);
}
 

@Nonnull
@Override
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	long checkpointId,
	long timestamp,
	@Nonnull CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) {
	return new FutureTask<>(() ->
		SnapshotResult.of(new MockKeyedStateHandle<>(copy(stateValues, stateSnapshotFilters))));
}