org.apache.flink.streaming.api.operators.OperatorSnapshotFutures Java Examples

private void cleanup() throws Exception {
	LOG.debug(
		"Cleanup AsyncCheckpointRunnable for checkpoint {} of {}.",
		checkpointMetaData.getCheckpointId(),
		owner.getName());

	Exception exception = null;

	// clean up ongoing operator snapshot results and non partitioned state handles
	for (OperatorSnapshotFutures operatorSnapshotResult : operatorSnapshotsInProgress.values()) {
		if (operatorSnapshotResult != null) {
			try {
				operatorSnapshotResult.cancel();
			} catch (Exception cancelException) {
				exception = ExceptionUtils.firstOrSuppressed(cancelException, exception);
			}
		}
	}

	if (null != exception) {
		throw exception;
	}
}
 

private void cleanup(
		Map<OperatorID, OperatorSnapshotFutures> operatorSnapshotsInProgress,
		CheckpointMetaData metadata,
		CheckpointMetrics metrics,
		Exception ex) {

	channelStateWriter.abort(metadata.getCheckpointId(), ex, true);
	for (OperatorSnapshotFutures operatorSnapshotResult : operatorSnapshotsInProgress.values()) {
		if (operatorSnapshotResult != null) {
			try {
				operatorSnapshotResult.cancel();
			} catch (Exception e) {
				LOG.warn("Could not properly cancel an operator snapshot result.", e);
			}
		}
	}

	if (LOG.isDebugEnabled()) {
		LOG.debug(
			"{} - did NOT finish synchronous part of checkpoint {}. Alignment duration: {} ms, snapshot duration {} ms",
			taskName, metadata.getCheckpointId(),
			metrics.getAlignmentDurationNanos() / 1_000_000,
			metrics.getSyncDurationMillis());
	}
}
 

private void cleanup() throws Exception {
	LOG.debug(
		"Cleanup AsyncCheckpointRunnable for checkpoint {} of {}.",
		checkpointMetaData.getCheckpointId(),
		owner.getName());

	Exception exception = null;

	// clean up ongoing operator snapshot results and non partitioned state handles
	for (OperatorSnapshotFutures operatorSnapshotResult : operatorSnapshotsInProgress.values()) {
		if (operatorSnapshotResult != null) {
			try {
				operatorSnapshotResult.cancel();
			} catch (Exception cancelException) {
				exception = ExceptionUtils.firstOrSuppressed(cancelException, exception);
			}
		}
	}

	if (null != exception) {
		throw exception;
	}
}
 

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

AsyncCheckpointRunnable(
		Map<OperatorID, OperatorSnapshotFutures> operatorSnapshotsInProgress,
		CheckpointMetaData checkpointMetaData,
		CheckpointMetrics checkpointMetrics,
		long asyncStartNanos,
		String taskName,
		Consumer<AsyncCheckpointRunnable> register,
		Consumer<AsyncCheckpointRunnable> unregister,
		Environment taskEnvironment,
		AsyncExceptionHandler asyncExceptionHandler) {

	this.operatorSnapshotsInProgress = checkNotNull(operatorSnapshotsInProgress);
	this.checkpointMetaData = checkNotNull(checkpointMetaData);
	this.checkpointMetrics = checkNotNull(checkpointMetrics);
	this.asyncStartNanos = asyncStartNanos;
	this.taskName = checkNotNull(taskName);
	this.registerConsumer = register;
	this.unregisterConsumer = unregister;
	this.taskEnvironment = checkNotNull(taskEnvironment);
	this.asyncExceptionHandler = checkNotNull(asyncExceptionHandler);
}
 

private void cleanup() throws Exception {
	LOG.debug(
		"Cleanup AsyncCheckpointRunnable for checkpoint {} of {}.",
		checkpointMetaData.getCheckpointId(),
		taskName);

	Exception exception = null;

	// clean up ongoing operator snapshot results and non partitioned state handles
	for (OperatorSnapshotFutures operatorSnapshotResult : operatorSnapshotsInProgress.values()) {
		if (operatorSnapshotResult != null) {
			try {
				operatorSnapshotResult.cancel();
			} catch (Exception cancelException) {
				exception = ExceptionUtils.firstOrSuppressed(cancelException, exception);
			}
		}
	}

	if (null != exception) {
		throw exception;
	}
}
 

