Java Code Examples for org.apache.flink.runtime.state.DoneFuture#of()

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

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

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

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

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

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

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

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

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

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

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

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