org.apache.flink.runtime.checkpoint.CheckpointOptions Java Examples

@Override
public boolean triggerCheckpoint(CheckpointMetaData checkpointMetaData, CheckpointOptions checkpointOptions) throws Exception {
	try {
		// No alignment if we inject a checkpoint
		CheckpointMetrics checkpointMetrics = new CheckpointMetrics()
				.setBytesBufferedInAlignment(0L)
				.setAlignmentDurationNanos(0L);

		return performCheckpoint(checkpointMetaData, checkpointOptions, checkpointMetrics);
	}
	catch (Exception e) {
		// propagate exceptions only if the task is still in "running" state
		if (isRunning) {
			throw new Exception("Could not perform checkpoint " + checkpointMetaData.getCheckpointId() +
				" for operator " + getName() + '.', e);
		} else {
			LOG.debug("Could not perform checkpoint {} for operator {} while the " +
				"invokable was not in state running.", checkpointMetaData.getCheckpointId(), getName(), e);
			return false;
		}
	}
}
 

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

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

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

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

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

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

@Override
public Optional<BufferOrEvent> getNextBufferOrEvent() throws IOException, InterruptedException {
	currentChannel = (currentChannel + 1) % numberOfChannels;

	if (barrierGens[currentChannel].isNextBarrier()) {
		return Optional.of(
			new BufferOrEvent(
				new CheckpointBarrier(
					++currentBarriers[currentChannel],
					System.currentTimeMillis(),
					CheckpointOptions.forCheckpointWithDefaultLocation()),
				currentChannel));
	} else {
		Buffer buffer = bufferPools[currentChannel].requestBuffer();
		buffer.getMemorySegment().putLong(0, c++);
		return Optional.of(new BufferOrEvent(buffer, currentChannel));
	}
}
 

@Override
public Future<Boolean> triggerCheckpointAsync(
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		boolean advanceToEndOfEventTime) {

	CompletableFuture<Boolean> result = new CompletableFuture<>();
	mainMailboxExecutor.execute(
			() -> {
				try {
					result.complete(triggerCheckpoint(checkpointMetaData, checkpointOptions, advanceToEndOfEventTime));
				}
				catch (Exception ex) {
					// Report the failure both via the Future result but also to the mailbox
					result.completeExceptionally(ex);
					throw ex;
				}
			},
			"checkpoint %s with %s",
		checkpointMetaData,
		checkpointOptions);
	return result;
}
 

@Test
public void testCheckpointCallsInOrder() throws Exception {

	Task task = createTask(CheckpointsInOrderInvokable.class);
	try (TaskCleaner ignored = new TaskCleaner(task)) {
		task.startTaskThread();

		awaitLatch.await();

		for (int i = 1; i <= numCalls; i++) {
			task.triggerCheckpointBarrier(i, 156865867234L, CheckpointOptions.forCheckpointWithDefaultLocation(), false);
		}

		triggerLatch.await();

		assertFalse(task.isCanceledOrFailed());

		ExecutionState currentState = task.getExecutionState();
		assertThat(currentState, isOneOf(ExecutionState.RUNNING, ExecutionState.FINISHED));
	}
}
 

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

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

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

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

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

	final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);

	long checkpointId = 42L;
	subtaskCheckpointCoordinator.checkpointState(
		new CheckpointMetaData(checkpointId, System.currentTimeMillis()),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		new CheckpointMetrics(),
		operatorChain,
		() -> true);
	subtaskCheckpointCoordinator.notifyCheckpointAborted(checkpointId, operatorChain, () -> true);
	assertEquals(0, subtaskCheckpointCoordinator.getAbortedCheckpointSize());
	assertEquals(checkpointId, stateManager.getNotifiedAbortedCheckpointId());
}
 

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