@Test
public void testNotifyCheckpointAbortedDuringAsyncPhase() throws Exception {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();
	SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder()
		.setEnvironment(mockEnvironment)
		.setExecutor(Executors.newSingleThreadExecutor())
		.setUnalignedCheckpointEnabled(true)
		.build();

	final BlockingRunnableFuture rawKeyedStateHandleFuture = new BlockingRunnableFuture();
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		DoneFuture.of(SnapshotResult.empty()),
		rawKeyedStateHandleFuture,
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));

	final OperatorChain<String, AbstractStreamOperator<String>> operatorChain = operatorChain(new CheckpointOperator(operatorSnapshotResult));

	long checkpointId = 42L;
	subtaskCheckpointCoordinator.getChannelStateWriter().start(checkpointId, CheckpointOptions.forCheckpointWithDefaultLocation());
	subtaskCheckpointCoordinator.checkpointState(
		new CheckpointMetaData(checkpointId, System.currentTimeMillis()),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		new CheckpointMetrics(),
		operatorChain,
		() -> true);
	rawKeyedStateHandleFuture.awaitRun();
	assertEquals(1, subtaskCheckpointCoordinator.getAsyncCheckpointRunnableSize());
	assertFalse(rawKeyedStateHandleFuture.isCancelled());

	subtaskCheckpointCoordinator.notifyCheckpointAborted(checkpointId, operatorChain, () -> true);
	assertTrue(rawKeyedStateHandleFuture.isCancelled());
	assertEquals(0, subtaskCheckpointCoordinator.getAsyncCheckpointRunnableSize());
}
 

@Test
public void testNotifyCheckpointAbortedBeforeAsyncPhase() throws Exception {
	TestTaskStateManager stateManager = new TestTaskStateManager();
	MockEnvironment mockEnvironment = MockEnvironment.builder().setTaskStateManager(stateManager).build();
	SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder()
		.setEnvironment(mockEnvironment)
		.setUnalignedCheckpointEnabled(true)
		.build();

	CheckpointOperator checkpointOperator = new CheckpointOperator(new OperatorSnapshotFutures());

	final OperatorChain<String, AbstractStreamOperator<String>> operatorChain = operatorChain(checkpointOperator);

	long checkpointId = 42L;
	// notify checkpoint aborted before execution.
	subtaskCheckpointCoordinator.notifyCheckpointAborted(checkpointId, operatorChain, () -> true);
	assertEquals(1, subtaskCheckpointCoordinator.getAbortedCheckpointSize());

	subtaskCheckpointCoordinator.getChannelStateWriter().start(checkpointId, CheckpointOptions.forCheckpointWithDefaultLocation());
	subtaskCheckpointCoordinator.checkpointState(
		new CheckpointMetaData(checkpointId, System.currentTimeMillis()),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		new CheckpointMetrics(),
		operatorChain,
		() -> true);
	assertFalse(checkpointOperator.isCheckpointed());
	assertEquals(-1, stateManager.getReportedCheckpointId());
	assertEquals(0, subtaskCheckpointCoordinator.getAbortedCheckpointSize());
	assertEquals(0, subtaskCheckpointCoordinator.getAsyncCheckpointRunnableSize());
}
 

