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

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

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

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

public OperatorSnapshotFutures() {
	this(
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));
}
 

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

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

public OperatorSnapshotFutures() {
	this(
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));
}
 

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

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

public OperatorSnapshotFutures() {
	this(
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));
}
 

/**
 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
 * the StreamStateHandle result is retrievable that the state handle are discarded.
 */
@Test
public void testCancelAndCleanup() throws Exception {
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures();

	operatorSnapshotResult.cancel();

	KeyedStateHandle keyedManagedStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateManagedResult =
		SnapshotResult.of(keyedManagedStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture =
		spy(DoneFuture.of(keyedStateManagedResult));

	KeyedStateHandle keyedRawStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateRawResult =
		SnapshotResult.of(keyedRawStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture =
		spy(DoneFuture.of(keyedStateRawResult));

	OperatorStateHandle operatorManagedStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateManagedResult =
		SnapshotResult.of(operatorManagedStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture =
		spy(DoneFuture.of(operatorStateManagedResult));

	OperatorStateHandle operatorRawStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateRawResult =
		SnapshotResult.of(operatorRawStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture =
		spy(DoneFuture.of(operatorStateRawResult));

	operatorSnapshotResult = new OperatorSnapshotFutures(
		keyedStateManagedFuture,
		keyedStateRawFuture,
		operatorStateManagedFuture,
		operatorStateRawFuture);

	operatorSnapshotResult.cancel();

	verify(keyedStateManagedFuture).cancel(true);
	verify(keyedStateRawFuture).cancel(true);
	verify(operatorStateManagedFuture).cancel(true);
	verify(operatorStateRawFuture).cancel(true);

	verify(keyedManagedStateHandle).discardState();
	verify(keyedRawStateHandle).discardState();
	verify(operatorManagedStateHandle).discardState();
	verify(operatorRawStateHandle).discardState();
}
 

private static <T extends StateObject> RunnableFuture<SnapshotResult<T>> createSnapshotResult(Class<T> clazz) {
	return DoneFuture.of(SnapshotResult.withLocalState(mock(clazz), mock(clazz)));
}
 

/**
 * Test that the runnable futures are executed and the result is correctly extracted.
 */
@Test
public void testRunAndExtract() throws Exception{

	Random random = new Random(0x42);

	KeyedStateHandle keyedTemplate =
		StateHandleDummyUtil.createNewKeyedStateHandle(new KeyGroupRange(0, 0));
	OperatorStateHandle operatorTemplate =
		StateHandleDummyUtil.createNewOperatorStateHandle(2, random);

	SnapshotResult<KeyedStateHandle> snapKeyMan = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate),
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate));

	SnapshotResult<KeyedStateHandle> snapKeyRaw = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate),
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate));

	SnapshotResult<OperatorStateHandle> snapOpMan = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate),
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate));

	SnapshotResult<OperatorStateHandle> snapOpRaw = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate),
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate));

	DoneFuture<SnapshotResult<KeyedStateHandle>> managedKeyed = new PseudoNotDoneFuture<>(snapKeyMan);
	DoneFuture<SnapshotResult<KeyedStateHandle>> rawKeyed = new PseudoNotDoneFuture<>(snapKeyRaw);
	DoneFuture<SnapshotResult<OperatorStateHandle>> managedOp = new PseudoNotDoneFuture<>(snapOpMan);
	DoneFuture<SnapshotResult<OperatorStateHandle>> rawOp = new PseudoNotDoneFuture<>(snapOpRaw);

	Assert.assertFalse(managedKeyed.isDone());
	Assert.assertFalse(rawKeyed.isDone());
	Assert.assertFalse(managedOp.isDone());
	Assert.assertFalse(rawOp.isDone());

	OperatorSnapshotFutures futures = new OperatorSnapshotFutures(managedKeyed, rawKeyed, managedOp, rawOp);
	OperatorSnapshotFinalizer operatorSnapshotFinalizer = new OperatorSnapshotFinalizer(futures);

	Assert.assertTrue(managedKeyed.isDone());
	Assert.assertTrue(rawKeyed.isDone());
	Assert.assertTrue(managedOp.isDone());
	Assert.assertTrue(rawOp.isDone());

	OperatorSubtaskState jobManagerOwnedState = operatorSnapshotFinalizer.getJobManagerOwnedState();
	Assert.assertTrue(checkResult(snapKeyMan.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getManagedKeyedState()));
	Assert.assertTrue(checkResult(snapKeyRaw.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getRawKeyedState()));
	Assert.assertTrue(checkResult(snapOpMan.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getManagedOperatorState()));
	Assert.assertTrue(checkResult(snapOpRaw.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getRawOperatorState()));

	OperatorSubtaskState taskLocalState = operatorSnapshotFinalizer.getTaskLocalState();
	Assert.assertTrue(checkResult(snapKeyMan.getTaskLocalSnapshot(), taskLocalState.getManagedKeyedState()));
	Assert.assertTrue(checkResult(snapKeyRaw.getTaskLocalSnapshot(), taskLocalState.getRawKeyedState()));
	Assert.assertTrue(checkResult(snapOpMan.getTaskLocalSnapshot(), taskLocalState.getManagedOperatorState()));
	Assert.assertTrue(checkResult(snapOpRaw.getTaskLocalSnapshot(), taskLocalState.getRawOperatorState()));
}
 

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

	final StreamOperator<?> 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.triggerCheckpoint(
		new CheckpointMetaData(checkpointId, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	rawKeyedStateHandleFuture.awaitRun();

	task.streamTask.cancel();

	final FutureUtils.ConjunctFuture<Void> discardFuture = FutureUtils.waitForAll(
		Arrays.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);
}
 

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

private static <T extends StateObject> RunnableFuture<SnapshotResult<T>> createSnapshotResult(Class<T> clazz) {
	return DoneFuture.of(SnapshotResult.withLocalState(mock(clazz), mock(clazz)));
}
 

/**
 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
 * the StreamStateHandle result is retrievable that the state handle are discarded.
 */
@Test
public void testCancelAndCleanup() throws Exception {
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures();

	operatorSnapshotResult.cancel();

	KeyedStateHandle keyedManagedStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateManagedResult = SnapshotResult.of(keyedManagedStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture = spy(DoneFuture.of(keyedStateManagedResult));

	KeyedStateHandle keyedRawStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateRawResult = SnapshotResult.of(keyedRawStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture = spy(DoneFuture.of(keyedStateRawResult));

	OperatorStateHandle operatorManagedStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateManagedResult = SnapshotResult.of(operatorManagedStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture = spy(DoneFuture.of(operatorStateManagedResult));

	OperatorStateHandle operatorRawStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateRawResult = SnapshotResult.of(operatorRawStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture = spy(DoneFuture.of(operatorStateRawResult));

	InputChannelStateHandle inputChannelRawStateHandle = mock(InputChannelStateHandle.class);
	SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannelStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(inputChannelRawStateHandle));
	Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>> inputChannelStateRawFuture = spy(DoneFuture.of(inputChannelStateRawResult));

	ResultSubpartitionStateHandle resultSubpartitionRawStateHandle = mock(ResultSubpartitionStateHandle.class);
	SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartitionStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(resultSubpartitionRawStateHandle));
	Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>> resultSubpartitionStateRawFuture = spy(DoneFuture.of(resultSubpartitionStateRawResult));

	operatorSnapshotResult = new OperatorSnapshotFutures(
		keyedStateManagedFuture,
		keyedStateRawFuture,
		operatorStateManagedFuture,
		operatorStateRawFuture,
		inputChannelStateRawFuture,
		resultSubpartitionStateRawFuture);

	operatorSnapshotResult.cancel();

	verify(keyedStateManagedFuture).cancel(true);
	verify(keyedStateRawFuture).cancel(true);
	verify(operatorStateManagedFuture).cancel(true);
	verify(operatorStateRawFuture).cancel(true);
	verify(inputChannelStateRawFuture).cancel(true);
	verify(resultSubpartitionStateRawFuture).cancel(true);

	verify(keyedManagedStateHandle).discardState();
	verify(keyedRawStateHandle).discardState();
	verify(operatorManagedStateHandle).discardState();
	verify(operatorRawStateHandle).discardState();
	verify(inputChannelRawStateHandle).discardState();
	verify(resultSubpartitionRawStateHandle).discardState();
}
 

/**
 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
 * the StreamStateHandle result is retrievable that the state handle are discarded.
 */
@Test
public void testCancelAndCleanup() throws Exception {
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures();

	operatorSnapshotResult.cancel();

	KeyedStateHandle keyedManagedStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateManagedResult =
		SnapshotResult.of(keyedManagedStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture =
		spy(DoneFuture.of(keyedStateManagedResult));

	KeyedStateHandle keyedRawStateHandle = mock(KeyedStateHandle.class);
	SnapshotResult<KeyedStateHandle> keyedStateRawResult =
		SnapshotResult.of(keyedRawStateHandle);
	RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture =
		spy(DoneFuture.of(keyedStateRawResult));

	OperatorStateHandle operatorManagedStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateManagedResult =
		SnapshotResult.of(operatorManagedStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture =
		spy(DoneFuture.of(operatorStateManagedResult));

	OperatorStateHandle operatorRawStateHandle = mock(OperatorStreamStateHandle.class);
	SnapshotResult<OperatorStateHandle> operatorStateRawResult =
		SnapshotResult.of(operatorRawStateHandle);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture =
		spy(DoneFuture.of(operatorStateRawResult));

	operatorSnapshotResult = new OperatorSnapshotFutures(
		keyedStateManagedFuture,
		keyedStateRawFuture,
		operatorStateManagedFuture,
		operatorStateRawFuture);

	operatorSnapshotResult.cancel();

	verify(keyedStateManagedFuture).cancel(true);
	verify(keyedStateRawFuture).cancel(true);
	verify(operatorStateManagedFuture).cancel(true);
	verify(operatorStateRawFuture).cancel(true);

	verify(keyedManagedStateHandle).discardState();
	verify(keyedRawStateHandle).discardState();
	verify(operatorManagedStateHandle).discardState();
	verify(operatorRawStateHandle).discardState();
}
 

private static <T extends StateObject> RunnableFuture<SnapshotResult<T>> createSnapshotResult(Class<T> clazz) {
	return DoneFuture.of(SnapshotResult.withLocalState(mock(clazz), mock(clazz)));
}
 

private static <T extends StateObject> RunnableFuture<SnapshotResult<StateObjectCollection<T>>> createSnapshotCollectionResult(Class<T> clazz) {
	return DoneFuture.of(SnapshotResult.withLocalState(singleton(mock(clazz)), singleton(mock(clazz))));
}
 

/**
 * Test that the runnable futures are executed and the result is correctly extracted.
 */
@Test
public void testRunAndExtract() throws Exception{

	Random random = new Random(0x42);

	KeyedStateHandle keyedTemplate =
		StateHandleDummyUtil.createNewKeyedStateHandle(new KeyGroupRange(0, 0));
	OperatorStateHandle operatorTemplate =
		StateHandleDummyUtil.createNewOperatorStateHandle(2, random);

	SnapshotResult<KeyedStateHandle> snapKeyMan = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate),
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate));

	SnapshotResult<KeyedStateHandle> snapKeyRaw = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate),
		StateHandleDummyUtil.deepDummyCopy(keyedTemplate));

	SnapshotResult<OperatorStateHandle> snapOpMan = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate),
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate));

	SnapshotResult<OperatorStateHandle> snapOpRaw = SnapshotResult.withLocalState(
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate),
		StateHandleDummyUtil.deepDummyCopy(operatorTemplate));

	DoneFuture<SnapshotResult<KeyedStateHandle>> managedKeyed = new PseudoNotDoneFuture<>(snapKeyMan);
	DoneFuture<SnapshotResult<KeyedStateHandle>> rawKeyed = new PseudoNotDoneFuture<>(snapKeyRaw);
	DoneFuture<SnapshotResult<OperatorStateHandle>> managedOp = new PseudoNotDoneFuture<>(snapOpMan);
	DoneFuture<SnapshotResult<OperatorStateHandle>> rawOp = new PseudoNotDoneFuture<>(snapOpRaw);

	Assert.assertFalse(managedKeyed.isDone());
	Assert.assertFalse(rawKeyed.isDone());
	Assert.assertFalse(managedOp.isDone());
	Assert.assertFalse(rawOp.isDone());

	OperatorSnapshotFutures futures = new OperatorSnapshotFutures(managedKeyed, rawKeyed, managedOp, rawOp);
	OperatorSnapshotFinalizer operatorSnapshotFinalizer = new OperatorSnapshotFinalizer(futures);

	Assert.assertTrue(managedKeyed.isDone());
	Assert.assertTrue(rawKeyed.isDone());
	Assert.assertTrue(managedOp.isDone());
	Assert.assertTrue(rawOp.isDone());

	OperatorSubtaskState jobManagerOwnedState = operatorSnapshotFinalizer.getJobManagerOwnedState();
	Assert.assertTrue(checkResult(snapKeyMan.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getManagedKeyedState()));
	Assert.assertTrue(checkResult(snapKeyRaw.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getRawKeyedState()));
	Assert.assertTrue(checkResult(snapOpMan.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getManagedOperatorState()));
	Assert.assertTrue(checkResult(snapOpRaw.getJobManagerOwnedSnapshot(), jobManagerOwnedState.getRawOperatorState()));

	OperatorSubtaskState taskLocalState = operatorSnapshotFinalizer.getTaskLocalState();
	Assert.assertTrue(checkResult(snapKeyMan.getTaskLocalSnapshot(), taskLocalState.getManagedKeyedState()));
	Assert.assertTrue(checkResult(snapKeyRaw.getTaskLocalSnapshot(), taskLocalState.getRawKeyedState()));
	Assert.assertTrue(checkResult(snapOpMan.getTaskLocalSnapshot(), taskLocalState.getManagedOperatorState()));
	Assert.assertTrue(checkResult(snapOpRaw.getTaskLocalSnapshot(), taskLocalState.getRawOperatorState()));
}
 

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