@Nonnull
@Override
protected RunnableFuture<SnapshotResult<KeyedStateHandle>> doSnapshot(
	long checkpointId,
	long checkpointTimestamp,
	@Nonnull CheckpointStreamFactory checkpointStreamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	final SnapshotDirectory snapshotDirectory = prepareLocalSnapshotDirectory(checkpointId);
	LOG.trace("Local RocksDB checkpoint goes to backup path {}.", snapshotDirectory);

	final List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = new ArrayList<>(kvStateInformation.size());
	final Set<StateHandleID> baseSstFiles = snapshotMetaData(checkpointId, stateMetaInfoSnapshots);

	takeDBNativeCheckpoint(snapshotDirectory);

	final RocksDBIncrementalSnapshotOperation snapshotOperation =
		new RocksDBIncrementalSnapshotOperation(
			checkpointId,
			checkpointStreamFactory,
			snapshotDirectory,
			baseSstFiles,
			stateMetaInfoSnapshots);

	return snapshotOperation.toAsyncSnapshotFutureTask(cancelStreamRegistry);
}
 

@Override
public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) {
	ACTUAL_ORDER_TRACKING.add("OPERATOR::setup");
	super.setup(containingTask, config, output);
	if (simulateCheckpointing) {
		testCheckpointer = new Thread() {
			@Override
			public void run() {
				try {
					runStarted.await();
					if (getContainingTask().isCanceled() || getContainingTask().triggerCheckpoint(
							new CheckpointMetaData(0, System.currentTimeMillis()),
							CheckpointOptions.forCheckpointWithDefaultLocation(),
							false)) {
						LifecycleTrackingStreamSource.runFinish.trigger();
					}
				} catch (Exception e) {
					e.printStackTrace();
					Assert.fail();
				}
			}
		};
		testCheckpointer.start();
	}
}
 

@Test
public void testEmptyStateCheckpointing() {

	try {
		CheckpointStreamFactory streamFactory = createStreamFactory();
		SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

		ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);

		// draw a snapshot
		KeyedStateHandle snapshot =
			runSnapshot(
				backend.snapshot(682375462379L, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
				sharedStateRegistry);
		assertNull(snapshot);
		backend.dispose();

		backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot);
		backend.dispose();
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

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

	long startTime = System.currentTimeMillis();

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

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

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

	chosenSnapshotStrategy.logSyncCompleted(streamFactory, startTime);

	return snapshotRunner;
}
 

@Test
public void testSnapshotEmpty() throws Exception {
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();

	final OperatorStateBackend operatorStateBackend =
			abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", emptyStateHandles, cancelStreamRegistry);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot =
			operatorStateBackend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtils.runIfNotDoneAndGet(snapshot);
	OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
	assertNull(stateHandle);
}
 

/**
 * Asserts that {@link AbstractInvokable#triggerCheckpoint(CheckpointMetaData, CheckpointOptions)},
 * {@link AbstractInvokable#notifyCheckpointComplete(long)}, and {@link StoppableTask#stop()} are
 * invoked by a thread whose context class loader is set to the user code class loader.
 */
@Test
public void testSetsUserCodeClassLoader() throws Exception {
	numCalls = 1;

	Task task = createTask(ContextClassLoaderInterceptingInvokable.class);
	try (TaskCleaner ignored = new TaskCleaner(task)) {
		task.startTaskThread();

		awaitLatch.await();

		task.triggerCheckpointBarrier(1, 1, CheckpointOptions.forCheckpointWithDefaultLocation());
		task.notifyCheckpointComplete(1);
		task.stopExecution();

		triggerLatch.await();
		notifyCheckpointCompleteLatch.await();
		stopLatch.await();

		assertThat(classLoaders, hasSize(greaterThanOrEqualTo(3)));
		assertThat(classLoaders, everyItem(instanceOf(TestUserCodeClassLoader.class)));
	}
}
 