@Test
public void testDecliningCheckpointStreamOperator() throws Exception {
	DeclineDummyEnvironment declineDummyEnvironment = new DeclineDummyEnvironment();

	// mock the returned snapshots
	OperatorSnapshotFutures operatorSnapshotResult1 = mock(OperatorSnapshotFutures.class);
	OperatorSnapshotFutures operatorSnapshotResult2 = mock(OperatorSnapshotFutures.class);

	final Exception testException = new ExpectedTestException();

	RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
		declineDummyEnvironment,
		operatorChain(
			streamOperatorWithSnapshotException(testException),
			streamOperatorWithSnapshot(operatorSnapshotResult1),
			streamOperatorWithSnapshot(operatorSnapshotResult2)
			)));
	MockStreamTask streamTask = task.streamTask;

	waitTaskIsRunning(streamTask, task.invocationFuture);

	streamTask.triggerCheckpointAsync(
		new CheckpointMetaData(42L, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	try {
		task.waitForTaskCompletion(false);
	}
	catch (Exception ex) {
		if (!ExceptionUtils.findThrowable(ex, ExpectedTestException.class).isPresent()) {
			throw ex;
		}
	}

	verify(operatorSnapshotResult1).cancel();
	verify(operatorSnapshotResult2).cancel();
}
 

/**
 * Tests that uncaught exceptions in the async part of a checkpoint operation are forwarded
 * to the uncaught exception handler. See <a href="https://issues.apache.org/jira/browse/FLINK-12889">FLINK-12889</a>.
 */
@Test
public void testUncaughtExceptionInAsynchronousCheckpointingOperation() throws Exception {
	final RuntimeException failingCause = new RuntimeException("Test exception");
	FailingDummyEnvironment failingDummyEnvironment = new FailingDummyEnvironment(failingCause);

	// mock the returned snapshots
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		ExceptionallyDoneFuture.of(failingCause),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));

	final TestingUncaughtExceptionHandler uncaughtExceptionHandler = new TestingUncaughtExceptionHandler();

	RunningTask<MockStreamTask> task = runTask(() -> new MockStreamTask(
		failingDummyEnvironment,
		operatorChain(streamOperatorWithSnapshot(operatorSnapshotResult)),
		uncaughtExceptionHandler));
	MockStreamTask streamTask = task.streamTask;

	waitTaskIsRunning(streamTask, task.invocationFuture);

	streamTask.triggerCheckpointAsync(
		new CheckpointMetaData(42L, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	final Throwable uncaughtException = uncaughtExceptionHandler.waitForUncaughtException();
	assertThat(uncaughtException, is(failingCause));

	streamTask.finishInput();
	task.waitForTaskCompletion(false);
}
 

private static <T> OneInputStreamOperator<T, T> streamOperatorWithSnapshot(OperatorSnapshotFutures operatorSnapshotResult) throws Exception {
	@SuppressWarnings("unchecked")
	OneInputStreamOperator<T, T> operator = mock(OneInputStreamOperator.class);
	when(operator.getOperatorID()).thenReturn(new OperatorID());

	when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class)))
		.thenReturn(operatorSnapshotResult);

	return operator;
}
 

private OperatorSnapshotFutures buildOperatorSnapshotFutures(
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		OperatorChain<?, ?> operatorChain,
		StreamOperator<?> op,
		Supplier<Boolean> isCanceled,
		ChannelStateWriteResult channelStateWriteResult,
		CheckpointStreamFactory storage) throws Exception {
	OperatorSnapshotFutures snapshotInProgress = checkpointStreamOperator(
		op,
		checkpointMetaData,
		checkpointOptions,
		storage,
		isCanceled);
	if (op == operatorChain.getHeadOperator()) {
		snapshotInProgress.setInputChannelStateFuture(
			channelStateWriteResult
				.getInputChannelStateHandles()
				.thenApply(StateObjectCollection::new)
				.thenApply(SnapshotResult::of));
	}
	if (op == operatorChain.getTailOperator()) {
		snapshotInProgress.setResultSubpartitionStateFuture(
			channelStateWriteResult
				.getResultSubpartitionStateHandles()
				.thenApply(StateObjectCollection::new)
				.thenApply(SnapshotResult::of));
	}
	return snapshotInProgress;
}
 

AsyncCheckpointRunnable(
	StreamTask<?, ?> owner,
	Map<OperatorID, OperatorSnapshotFutures> operatorSnapshotsInProgress,
	CheckpointMetaData checkpointMetaData,
	CheckpointMetrics checkpointMetrics,
	long asyncStartNanos) {

	this.owner = Preconditions.checkNotNull(owner);
	this.operatorSnapshotsInProgress = Preconditions.checkNotNull(operatorSnapshotsInProgress);
	this.checkpointMetaData = Preconditions.checkNotNull(checkpointMetaData);
	this.checkpointMetrics = Preconditions.checkNotNull(checkpointMetrics);
	this.asyncStartNanos = asyncStartNanos;
}
 

