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

public void serializeOld(SavepointV1 savepoint, DataOutputStream dos) throws IOException {
	dos.writeLong(savepoint.getCheckpointId());

	Collection<TaskState> taskStates = savepoint.getTaskStates();
	dos.writeInt(taskStates.size());

	for (TaskState taskState : savepoint.getTaskStates()) {
		// Vertex ID
		dos.writeLong(taskState.getJobVertexID().getLowerPart());
		dos.writeLong(taskState.getJobVertexID().getUpperPart());

		// Parallelism
		int parallelism = taskState.getParallelism();
		dos.writeInt(parallelism);
		dos.writeInt(taskState.getMaxParallelism());
		dos.writeInt(taskState.getChainLength());

		// Sub task states
		Map<Integer, SubtaskState> subtaskStateMap = taskState.getSubtaskStates();
		dos.writeInt(subtaskStateMap.size());
		for (Map.Entry<Integer, SubtaskState> entry : subtaskStateMap.entrySet()) {
			dos.writeInt(entry.getKey());
			serializeSubtaskState(entry.getValue(), dos);
		}
	}
}
 

public void serializeOld(SavepointV1 savepoint, DataOutputStream dos) throws IOException {
	dos.writeLong(savepoint.getCheckpointId());

	Collection<TaskState> taskStates = savepoint.getTaskStates();
	dos.writeInt(taskStates.size());

	for (TaskState taskState : savepoint.getTaskStates()) {
		// Vertex ID
		dos.writeLong(taskState.getJobVertexID().getLowerPart());
		dos.writeLong(taskState.getJobVertexID().getUpperPart());

		// Parallelism
		int parallelism = taskState.getParallelism();
		dos.writeInt(parallelism);
		dos.writeInt(taskState.getMaxParallelism());
		dos.writeInt(taskState.getChainLength());

		// Sub task states
		Map<Integer, SubtaskState> subtaskStateMap = taskState.getSubtaskStates();
		dos.writeInt(subtaskStateMap.size());
		for (Map.Entry<Integer, SubtaskState> entry : subtaskStateMap.entrySet()) {
			dos.writeInt(entry.getKey());
			serializeSubtaskState(entry.getValue(), dos);
		}
	}
}
 

@Override
public SavepointV2 deserialize(DataInputStream dis, ClassLoader cl) throws IOException {
	long checkpointId = dis.readLong();

	// Task states
	int numTaskStates = dis.readInt();
	List<TaskState> taskStates = new ArrayList<>(numTaskStates);

	for (int i = 0; i < numTaskStates; i++) {
		JobVertexID jobVertexId = new JobVertexID(dis.readLong(), dis.readLong());
		int parallelism = dis.readInt();
		int maxParallelism = dis.readInt();
		int chainLength = dis.readInt();

		// Add task state
		TaskState taskState = new TaskState(jobVertexId, parallelism, maxParallelism, chainLength);
		taskStates.add(taskState);

		// Sub task states
		int numSubTaskStates = dis.readInt();

		for (int j = 0; j < numSubTaskStates; j++) {
			int subtaskIndex = dis.readInt();
			SubtaskState subtaskState = deserializeSubtaskState(dis);
			taskState.putState(subtaskIndex, subtaskState);
		}
	}

	return new SavepointV2(checkpointId, taskStates);
}
 

@Override
public SavepointV2 deserialize(DataInputStream dis, ClassLoader cl) throws IOException {
	long checkpointId = dis.readLong();

	// Task states
	int numTaskStates = dis.readInt();
	List<TaskState> taskStates = new ArrayList<>(numTaskStates);

	for (int i = 0; i < numTaskStates; i++) {
		JobVertexID jobVertexId = new JobVertexID(dis.readLong(), dis.readLong());
		int parallelism = dis.readInt();
		int maxParallelism = dis.readInt();
		int chainLength = dis.readInt();

		// Add task state
		TaskState taskState = new TaskState(jobVertexId, parallelism, maxParallelism, chainLength);
		taskStates.add(taskState);

		// Sub task states
		int numSubTaskStates = dis.readInt();

		for (int j = 0; j < numSubTaskStates; j++) {
			int subtaskIndex = dis.readInt();
			SubtaskState subtaskState = deserializeSubtaskState(dis);
			taskState.putState(subtaskIndex, subtaskState);
		}
	}

	return new SavepointV2(checkpointId, taskStates);
}
 