@Test
public void testCheckpointBarrierSerialization() throws Exception {
	long id = Integer.MAX_VALUE + 123123L;
	long timestamp = Integer.MAX_VALUE + 1228L;

	CheckpointOptions checkpoint = CheckpointOptions.forCheckpointWithDefaultLocation();
	testCheckpointBarrierSerialization(id, timestamp, checkpoint);

	final byte[] reference = new byte[] { 15, 52, 52, 11, 0, 0, 0, 0, -1, -23, -19, 35 };

	CheckpointOptions savepoint = new CheckpointOptions(
			CheckpointType.SAVEPOINT, new CheckpointStorageLocationReference(reference));
	testCheckpointBarrierSerialization(id, timestamp, savepoint);
}
 

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

@Test
public void testPriorityQueueStateCreationFailsIfNewSerializerIsNotCompatible() throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();

	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {
		InternalPriorityQueue<TestType> internalPriorityQueue = backend.create(
			"testPriorityQueue", new TestType.V1TestTypeSerializer());

		internalPriorityQueue.add(new TestType("key-1", 123));
		internalPriorityQueue.add(new TestType("key-2", 346));
		internalPriorityQueue.add(new TestType("key-1", 777));

		KeyedStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		backend.dispose();

		backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot);
		backend.create(
			"testPriorityQueue", new TestType.IncompatibleTestTypeSerializer());

		Assert.fail("should have failed");
	} catch (Exception e) {
		Assert.assertTrue(ExceptionUtils.findThrowable(e, StateMigrationException.class).isPresent());
	} finally {
		backend.dispose();
	}
}
 

/**
 * Tests that the created StateSnapshotContextSynchronousImpl is closed in case of a failing
 * Operator#snapshotState(StateSnapshotContextSynchronousImpl) call.
 */
@Test
public void testFailingSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

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

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = mock(StateSnapshotContextSynchronousImpl.class);

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

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

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

	// lets fail when calling the actual snapshotState method
	doThrow(failingException).when(operator).snapshotState(eq(context));

	try {
		operator.snapshotState(
				checkpointId,
				timestamp,
				CheckpointOptions.forCheckpointWithDefaultLocation(),
				new MemCheckpointStreamFactory(Integer.MAX_VALUE));
		fail("Exception expected.");
	} catch (Exception e) {
		assertEquals(failingException, e.getCause());
	}

	verify(context).close();
}
 

private void testBroadcastStateValueUpgrade(
		MapStateDescriptor<Integer, TestType> initialAccessDescriptor,
		MapStateDescriptor<Integer, TestType> newAccessDescriptorAfterRestore) throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();

	OperatorStateBackend backend = createOperatorStateBackend();

	try {
		BroadcastState<Integer, TestType> state = backend.getBroadcastState(initialAccessDescriptor);

		state.put(3, new TestType("foo", 13));
		state.put(5, new TestType("bar", 278));

		OperatorStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()));
		backend.dispose();

		backend = restoreOperatorStateBackend(snapshot);

		state = backend.getBroadcastState(newAccessDescriptorAfterRestore);

		// the state backend should have decided whether or not it needs to perform state migration;
		// make sure that reading and writing each broadcast entry works with the new serializer
		Assert.assertEquals(new TestType("foo", 13), state.get(3));
		Assert.assertEquals(new TestType("bar", 278), state.get(5));
		state.put(17, new TestType("new-entry", 777));
	} finally {
		backend.dispose();
	}
}
 

@Override
public void triggerCheckpointOnBarrier(
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		CheckpointMetrics checkpointMetrics) throws Exception {
	assertTrue("wrong checkpoint id", nextExpectedCheckpointId == -1L ||
		nextExpectedCheckpointId == checkpointMetaData.getCheckpointId());

	assertTrue(checkpointMetaData.getTimestamp() > 0);
	assertTrue(checkpointMetrics.getBytesBufferedInAlignment() >= 0);
	assertTrue(checkpointMetrics.getAlignmentDurationNanos() >= 0);

	nextExpectedCheckpointId++;
	lastReportedBytesBufferedInAlignment = checkpointMetrics.getBytesBufferedInAlignment();
}
 

