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

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

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

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

@Override
protected SnapshotResult<KeyedStateHandle> callInternal() throws Exception {

	boolean completed = false;

	// Handle to the meta data file
	SnapshotResult<StreamStateHandle> metaStateHandle = null;
	// Handles to new sst files since the last completed checkpoint will go here
	final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
	// Handles to the misc files in the current snapshot will go here
	final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();

	try {

		metaStateHandle = materializeMetaData();

		// Sanity checks - they should never fail
		Preconditions.checkNotNull(metaStateHandle, "Metadata was not properly created.");
		Preconditions.checkNotNull(metaStateHandle.getJobManagerOwnedSnapshot(),
			"Metadata for job manager was not properly created.");

		uploadSstFiles(sstFiles, miscFiles);

		synchronized (materializedSstFiles) {
			materializedSstFiles.put(checkpointId, sstFiles.keySet());
		}

		final IncrementalRemoteKeyedStateHandle jmIncrementalKeyedStateHandle =
			new IncrementalRemoteKeyedStateHandle(
				backendUID,
				keyGroupRange,
				checkpointId,
				sstFiles,
				miscFiles,
				metaStateHandle.getJobManagerOwnedSnapshot());

		final DirectoryStateHandle directoryStateHandle = localBackupDirectory.completeSnapshotAndGetHandle();
		final SnapshotResult<KeyedStateHandle> snapshotResult;
		if (directoryStateHandle != null && metaStateHandle.getTaskLocalSnapshot() != null) {

			IncrementalLocalKeyedStateHandle localDirKeyedStateHandle =
				new IncrementalLocalKeyedStateHandle(
					backendUID,
					checkpointId,
					directoryStateHandle,
					keyGroupRange,
					metaStateHandle.getTaskLocalSnapshot(),
					sstFiles.keySet());

			snapshotResult = SnapshotResult.withLocalState(jmIncrementalKeyedStateHandle, localDirKeyedStateHandle);
		} else {
			snapshotResult = SnapshotResult.of(jmIncrementalKeyedStateHandle);
		}

		completed = true;

		return snapshotResult;
	} finally {
		if (!completed) {
			final List<StateObject> statesToDiscard =
				new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
			statesToDiscard.add(metaStateHandle);
			statesToDiscard.addAll(miscFiles.values());
			statesToDiscard.addAll(sstFiles.values());
			cleanupIncompleteSnapshot(statesToDiscard);
		}
	}
}
 

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

@Override
protected SnapshotResult<KeyedStateHandle> callInternal() throws Exception {

	boolean completed = false;

	// Handle to the meta data file
	SnapshotResult<StreamStateHandle> metaStateHandle = null;
	// Handles to new sst files since the last completed checkpoint will go here
	final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
	// Handles to the misc files in the current snapshot will go here
	final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();

	try {

		metaStateHandle = materializeMetaData();

		// Sanity checks - they should never fail
		Preconditions.checkNotNull(metaStateHandle, "Metadata was not properly created.");
		Preconditions.checkNotNull(metaStateHandle.getJobManagerOwnedSnapshot(),
			"Metadata for job manager was not properly created.");

		uploadSstFiles(sstFiles, miscFiles);

		synchronized (materializedSstFiles) {
			materializedSstFiles.put(checkpointId, sstFiles.keySet());
		}

		final IncrementalRemoteKeyedStateHandle jmIncrementalKeyedStateHandle =
			new IncrementalRemoteKeyedStateHandle(
				backendUID,
				keyGroupRange,
				checkpointId,
				sstFiles,
				miscFiles,
				metaStateHandle.getJobManagerOwnedSnapshot());

		final DirectoryStateHandle directoryStateHandle = localBackupDirectory.completeSnapshotAndGetHandle();
		final SnapshotResult<KeyedStateHandle> snapshotResult;
		if (directoryStateHandle != null && metaStateHandle.getTaskLocalSnapshot() != null) {

			IncrementalLocalKeyedStateHandle localDirKeyedStateHandle =
				new IncrementalLocalKeyedStateHandle(
					backendUID,
					checkpointId,
					directoryStateHandle,
					keyGroupRange,
					metaStateHandle.getTaskLocalSnapshot(),
					sstFiles.keySet());

			snapshotResult = SnapshotResult.withLocalState(jmIncrementalKeyedStateHandle, localDirKeyedStateHandle);
		} else {
			snapshotResult = SnapshotResult.of(jmIncrementalKeyedStateHandle);
		}

		completed = true;

		return snapshotResult;
	} finally {
		if (!completed) {
			final List<StateObject> statesToDiscard =
				new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
			statesToDiscard.add(metaStateHandle);
			statesToDiscard.addAll(miscFiles.values());
			statesToDiscard.addAll(sstFiles.values());
			cleanupIncompleteSnapshot(statesToDiscard);
		}
	}
}
 

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

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

@Override
protected SnapshotResult<KeyedStateHandle> callInternal() throws Exception {

	boolean completed = false;

	// Handle to the meta data file
	SnapshotResult<StreamStateHandle> metaStateHandle = null;
	// Handles to new sst files since the last completed checkpoint will go here
	final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
	// Handles to the misc files in the current snapshot will go here
	final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();

	try {

		metaStateHandle = materializeMetaData();

		// Sanity checks - they should never fail
		Preconditions.checkNotNull(metaStateHandle, "Metadata was not properly created.");
		Preconditions.checkNotNull(metaStateHandle.getJobManagerOwnedSnapshot(),
			"Metadata for job manager was not properly created.");

		uploadSstFiles(sstFiles, miscFiles);

		synchronized (materializedSstFiles) {
			materializedSstFiles.put(checkpointId, sstFiles.keySet());
		}

		final IncrementalRemoteKeyedStateHandle jmIncrementalKeyedStateHandle =
			new IncrementalRemoteKeyedStateHandle(
				backendUID,
				keyGroupRange,
				checkpointId,
				sstFiles,
				miscFiles,
				metaStateHandle.getJobManagerOwnedSnapshot());

		final DirectoryStateHandle directoryStateHandle = localBackupDirectory.completeSnapshotAndGetHandle();
		final SnapshotResult<KeyedStateHandle> snapshotResult;
		if (directoryStateHandle != null && metaStateHandle.getTaskLocalSnapshot() != null) {

			IncrementalLocalKeyedStateHandle localDirKeyedStateHandle =
				new IncrementalLocalKeyedStateHandle(
					backendUID,
					checkpointId,
					directoryStateHandle,
					keyGroupRange,
					metaStateHandle.getTaskLocalSnapshot(),
					sstFiles.keySet());

			snapshotResult = SnapshotResult.withLocalState(jmIncrementalKeyedStateHandle, localDirKeyedStateHandle);
		} else {
			snapshotResult = SnapshotResult.of(jmIncrementalKeyedStateHandle);
		}

		completed = true;

		return snapshotResult;
	} finally {
		if (!completed) {
			final List<StateObject> statesToDiscard =
				new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
			statesToDiscard.add(metaStateHandle);
			statesToDiscard.addAll(miscFiles.values());
			statesToDiscard.addAll(sstFiles.values());
			cleanupIncompleteSnapshot(statesToDiscard);
		}
	}
}
 

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

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