private void finishAndReportAsync(Map<OperatorID, OperatorSnapshotFutures> snapshotFutures, CheckpointMetaData metadata, CheckpointMetrics metrics, CheckpointOptions options) {
	// we are transferring ownership over snapshotInProgressList for cleanup to the thread, active on submit
	executorService.execute(new AsyncCheckpointRunnable(
		snapshotFutures,
		metadata,
		metrics,
		System.nanoTime(),
		taskName,
		registerConsumer(),
		unregisterConsumer(),
		env,
		asyncExceptionHandler));
}
 

/**
 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions, org.apache.flink.runtime.state.CheckpointStreamFactory)}.
 */
public OperatorSnapshotFinalizer snapshotWithLocalState(long checkpointId, long timestamp) throws Exception {

	OperatorSnapshotFutures operatorStateResult = operator.snapshotState(
		checkpointId,
		timestamp,
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		checkpointStorage.resolveCheckpointStorageLocation(checkpointId, CheckpointStorageLocationReference.getDefault()));

	return new OperatorSnapshotFinalizer(operatorStateResult);
}
 

public static <OUT, OP extends StreamOperator<OUT>> TaggedOperatorSubtaskState snapshot(
	OP operator,
	int index,
	long timestamp,
	boolean isExactlyOnceMode,
	boolean isUnalignedCheckpoint,
	CheckpointStorageWorkerView checkpointStorage,
	Path savepointPath) throws Exception {

	CheckpointOptions options = new CheckpointOptions(
		CheckpointType.SAVEPOINT,
		AbstractFsCheckpointStorage.encodePathAsReference(savepointPath),
		isExactlyOnceMode,
		isUnalignedCheckpoint);

	operator.prepareSnapshotPreBarrier(CHECKPOINT_ID);

	CheckpointStreamFactory storage = checkpointStorage.resolveCheckpointStorageLocation(
		CHECKPOINT_ID,
		options.getTargetLocation());

	OperatorSnapshotFutures snapshotInProgress = operator.snapshotState(
		CHECKPOINT_ID,
		timestamp,
		options,
		storage);

	OperatorSubtaskState state = new OperatorSnapshotFinalizer(snapshotInProgress).getJobManagerOwnedState();

	operator.notifyCheckpointComplete(CHECKPOINT_ID);
	return new TaggedOperatorSubtaskState(index, state);
}
 

/**
 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions, org.apache.flink.runtime.state.CheckpointStreamFactory)}.
 */
public OperatorSnapshotFinalizer snapshotWithLocalState(long checkpointId, long timestamp) throws Exception {

	OperatorSnapshotFutures operatorStateResult = operator.snapshotState(
		checkpointId,
		timestamp,
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		checkpointStorage.resolveCheckpointStorageLocation(checkpointId, CheckpointStorageLocationReference.getDefault()));

	return new OperatorSnapshotFinalizer(operatorStateResult);
}
 

private static StreamOperator<?> streamOperatorWithSnapshot(OperatorSnapshotFutures operatorSnapshotResult) throws Exception {
	StreamOperator<?> operator = mock(StreamOperator.class);
	when(operator.getOperatorID()).thenReturn(new OperatorID());

	when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class)))
		.thenReturn(operatorSnapshotResult);

	return operator;
}
 

/**
 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions, org.apache.flink.runtime.state.CheckpointStreamFactory)}.
 */
public OperatorSnapshotFinalizer snapshotWithLocalState(long checkpointId, long timestamp) throws Exception {

	OperatorSnapshotFutures operatorStateResult = operator.snapshotState(
		checkpointId,
		timestamp,
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		checkpointStorage.resolveCheckpointStorageLocation(checkpointId, CheckpointStorageLocationReference.getDefault()));

	return new OperatorSnapshotFinalizer(operatorStateResult);
}
 

@SuppressWarnings("deprecation")
private void checkpointStreamOperator(StreamOperator<?> op) throws Exception {
	if (null != op) {

		OperatorSnapshotFutures snapshotInProgress = op.snapshotState(
				checkpointMetaData.getCheckpointId(),
				checkpointMetaData.getTimestamp(),
				checkpointOptions,
				storageLocation);
		operatorSnapshotsInProgress.put(op.getOperatorID(), snapshotInProgress);
	}
}
 