@Override
public Boolean call() throws Exception {
	for (int i = 0; i < numCheckpoints; i++) {
		long currentCheckpointId = checkpointId.getAndIncrement();
		CheckpointMetaData checkpointMetaData = new CheckpointMetaData(currentCheckpointId, 0L);
		sourceTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation());
		Thread.sleep(checkpointInterval);
	}
	return true;
}
 

public CheckpointingOperation(
		StreamTask<?, ?> owner,
		CheckpointMetaData checkpointMetaData,
		CheckpointOptions checkpointOptions,
		CheckpointStreamFactory checkpointStorageLocation,
		CheckpointMetrics checkpointMetrics) {

	this.owner = Preconditions.checkNotNull(owner);
	this.checkpointMetaData = Preconditions.checkNotNull(checkpointMetaData);
	this.checkpointOptions = Preconditions.checkNotNull(checkpointOptions);
	this.checkpointMetrics = Preconditions.checkNotNull(checkpointMetrics);
	this.storageLocation = Preconditions.checkNotNull(checkpointStorageLocation);
	this.allOperators = owner.operatorChain.getAllOperators();
	this.operatorSnapshotsInProgress = new HashMap<>(allOperators.length);
}
 

/**
 * Calls the invokable to trigger a checkpoint.
 *
 * @param checkpointID The ID identifying the checkpoint.
 * @param checkpointTimestamp The timestamp associated with the checkpoint.
 * @param checkpointOptions Options for performing this checkpoint.
 * @param advanceToEndOfEventTime Flag indicating if the source should inject a {@code MAX_WATERMARK} in the pipeline
 *                           to fire any registered event-time timers.
 */
public void triggerCheckpointBarrier(
		final long checkpointID,
		final long checkpointTimestamp,
		final CheckpointOptions checkpointOptions,
		final boolean advanceToEndOfEventTime) {

	final AbstractInvokable invokable = this.invokable;
	final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointID, checkpointTimestamp);

	if (executionState == ExecutionState.RUNNING && invokable != null) {
		try {
			invokable.triggerCheckpointAsync(checkpointMetaData, checkpointOptions, advanceToEndOfEventTime);
		}
		catch (RejectedExecutionException ex) {
			// This may happen if the mailbox is closed. It means that the task is shutting down, so we just ignore it.
			LOG.debug(
				"Triggering checkpoint {} for {} ({}) was rejected by the mailbox",
				checkpointID, taskNameWithSubtask, executionId);
		}
		catch (Throwable t) {
			if (getExecutionState() == ExecutionState.RUNNING) {
				failExternally(new Exception(
					"Error while triggering checkpoint " + checkpointID + " for " +
						taskNameWithSubtask, t));
			} else {
				LOG.debug("Encountered error while triggering checkpoint {} for " +
					"{} ({}) while being not in state running.", checkpointID,
					taskNameWithSubtask, executionId, t);
			}
		}
	}
	else {
		LOG.debug("Declining checkpoint request for non-running task {} ({}).", taskNameWithSubtask, executionId);

		// send back a message that we did not do the checkpoint
		checkpointResponder.declineCheckpoint(jobId, executionId, checkpointID,
				new CheckpointException("Task name with subtask : " + taskNameWithSubtask, CheckpointFailureReason.CHECKPOINT_DECLINED_TASK_NOT_READY));
	}
}
 

public TriggerCheckpoint(
		JobID job,
		ExecutionAttemptID taskExecutionId,
		long checkpointId,
		long timestamp,
		CheckpointOptions checkpointOptions) {

	super(job, taskExecutionId, checkpointId);
	this.timestamp = timestamp;
	this.checkpointOptions = checkNotNull(checkpointOptions);
}
 

/**
 * Tests broadcasting events when no records have been emitted yet.
 */
