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

@Test
public void testRestoreWithoutState() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new StatelessOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();
	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new StatelessOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 

@Test
public void testRestoreTailWithNewId() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		new OperatorID(44L, 44L),
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();
	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		new OperatorID(4444L, 4444L),
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(Collections.singleton(headOperatorID), RESTORED_OPERATORS);
}
 

@Test
public void testMultipleStatefulOperatorChainedSnapshotAndRestore() throws Exception {

	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator("head"),
		tailOperatorID,
		new CounterOperator("tail"),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();

	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator("head"),
		tailOperatorID,
		new CounterOperator("tail"),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 

@Test
public void testRestoreTailWithNewId() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		new OperatorID(44L, 44L),
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();
	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		new OperatorID(4444L, 4444L),
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(Collections.singleton(headOperatorID), RESTORED_OPERATORS);
}
 

@Override
public Future<Boolean> triggerCheckpointAsync(final CheckpointMetaData checkpointMetaData, final CheckpointOptions checkpointOptions, final boolean advanceToEndOfEventTime) {
	final TaskStateSnapshot checkpointStateHandles = new TaskStateSnapshot();
	checkpointStateHandles.putSubtaskStateByOperatorID(
		OperatorID.fromJobVertexID(getEnvironment().getJobVertexId()),
		new OperatorSubtaskState());

	getEnvironment().acknowledgeCheckpoint(
		checkpointMetaData.getCheckpointId(),
		new CheckpointMetrics(),
		checkpointStateHandles);

	triggerCheckpointLatch.countDown();

	return CompletableFuture.completedFuture(true);
}
 

public void restoreLatestCheckpointState(Map<Long, TaskStateSnapshot> taskStateSnapshotsByCheckpointId) {

		if (taskStateSnapshotsByCheckpointId == null
			|| taskStateSnapshotsByCheckpointId.isEmpty()) {
			return;
		}

		long latestId = -1;

		for (long id : taskStateSnapshotsByCheckpointId.keySet()) {
			if (id > latestId) {
				latestId = id;
			}
		}

		setReportedCheckpointId(latestId);
		setJobManagerTaskStateSnapshotsByCheckpointId(taskStateSnapshotsByCheckpointId);
	}
 