/**
 * Creates a random collection of TaskState objects containing various types of state handles.
 */
public static Collection<TaskState> createTaskStates(
		Random random,
		int numTaskStates,
		int numSubtasksPerTask) {

	List<TaskState> taskStates = new ArrayList<>(numTaskStates);

	for (int stateIdx = 0; stateIdx < numTaskStates; ++stateIdx) {

		int chainLength = 1 + random.nextInt(8);

		TaskState taskState = new TaskState(new JobVertexID(), numSubtasksPerTask, 128, chainLength);

		int noNonPartitionableStateAtIndex = random.nextInt(chainLength);
		int noOperatorStateBackendAtIndex = random.nextInt(chainLength);
		int noOperatorStateStreamAtIndex = random.nextInt(chainLength);

		boolean hasKeyedBackend = random.nextInt(4) != 0;
		boolean hasKeyedStream = random.nextInt(4) != 0;

		for (int subtaskIdx = 0; subtaskIdx < numSubtasksPerTask; subtaskIdx++) {

			List<OperatorStateHandle> operatorStatesBackend = new ArrayList<>(chainLength);
			List<OperatorStateHandle> operatorStatesStream = new ArrayList<>(chainLength);

			for (int chainIdx = 0; chainIdx < chainLength; ++chainIdx) {

				StreamStateHandle operatorStateBackend =
						new ByteStreamStateHandle("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
				StreamStateHandle operatorStateStream =
						new ByteStreamStateHandle("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
				Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
				offsetsMap.put("A", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
				offsetsMap.put("B", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
				offsetsMap.put("C", new OperatorStateHandle.StateMetaInfo(new long[]{60, 70, 80}, OperatorStateHandle.Mode.UNION));

				if (chainIdx != noOperatorStateBackendAtIndex) {
					OperatorStateHandle operatorStateHandleBackend =
							new OperatorStreamStateHandle(offsetsMap, operatorStateBackend);
					operatorStatesBackend.add(operatorStateHandleBackend);
				}

				if (chainIdx != noOperatorStateStreamAtIndex) {
					OperatorStateHandle operatorStateHandleStream =
							new OperatorStreamStateHandle(offsetsMap, operatorStateStream);
					operatorStatesStream.add(operatorStateHandleStream);
				}
			}

			KeyGroupsStateHandle keyedStateBackend = null;
			KeyGroupsStateHandle keyedStateStream = null;

			if (hasKeyedBackend) {
				keyedStateBackend = createDummyKeyGroupStateHandle(random);
			}

			if (hasKeyedStream) {
				keyedStateStream = createDummyKeyGroupStateHandle(random);
			}

			taskState.putState(subtaskIdx, new SubtaskState(
					new ChainedStateHandle<>(operatorStatesBackend),
					new ChainedStateHandle<>(operatorStatesStream),
					keyedStateStream,
					keyedStateBackend));
		}

		taskStates.add(taskState);
	}

	return taskStates;
}
 

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	Environment mockEnvironment = spy(new MockEnvironmentBuilder().build());

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	when(mockEnvironment.getTaskStateManager()).thenReturn(taskStateManager);

	StreamTask<?, ?> streamTask = new EmptyStreamTask(mockEnvironment);
	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);

	// mock the operators
	StreamOperator<?> statelessOperator =
			mock(StreamOperator.class);

	final OperatorID operatorID = new OperatorID();
	when(statelessOperator.getOperatorID()).thenReturn(operatorID);

	// mock the returned empty snapshot result (all state handles are null)
	OperatorSnapshotFutures statelessOperatorSnapshotResult = new OperatorSnapshotFutures();
	when(statelessOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class)))
			.thenReturn(statelessOperatorSnapshotResult);

	// set up the task
	StreamOperator<?>[] streamOperators = {statelessOperator};
	OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
	when(operatorChain.getAllOperators()).thenReturn(streamOperators);

	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, "configuration", new StreamConfig(new Configuration()));
	Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", Executors.newCachedThreadPool());
	Whitebox.setInternalState(streamTask, "checkpointStorage", new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE));

	streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation());
	checkpointCompletedLatch.await(30, TimeUnit.SECONDS);
	streamTask.cancel();

	// ensure that 'null' was acknowledged as subtask state
	Assert.assertNull(checkpointResult.get(0));
}
 