@Test
public void testBroadcastEventNoRecords() throws Exception {
	int numberOfChannels = 4;
	int bufferSize = 32;

	@SuppressWarnings("unchecked")
	Queue<BufferConsumer>[] queues = new Queue[numberOfChannels];
	for (int i = 0; i < numberOfChannels; i++) {
		queues[i] = new ArrayDeque<>();
	}

	TestPooledBufferProvider bufferProvider = new TestPooledBufferProvider(Integer.MAX_VALUE, bufferSize);

	ResultPartitionWriter partitionWriter = new CollectingPartitionWriter(queues, bufferProvider);
	RecordWriter<ByteArrayIO> writer = new RecordWriterBuilder().build(partitionWriter);
	CheckpointBarrier barrier = new CheckpointBarrier(Integer.MAX_VALUE + 919192L, Integer.MAX_VALUE + 18828228L, CheckpointOptions.forCheckpointWithDefaultLocation());

	// No records emitted yet, broadcast should not request a buffer
	writer.broadcastEvent(barrier);

	assertEquals(0, bufferProvider.getNumberOfCreatedBuffers());

	for (int i = 0; i < numberOfChannels; i++) {
		assertEquals(1, queues[i].size());
		BufferOrEvent boe = parseBuffer(queues[i].remove(), i);
		assertTrue(boe.isEvent());
		assertEquals(barrier, boe.getEvent());
		assertEquals(0, queues[i].size());
	}
}
 

@Test
public void testNoNotifyOnSavepoint() throws IOException {
	TestBufferReceivedListener listener = new TestBufferReceivedListener();
	SingleInputGate inputGate = new SingleInputGateBuilder().build();
	inputGate.registerBufferReceivedListener(listener);
	LocalInputChannel channel = InputChannelBuilder.newBuilder().buildLocalChannel(inputGate);
	CheckpointBarrier barrier = new CheckpointBarrier(123L, 123L, new CheckpointOptions(SAVEPOINT, CheckpointStorageLocationReference.getDefault()));
	channel.notifyPriorityEvent(new BufferConsumer(toBuffer(barrier).getMemorySegment(), FreeingBufferRecycler.INSTANCE, getDataType(barrier)));
	channel.checkError();
	assertTrue(listener.notifiedOnBarriers.isEmpty());
}
 

@Test
public void testCompletingSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		Thread asyncSnapshotThread = new Thread(snapshot);
		asyncSnapshotThread.start();
		waiter.await(); // wait for snapshot to run
		waiter.reset();
		runStateUpdates();
		blocker.trigger(); // allow checkpointing to start writing
		waiter.await(); // wait for snapshot stream writing to run

		SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();
		KeyedStateHandle keyedStateHandle = snapshotResult.getJobManagerOwnedSnapshot();
		assertNotNull(keyedStateHandle);
		assertTrue(keyedStateHandle.getStateSize() > 0);
		assertEquals(2, keyedStateHandle.getKeyGroupRange().getNumberOfKeyGroups());

		for (BlockingCheckpointOutputStream stream : testStreamFactory.getAllCreatedStreams()) {
			assertTrue(stream.isClosed());
		}

		asyncSnapshotThread.join();
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

private void testBroadcastStateKeyUpgrade(
		MapStateDescriptor<TestType, Integer> initialAccessDescriptor,
		MapStateDescriptor<TestType, Integer> newAccessDescriptorAfterRestore) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();

	OperatorStateBackend backend = createOperatorStateBackend();

	try {
		BroadcastState<TestType, Integer> state = backend.getBroadcastState(initialAccessDescriptor);

		state.put(new TestType("foo", 13), 3);
		state.put(new TestType("bar", 278), 5);

		OperatorStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()));
		backend.dispose();

		backend = restoreOperatorStateBackend(snapshot);

		state = backend.getBroadcastState(newAccessDescriptorAfterRestore);

		// the state backend should have decided whether or not it needs to perform state migration;
		// make sure that reading and writing each broadcast entry works with the new serializer
		assertEquals((Integer) 3, state.get(new TestType("foo", 13)));
		assertEquals((Integer) 5, state.get(new TestType("bar", 278)));
		state.put(new TestType("new-entry", 777), 17);
	} finally {
		backend.dispose();
	}
}
 

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