public static <OUT, OP extends StreamOperator<OUT>> TaggedOperatorSubtaskState snapshot(
	OP operator,
	int index,
	long timestamp,
	CheckpointStorageWorkerView checkpointStorage,
	Path savepointPath) throws Exception {

	CheckpointOptions options = new CheckpointOptions(
		CheckpointType.SAVEPOINT,
		AbstractFsCheckpointStorage.encodePathAsReference(savepointPath));

	operator.prepareSnapshotPreBarrier(CHECKPOINT_ID);

	CheckpointStreamFactory storage = checkpointStorage.resolveCheckpointStorageLocation(
		CHECKPOINT_ID,
		options.getTargetLocation());

	OperatorSnapshotFutures snapshotInProgress = operator.snapshotState(
		CHECKPOINT_ID,
		timestamp,
		options,
		storage);

	OperatorSubtaskState state = new OperatorSnapshotFinalizer(snapshotInProgress).getJobManagerOwnedState();

	operator.notifyCheckpointComplete(CHECKPOINT_ID);
	return new TaggedOperatorSubtaskState(index, state);
}
 

AsyncCheckpointRunnable(
	StreamTask<?, ?> owner,
	Map<OperatorID, OperatorSnapshotFutures> operatorSnapshotsInProgress,
	CheckpointMetaData checkpointMetaData,
	CheckpointMetrics checkpointMetrics,
	long asyncStartNanos) {

	this.owner = Preconditions.checkNotNull(owner);
	this.operatorSnapshotsInProgress = Preconditions.checkNotNull(operatorSnapshotsInProgress);
	this.checkpointMetaData = Preconditions.checkNotNull(checkpointMetaData);
	this.checkpointMetrics = Preconditions.checkNotNull(checkpointMetrics);
	this.asyncStartNanos = asyncStartNanos;
}
 

@SuppressWarnings("deprecation")
private void checkpointStreamOperator(StreamOperator<?> op) throws Exception {
	if (null != op) {

		OperatorSnapshotFutures snapshotInProgress = op.snapshotState(
				checkpointMetaData.getCheckpointId(),
				checkpointMetaData.getTimestamp(),
				checkpointOptions,
				storageLocation);
		operatorSnapshotsInProgress.put(op.getOperatorID(), snapshotInProgress);
	}
}
 