@Test
public void testMultipleStatefulOperatorChainedSnapshotAndRestore() throws Exception {

	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator("head"),
		tailOperatorID,
		new CounterOperator("tail"),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();

	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator("head"),
		tailOperatorID,
		new CounterOperator("tail"),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 

@Override
public boolean triggerCheckpoint(final CheckpointMetaData checkpointMetaData, final CheckpointOptions checkpointOptions, final boolean advanceToEndOfEventTime) {
	final TaskStateSnapshot checkpointStateHandles = new TaskStateSnapshot();
	checkpointStateHandles.putSubtaskStateByOperatorID(
		OperatorID.fromJobVertexID(getEnvironment().getJobVertexId()),
		new OperatorSubtaskState());

	getEnvironment().acknowledgeCheckpoint(
		checkpointMetaData.getCheckpointId(),
		new CheckpointMetrics(),
		checkpointStateHandles);

	triggerCheckpointLatch.countDown();

	return true;
}
 

/**
 * Test checkpoint pruning.
 */
@Test
public void pruneCheckpoints() throws Exception {

	final int chkCount = 3;

	List<TaskStateSnapshot> taskStateSnapshots = storeStates(chkCount);

	// test retrieve with pruning
	taskLocalStateStore.pruneMatchingCheckpoints((long chk) -> chk != chkCount - 1);

	for (int i = 0; i < chkCount - 1; ++i) {
		Assert.assertNull(taskLocalStateStore.retrieveLocalState(i));
	}

	checkStoredAsExpected(taskStateSnapshots, chkCount - 1, chkCount);
}
 

@Override
public void confirmCheckpoint(long confirmedCheckpointId) {
	Preconditions.checkState(!disposed);
	Iterator<Map.Entry<Long, TaskStateSnapshot>> iterator = taskStateSnapshotsByCheckpointID.entrySet().iterator();
	while (iterator.hasNext()) {
		Map.Entry<Long, TaskStateSnapshot> entry = iterator.next();
		if (entry.getKey() < confirmedCheckpointId) {
			iterator.remove();
			try {
				entry.getValue().discardState();
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
		} else {
			break;
		}
	}
}
 

@Override
public void acknowledgeCheckpoint(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	long checkpointId,
	CheckpointMetrics checkpointMetrics,
	TaskStateSnapshot subtaskState) {

	AcknowledgeReport acknowledgeReport = new AcknowledgeReport(
		jobID,
		executionAttemptID,
		checkpointId,
		checkpointMetrics,
		subtaskState);

	acknowledgeReports.add(acknowledgeReport);

	if (acknowledgeLatch != null) {
		acknowledgeLatch.trigger();
	}
}
 

@Test
public void testRestoreWithoutState() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new StatelessOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();
	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new StatelessOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 

/**
 * Pruning the useless checkpoints, it should be called only when holding the {@link #lock}.
 */
private void pruneCheckpoints(LongPredicate pruningChecker, boolean breakOnceCheckerFalse) {

	final List<Map.Entry<Long, TaskStateSnapshot>> toRemove = new ArrayList<>();

	synchronized (lock) {

		Iterator<Map.Entry<Long, TaskStateSnapshot>> entryIterator =
			storedTaskStateByCheckpointID.entrySet().iterator();

		while (entryIterator.hasNext()) {

			Map.Entry<Long, TaskStateSnapshot> snapshotEntry = entryIterator.next();
			long entryCheckpointId = snapshotEntry.getKey();

			if (pruningChecker.test(entryCheckpointId)) {
				toRemove.add(snapshotEntry);
				entryIterator.remove();
			} else if (breakOnceCheckerFalse) {
				break;
			}
		}
	}

	asyncDiscardLocalStateForCollection(toRemove);
}
 

/**
 * Let the checkpoint coordinator to receive all acknowledges from given executionVertexes so that to complete the expected checkpoint.
 */
private void acknowledgeAllCheckpoints(CheckpointCoordinator checkpointCoordinator, Iterator<ExecutionVertex> executionVertexes) throws IOException, CheckpointException {
	while (executionVertexes.hasNext()) {
		ExecutionVertex executionVertex = executionVertexes.next();
		for (int index = 0; index < executionVertex.getJobVertex().getParallelism(); index++) {
			JobVertexID jobVertexID = executionVertex.getJobvertexId();
			OperatorStateHandle opStateBackend = CheckpointCoordinatorTest.generatePartitionableStateHandle(jobVertexID, index, 2, 8, false);
			OperatorSubtaskState operatorSubtaskState = new OperatorSubtaskState(opStateBackend, null, null, null);
			TaskStateSnapshot taskOperatorSubtaskStates = new TaskStateSnapshot();
			taskOperatorSubtaskStates.putSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID), operatorSubtaskState);

			AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
				executionVertex.getJobId(),
				executionVertex.getJobVertex().getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
				checkpointId,
				new CheckpointMetrics(),
				taskOperatorSubtaskStates);

			checkpointCoordinator.receiveAcknowledgeMessage(acknowledgeCheckpoint, "Unknown location");
		}
	}
}
 

@Override
public void reportTaskStateSnapshots(
	@Nonnull CheckpointMetaData checkpointMetaData,
	@Nonnull CheckpointMetrics checkpointMetrics,
	@Nullable TaskStateSnapshot acknowledgedState,
	@Nullable TaskStateSnapshot localState) {

	long checkpointId = checkpointMetaData.getCheckpointId();

	localStateStore.storeLocalState(checkpointId, localState);

	checkpointResponder.acknowledgeCheckpoint(
		jobId,
		executionAttemptID,
		checkpointId,
		checkpointMetrics,
		acknowledgedState);
}
 