/**
 * Creates a random collection of TaskState objects containing various types of state handles.
 */
public static Collection<TaskState> createTaskStates(
		Random random,
		int numTaskStates,
		int numSubtasksPerTask) {

	List<TaskState> taskStates = new ArrayList<>(numTaskStates);

	for (int stateIdx = 0; stateIdx < numTaskStates; ++stateIdx) {

		int chainLength = 1 + random.nextInt(8);

		TaskState taskState = new TaskState(new JobVertexID(), numSubtasksPerTask, 128, chainLength);

		int noNonPartitionableStateAtIndex = random.nextInt(chainLength);
		int noOperatorStateBackendAtIndex = random.nextInt(chainLength);
		int noOperatorStateStreamAtIndex = random.nextInt(chainLength);

		boolean hasKeyedBackend = random.nextInt(4) != 0;
		boolean hasKeyedStream = random.nextInt(4) != 0;

		for (int subtaskIdx = 0; subtaskIdx < numSubtasksPerTask; subtaskIdx++) {

			List<OperatorStateHandle> operatorStatesBackend = new ArrayList<>(chainLength);
			List<OperatorStateHandle> operatorStatesStream = new ArrayList<>(chainLength);

			for (int chainIdx = 0; chainIdx < chainLength; ++chainIdx) {

				StreamStateHandle operatorStateBackend =
						new ByteStreamStateHandle("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
				StreamStateHandle operatorStateStream =
						new ByteStreamStateHandle("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
				Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
				offsetsMap.put("A", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
				offsetsMap.put("B", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
				offsetsMap.put("C", new OperatorStateHandle.StateMetaInfo(new long[]{60, 70, 80}, OperatorStateHandle.Mode.UNION));

				if (chainIdx != noOperatorStateBackendAtIndex) {
					OperatorStateHandle operatorStateHandleBackend =
							new OperatorStreamStateHandle(offsetsMap, operatorStateBackend);
					operatorStatesBackend.add(operatorStateHandleBackend);
				}

				if (chainIdx != noOperatorStateStreamAtIndex) {
					OperatorStateHandle operatorStateHandleStream =
							new OperatorStreamStateHandle(offsetsMap, operatorStateStream);
					operatorStatesStream.add(operatorStateHandleStream);
				}
			}

			KeyGroupsStateHandle keyedStateBackend = null;
			KeyGroupsStateHandle keyedStateStream = null;

			if (hasKeyedBackend) {
				keyedStateBackend = createDummyKeyGroupStateHandle(random);
			}

			if (hasKeyedStream) {
				keyedStateStream = createDummyKeyGroupStateHandle(random);
			}

			taskState.putState(subtaskIdx, new SubtaskState(
					new ChainedStateHandle<>(operatorStatesBackend),
					new ChainedStateHandle<>(operatorStatesStream),
					keyedStateStream,
					keyedStateBackend));
		}

		taskStates.add(taskState);
	}

	return taskStates;
}
 

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	StreamOperator<?> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpoint(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}
 

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	OneInputStreamOperator<String, String> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpointAsync(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}