@Test
public void testDecliningCheckpointStreamOperator() throws Exception {
	CheckpointExceptionHandlerTest.DeclineDummyEnvironment declineDummyEnvironment = new CheckpointExceptionHandlerTest.DeclineDummyEnvironment();

	// mock the returned snapshots
	OperatorSnapshotFutures operatorSnapshotResult1 = mock(OperatorSnapshotFutures.class);
	OperatorSnapshotFutures operatorSnapshotResult2 = mock(OperatorSnapshotFutures.class);

	final Exception testException = new Exception("Test exception");

	RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
		declineDummyEnvironment,
		operatorChain(
			streamOperatorWithSnapshot(operatorSnapshotResult1),
			streamOperatorWithSnapshot(operatorSnapshotResult2),
			streamOperatorWithSnapshotException(testException))));
	MockStreamTask streamTask = task.streamTask;

	waitTaskIsRunning(streamTask, task.invocationFuture);

	streamTask.triggerCheckpoint(
		new CheckpointMetaData(42L, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	assertEquals(testException, declineDummyEnvironment.getLastDeclinedCheckpointCause());

	verify(operatorSnapshotResult1).cancel();
	verify(operatorSnapshotResult2).cancel();

	task.streamTask.finishInput();
	task.waitForTaskCompletion(false);
}
 

/**
 * Tests that uncaught exceptions in the async part of a checkpoint operation are forwarded
 * to the uncaught exception handler. See <a href="https://issues.apache.org/jira/browse/FLINK-12889">FLINK-12889</a>.
 */
@Test
public void testUncaughtExceptionInAsynchronousCheckpointingOperation() throws Exception {
	final RuntimeException failingCause = new RuntimeException("Test exception");
	FailingDummyEnvironment failingDummyEnvironment = new FailingDummyEnvironment(failingCause);

	// mock the returned snapshots
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		ExceptionallyDoneFuture.of(failingCause),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));

	final TestingUncaughtExceptionHandler uncaughtExceptionHandler = new TestingUncaughtExceptionHandler();

	RunningTask<MockStreamTask> task = runTask(() -> new MockStreamTask(
		failingDummyEnvironment,
		operatorChain(streamOperatorWithSnapshot(operatorSnapshotResult)),
		uncaughtExceptionHandler));
	MockStreamTask streamTask = task.streamTask;

	waitTaskIsRunning(streamTask, task.invocationFuture);

	streamTask.triggerCheckpoint(
		new CheckpointMetaData(42L, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	final Throwable uncaughtException = uncaughtExceptionHandler.waitForUncaughtException();
	assertThat(uncaughtException, is(failingCause));

	streamTask.finishInput();
	task.waitForTaskCompletion(false);
}
 

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation discards the state handles of a not yet
 * acknowledged checkpoint and prevents sending an acknowledge message to the
 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
 * before Environment.acknowledgeCheckpoint().
 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final OneShotLatch createSubtask = new OneShotLatch();
	final OneShotLatch completeSubtask = new OneShotLatch();

	Environment mockEnvironment = spy(new MockEnvironmentBuilder().build());

	whenNew(OperatorSnapshotFinalizer.class).
		withAnyArguments().
		thenAnswer((Answer<OperatorSnapshotFinalizer>) invocation -> {
				createSubtask.trigger();
				completeSubtask.await();
				Object[] arguments = invocation.getArguments();
				return new OperatorSnapshotFinalizer((OperatorSnapshotFutures) arguments[0]);
			}
		);

	StreamTask<?, ?> streamTask = new EmptyStreamTask(mockEnvironment);
	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);

	final StreamOperator<?> streamOperator = mock(StreamOperator.class);
	final OperatorID operatorID = new OperatorID();
	when(streamOperator.getOperatorID()).thenReturn(operatorID);

	KeyedStateHandle managedKeyedStateHandle = mock(KeyedStateHandle.class);
	KeyedStateHandle rawKeyedStateHandle = mock(KeyedStateHandle.class);
	OperatorStateHandle managedOperatorStateHandle = mock(OperatorStreamStateHandle.class);
	OperatorStateHandle rawOperatorStateHandle = mock(OperatorStreamStateHandle.class);

	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)),
		DoneFuture.of(SnapshotResult.of(rawKeyedStateHandle)),
		DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)));

	when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult);

	StreamOperator<?>[] streamOperators = {streamOperator};

	OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
	when(operatorChain.getAllOperators()).thenReturn(streamOperators);

	CheckpointStorage checkpointStorage = new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE);

	ExecutorService executor = Executors.newFixedThreadPool(1);

	Whitebox.setInternalState(streamTask, "isRunning", true);
	Whitebox.setInternalState(streamTask, "lock", new Object());
	Whitebox.setInternalState(streamTask, "operatorChain", operatorChain);
	Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry());
	Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", executor);
	Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
	Whitebox.setInternalState(streamTask, "checkpointStorage", checkpointStorage);

	streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation());

	createSubtask.await();

	streamTask.cancel();

	completeSubtask.trigger();

	// wait for the completion of the async task
	executor.shutdown();

	if (!executor.awaitTermination(10000L, TimeUnit.MILLISECONDS)) {
		fail("Executor did not shut down within the given timeout. This indicates that the " +
			"checkpointing did not resume.");
	}

	// check that the checkpoint has not been acknowledged
	verify(mockEnvironment, never()).acknowledgeCheckpoint(eq(checkpointId), any(CheckpointMetrics.class), any(TaskStateSnapshot.class));

	// check that the state handles have been discarded
	verify(managedKeyedStateHandle).discardState();
	verify(rawKeyedStateHandle).discardState();
	verify(managedOperatorStateHandle).discardState();
	verify(rawOperatorStateHandle).discardState();
}
 

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	StreamOperator<?> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpoint(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}
 

/**
 * This tests the forwarding of jm and tm-local state from the futures reported by the backends, through the
 * async checkpointing thread to the {@link org.apache.flink.runtime.state.TaskStateManager}.
 */