@VisibleForTesting
TaskLocalStateStoreImpl(
	@Nonnull JobID jobID,
	@Nonnull AllocationID allocationID,
	@Nonnull JobVertexID jobVertexID,
	@Nonnegative int subtaskIndex,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull Executor discardExecutor,
	@Nonnull SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID,
	@Nonnull Object lock) {

	this.jobID = jobID;
	this.allocationID = allocationID;
	this.jobVertexID = jobVertexID;
	this.subtaskIndex = subtaskIndex;
	this.discardExecutor = discardExecutor;
	this.localRecoveryConfig = localRecoveryConfig;
	this.storedTaskStateByCheckpointID = storedTaskStateByCheckpointID;
	this.lock = lock;
	this.disposed = false;
}
 

private List<TaskStateSnapshot> storeStates(int count) {
	List<TaskStateSnapshot> taskStateSnapshots = new ArrayList<>(count);
	for (int i = 0; i < count; ++i) {
		OperatorID operatorID = new OperatorID();
		TaskStateSnapshot taskStateSnapshot = spy(new TaskStateSnapshot());
		OperatorSubtaskState operatorSubtaskState = new OperatorSubtaskState();
		taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
		taskLocalStateStore.storeLocalState(i, taskStateSnapshot);
		taskStateSnapshots.add(taskStateSnapshot);
	}
	return taskStateSnapshots;
}
 

public AcknowledgeCheckpoint(
		JobID job,
		ExecutionAttemptID taskExecutionId,
		long checkpointId,
		CheckpointMetrics checkpointMetrics,
		TaskStateSnapshot subtaskState) {

	super(job, taskExecutionId, checkpointId);

	this.subtaskState = subtaskState;
	this.checkpointMetrics = checkpointMetrics;
}
 

@Override
public void storeLocalState(
	@Nonnegative long checkpointId,
	@Nullable TaskStateSnapshot localState) {

	if (localState == null) {
		localState = NULL_DUMMY;
	}

	if (LOG.isTraceEnabled()) {
		LOG.debug(
			"Stored local state for checkpoint {} in subtask ({} - {} - {}) : {}.",
			checkpointId, jobID, jobVertexID, subtaskIndex, localState);
	} else if (LOG.isDebugEnabled()) {
		LOG.debug(
			"Stored local state for checkpoint {} in subtask ({} - {} - {})",
			checkpointId, jobID, jobVertexID, subtaskIndex);
	}

	Map.Entry<Long, TaskStateSnapshot> toDiscard = null;

	synchronized (lock) {
		if (disposed) {
			// we ignore late stores and simply discard the state.
			toDiscard = new AbstractMap.SimpleEntry<>(checkpointId, localState);
		} else {
			TaskStateSnapshot previous =
				storedTaskStateByCheckpointID.put(checkpointId, localState);

			if (previous != null) {
				toDiscard = new AbstractMap.SimpleEntry<>(checkpointId, previous);
			}
		}
	}

	if (toDiscard != null) {
		asyncDiscardLocalStateForCollection(Collections.singletonList(toDiscard));
	}
}
 

private void checkStoredAsExpected(List<TaskStateSnapshot> history, int off, int len) throws Exception {
	for (int i = off; i < len; ++i) {
		TaskStateSnapshot expected = history.get(i);
		Assert.assertTrue(expected == taskLocalStateStore.retrieveLocalState(i));
		Mockito.verify(expected, Mockito.never()).discardState();
	}
}
 