@Test
public void testReportingFromSnapshotToTaskStateManager() {

	TestTaskStateManager taskStateManager = new TestTaskStateManager();

	StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(
		new Configuration(),
		new Configuration(),
		new ExecutionConfig(),
		1024 * 1024,
		new MockInputSplitProvider(),
		0,
		taskStateManager);

	StreamTask testStreamTask = new StreamTaskTest.NoOpStreamTask(streamMockEnvironment);
	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(0L, 0L);
	CheckpointMetrics checkpointMetrics = new CheckpointMetrics();

	Map<OperatorID, OperatorSnapshotFutures> snapshots = new HashMap<>(1);
	OperatorSnapshotFutures osFuture = new OperatorSnapshotFutures();

	osFuture.setKeyedStateManagedFuture(createSnapshotResult(KeyedStateHandle.class));
	osFuture.setKeyedStateRawFuture(createSnapshotResult(KeyedStateHandle.class));
	osFuture.setOperatorStateManagedFuture(createSnapshotResult(OperatorStateHandle.class));
	osFuture.setOperatorStateRawFuture(createSnapshotResult(OperatorStateHandle.class));

	OperatorID operatorID = new OperatorID();
	snapshots.put(operatorID, osFuture);

	StreamTask.AsyncCheckpointRunnable checkpointRunnable =
		new StreamTask.AsyncCheckpointRunnable(
			testStreamTask,
			snapshots,
			checkpointMetaData,
			checkpointMetrics,
			0L);

	checkpointRunnable.run();

	TaskStateSnapshot lastJobManagerTaskStateSnapshot = taskStateManager.getLastJobManagerTaskStateSnapshot();
	TaskStateSnapshot lastTaskManagerTaskStateSnapshot = taskStateManager.getLastTaskManagerTaskStateSnapshot();

	OperatorSubtaskState jmState =
		lastJobManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);

	OperatorSubtaskState tmState =
		lastTaskManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);

	performCheck(osFuture.getKeyedStateManagedFuture(), jmState.getManagedKeyedState(), tmState.getManagedKeyedState());
	performCheck(osFuture.getKeyedStateRawFuture(), jmState.getRawKeyedState(), tmState.getRawKeyedState());
	performCheck(osFuture.getOperatorStateManagedFuture(), jmState.getManagedOperatorState(), tmState.getManagedOperatorState());
	performCheck(osFuture.getOperatorStateRawFuture(), jmState.getRawOperatorState(), tmState.getRawOperatorState());
}
 

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	OneInputStreamOperator<String, String> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpointAsync(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}
 

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation discards the state handles of a not yet
 * acknowledged checkpoint and prevents sending an acknowledge message to the
 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
 * before Environment.acknowledgeCheckpoint().
 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {

	final TestingKeyedStateHandle managedKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingKeyedStateHandle rawKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingOperatorStateHandle managedOperatorStateHandle = new TestingOperatorStateHandle();
	final TestingOperatorStateHandle rawOperatorStateHandle = new TestingOperatorStateHandle();

	final BlockingRunnableFuture<SnapshotResult<KeyedStateHandle>> rawKeyedStateHandleFuture = new BlockingRunnableFuture<>(2, SnapshotResult.of(rawKeyedStateHandle));
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)),
		rawKeyedStateHandleFuture,
		DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));

	final OneInputStreamOperator<String, String> streamOperator = streamOperatorWithSnapshot(operatorSnapshotResult);

	final AcknowledgeDummyEnvironment mockEnvironment = new AcknowledgeDummyEnvironment();

	RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
		mockEnvironment,
		operatorChain(streamOperator)));

	waitTaskIsRunning(task.streamTask, task.invocationFuture);

	final long checkpointId = 42L;
	task.streamTask.triggerCheckpointAsync(
		new CheckpointMetaData(checkpointId, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	rawKeyedStateHandleFuture.awaitRun();

	task.streamTask.cancel();

	final FutureUtils.ConjunctFuture<Void> discardFuture = FutureUtils.waitForAll(
		asList(
			managedKeyedStateHandle.getDiscardFuture(),
			rawKeyedStateHandle.getDiscardFuture(),
			managedOperatorStateHandle.getDiscardFuture(),
			rawOperatorStateHandle.getDiscardFuture()));

	// make sure that all state handles have been discarded
	discardFuture.get();

	try {
		mockEnvironment.getAcknowledgeCheckpointFuture().get(10L, TimeUnit.MILLISECONDS);
		fail("The checkpoint should not get acknowledged.");
	} catch (TimeoutException expected) {
		// future should not be completed
	}

	task.waitForTaskCompletion(true);
}