private void reportCompletedSnapshotStates(
	TaskStateSnapshot acknowledgedTaskStateSnapshot,
	TaskStateSnapshot localTaskStateSnapshot,
	long asyncDurationMillis) {

	TaskStateManager taskStateManager = owner.getEnvironment().getTaskStateManager();

	boolean hasAckState = acknowledgedTaskStateSnapshot.hasState();
	boolean hasLocalState = localTaskStateSnapshot.hasState();

	Preconditions.checkState(hasAckState || !hasLocalState,
		"Found cached state but no corresponding primary state is reported to the job " +
			"manager. This indicates a problem.");

	// we signal stateless tasks by reporting null, so that there are no attempts to assign empty state
	// to stateless tasks on restore. This enables simple job modifications that only concern
	// stateless without the need to assign them uids to match their (always empty) states.
	taskStateManager.reportTaskStateSnapshots(
		checkpointMetaData,
		checkpointMetrics,
		hasAckState ? acknowledgedTaskStateSnapshot : null,
		hasLocalState ? localTaskStateSnapshot : null);

	LOG.debug("{} - finished asynchronous part of checkpoint {}. Asynchronous duration: {} ms",
		owner.getName(), checkpointMetaData.getCheckpointId(), asyncDurationMillis);

	LOG.trace("{} - reported the following states in snapshot for checkpoint {}: {}.",
		owner.getName(), checkpointMetaData.getCheckpointId(), acknowledgedTaskStateSnapshot);
}
 

public void dispose() {
	if (!disposed) {
		disposed = true;
		for (TaskStateSnapshot stateSnapshot : taskStateSnapshotsByCheckpointID.values()) {
			try {
				stateSnapshot.discardState();
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
		}
		taskStateSnapshotsByCheckpointID.clear();
	}
}
 

/**
 * Tests pruning of previous checkpoints if a new checkpoint is confirmed.
 */
@Test
public void confirmCheckpoint() throws Exception {

	final int chkCount = 3;
	final int confirmed = chkCount - 1;
	List<TaskStateSnapshot> taskStateSnapshots = storeStates(chkCount);
	taskLocalStateStore.confirmCheckpoint(confirmed);
	checkPrunedAndDiscarded(taskStateSnapshots, 0, confirmed);
	checkStoredAsExpected(taskStateSnapshots, confirmed, chkCount);
}
 

ChannelStateReaderImpl(TaskStateSnapshot snapshot, ChannelStateSerializer serializer) {
	RefCountingFSDataInputStreamFactory streamFactory = new RefCountingFSDataInputStreamFactory(serializer);
	final HashMap<InputChannelInfo, ChannelStateStreamReader> inputChannelHandleReadersTmp = new HashMap<>();
	final HashMap<ResultSubpartitionInfo, ChannelStateStreamReader> resultSubpartitionHandleReadersTmp = new HashMap<>();
	for (Map.Entry<OperatorID, OperatorSubtaskState> e : snapshot.getSubtaskStateMappings()) {
		addReaders(inputChannelHandleReadersTmp, e.getValue().getInputChannelState(), streamFactory);
		addReaders(resultSubpartitionHandleReadersTmp, e.getValue().getResultSubpartitionState(), streamFactory);
	}
	inputChannelHandleReaders = inputChannelHandleReadersTmp; // memory barrier to allow another thread call clear()
	resultSubpartitionHandleReaders = resultSubpartitionHandleReadersTmp; // memory barrier to allow another thread call clear()
}
 

public AcknowledgeCheckpoint(
		JobID job,
		ExecutionAttemptID taskExecutionId,
		long checkpointId,
		CheckpointMetrics checkpointMetrics,
		TaskStateSnapshot subtaskState) {

	super(job, taskExecutionId, checkpointId);

	this.subtaskState = subtaskState;
	this.checkpointMetrics = checkpointMetrics;
}
 

@Override
public void acknowledgeCheckpoint(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	long checkpointId,
	CheckpointMetrics checkpointMetrics,
	TaskStateSnapshot subtaskState) {

	throw new RuntimeException("Unexpected call.");
}
 

@Test
public void testConfirmTaskCheckpointed() {
	try {
		AcknowledgeCheckpoint noState = new AcknowledgeCheckpoint(
				new JobID(), new ExecutionAttemptID(), 569345L);

		KeyGroupRange keyGroupRange = KeyGroupRange.of(42,42);

		TaskStateSnapshot checkpointStateHandles = new TaskStateSnapshot();
		checkpointStateHandles.putSubtaskStateByOperatorID(
			new OperatorID(),
			new OperatorSubtaskState(
				CheckpointCoordinatorTest.generatePartitionableStateHandle(new JobVertexID(), 0, 2, 8, false),
				null,
				CheckpointCoordinatorTest.generateKeyGroupState(keyGroupRange, Collections.singletonList(new MyHandle())),
				null
			)
		);

		AcknowledgeCheckpoint withState = new AcknowledgeCheckpoint(
				new JobID(),
				new ExecutionAttemptID(),
				87658976143L,
				new CheckpointMetrics(),
				checkpointStateHandles);

		testSerializabilityEqualsHashCode(noState);
		testSerializabilityEqualsHashCode(withState);
	} catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#initializeState()}.
 * Calls {@link org.apache.flink.streaming.api.operators.SetupableStreamOperator#setup(StreamTask, StreamConfig, Output)}
 * if it was not called before.
 *
 * @param jmOperatorStateHandles the primary state (owned by JM)
 * @param tmOperatorStateHandles the (optional) local state (owned by TM) or null.
 * @throws Exception
 */
public void initializeState(
	OperatorSubtaskState jmOperatorStateHandles,
	OperatorSubtaskState tmOperatorStateHandles) throws Exception {

	checkState(!initializeCalled, "TestHarness has already been initialized. Have you " +
		"opened this harness before initializing it?");
	if (!setupCalled) {
		setup();
	}

	if (jmOperatorStateHandles != null) {

		TaskStateSnapshot jmTaskStateSnapshot = new TaskStateSnapshot();
		jmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), jmOperatorStateHandles);

		taskStateManager.setReportedCheckpointId(0);
		taskStateManager.setJobManagerTaskStateSnapshotsByCheckpointId(
			Collections.singletonMap(0L, jmTaskStateSnapshot));

		if (tmOperatorStateHandles != null) {
			TaskStateSnapshot tmTaskStateSnapshot = new TaskStateSnapshot();
			tmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), tmOperatorStateHandles);
			taskStateManager.setTaskManagerTaskStateSnapshotsByCheckpointId(
				Collections.singletonMap(0L, tmTaskStateSnapshot));
		}
	}

	operator.initializeState(mockTask.createStreamTaskStateInitializer());
	initializeCalled = true;
}
 

/**
 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#initializeState()}.
 * Calls {@link org.apache.flink.streaming.api.operators.SetupableStreamOperator#setup(StreamTask, StreamConfig, Output)}
 * if it was not called before.
 *
 * @param jmOperatorStateHandles the primary state (owned by JM)
 * @param tmOperatorStateHandles the (optional) local state (owned by TM) or null.
 * @throws Exception
 */
public void initializeState(
	OperatorSubtaskState jmOperatorStateHandles,
	OperatorSubtaskState tmOperatorStateHandles) throws Exception {

	checkState(!initializeCalled, "TestHarness has already been initialized. Have you " +
		"opened this harness before initializing it?");
	if (!setupCalled) {
		setup();
	}

	if (jmOperatorStateHandles != null) {

		TaskStateSnapshot jmTaskStateSnapshot = new TaskStateSnapshot();
		jmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), jmOperatorStateHandles);

		taskStateManager.setReportedCheckpointId(0);
		taskStateManager.setJobManagerTaskStateSnapshotsByCheckpointId(
			Collections.singletonMap(0L, jmTaskStateSnapshot));

		if (tmOperatorStateHandles != null) {
			TaskStateSnapshot tmTaskStateSnapshot = new TaskStateSnapshot();
			tmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), tmOperatorStateHandles);
			taskStateManager.setTaskManagerTaskStateSnapshotsByCheckpointId(
				Collections.singletonMap(0L, tmTaskStateSnapshot));
		}
	}

	operator.initializeState();
	initializeCalled = true;
}
 

public void dispose() {
	if (!disposed) {
		disposed = true;
		for (TaskStateSnapshot stateSnapshot : taskStateSnapshotsByCheckpointID.values()) {
			try {
				stateSnapshot.discardState();
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
		}
		taskStateSnapshotsByCheckpointID.clear();
	}
}