org.apache.flink.runtime.state.TaskStateManager Java Examples

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {
	this(
		new JobID(),
		new ExecutionAttemptID(0L, 0L),
		jobConfig,
		taskConfig,
		executionConfig,
		memorySize,
		inputSplitProvider,
		bufferSize,
		taskStateManager);
}
 

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {
	this(
		new JobID(),
		new ExecutionAttemptID(0L, 0L),
		jobConfig,
		taskConfig,
		executionConfig,
		memorySize,
		inputSplitProvider,
		bufferSize,
		taskStateManager);
}
 

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {
	this(
		new JobID(),
		new ExecutionAttemptID(0L, 0L),
		jobConfig,
		taskConfig,
		executionConfig,
		memorySize,
		inputSplitProvider,
		bufferSize,
		taskStateManager);
}
 

public StreamMockEnvironment(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this.jobID = jobID;
	this.executionAttemptID = executionAttemptID;

	int subtaskIndex = 0;
	this.taskInfo = new TaskInfo(
		"", /* task name */
		1, /* num key groups / max parallelism */
		subtaskIndex, /* index of this subtask */
		1, /* num subtasks */
		0 /* attempt number */);
	this.jobConfiguration = jobConfig;
	this.taskConfiguration = taskConfig;
	this.inputs = new LinkedList<InputGate>();
	this.outputs = new LinkedList<ResultPartitionWriter>();
	this.memManager = new MemoryManager(memorySize, 1);
	this.ioManager = new IOManagerAsync();
	this.taskStateManager = Preconditions.checkNotNull(taskStateManager);
	this.aggregateManager = new TestGlobalAggregateManager();
	this.inputSplitProvider = inputSplitProvider;
	this.bufferSize = bufferSize;

	this.executionConfig = executionConfig;
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId());

	KvStateRegistry registry = new KvStateRegistry();
	this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	LocalRecoveryConfig localRecoveryConfig = taskStateManager.createLocalRecoveryConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);

	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		localRecoveryConfig,
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		taskStateManager.createLocalRecoveryConfig(),
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		taskStateManager.createLocalRecoveryConfig(),
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	LocalRecoveryConfig localRecoveryConfig = taskStateManager.createLocalRecoveryConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);

	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		localRecoveryConfig,
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Nullable
static StreamStateHandle readSnapshot(TaskStateManager stateManager, OperatorID operatorId) {
	final PrioritizedOperatorSubtaskState poss = stateManager.prioritizedOperatorState(operatorId);
	if (!poss.isRestored()) {
		return null;
	}

	final StateObjectCollection<OperatorStateHandle> opState =
		stateManager.prioritizedOperatorState(operatorId).getPrioritizedManagedOperatorState().get(0);
	final OperatorStateHandle handle = Iterators.getOnlyElement(opState.iterator());
	return handle.getDelegateStateHandle();
}
 

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this(jobConfig, taskConfig, new ExecutionConfig(), memorySize, inputSplitProvider, bufferSize, taskStateManager);
}
 

public StreamMockEnvironment(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this.jobID = jobID;
	this.executionAttemptID = executionAttemptID;

	int subtaskIndex = 0;
	this.taskInfo = new TaskInfo(
		"", /* task name */
		1, /* num key groups / max parallelism */
		subtaskIndex, /* index of this subtask */
		1, /* num subtasks */
		0 /* attempt number */);
	this.jobConfiguration = jobConfig;
	this.taskConfiguration = taskConfig;
	this.inputs = new LinkedList<InputGate>();
	this.outputs = new LinkedList<ResultPartitionWriter>();
	this.memManager = new MemoryManager(memorySize, 1);
	this.ioManager = new IOManagerAsync();
	this.taskStateManager = Preconditions.checkNotNull(taskStateManager);
	this.aggregateManager = new TestGlobalAggregateManager();
	this.inputSplitProvider = inputSplitProvider;
	this.bufferSize = bufferSize;

	this.executionConfig = executionConfig;
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId());

	KvStateRegistry registry = new KvStateRegistry();
	this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
}
 

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

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

private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				ProcessingTimeService processingTimeService,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 

public StreamMockEnvironment(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long offHeapMemorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this.jobID = jobID;
	this.executionAttemptID = executionAttemptID;

	int subtaskIndex = 0;
	this.taskInfo = new TaskInfo(
		"", /* task name */
		1, /* num key groups / max parallelism */
		subtaskIndex, /* index of this subtask */
		1, /* num subtasks */
		0 /* attempt number */);
	this.jobConfiguration = jobConfig;
	this.taskConfiguration = taskConfig;
	this.inputs = new LinkedList<>();
	this.outputs = new LinkedList<ResultPartitionWriter>();
	this.memManager = MemoryManagerBuilder.newBuilder().setMemorySize(offHeapMemorySize).build();
	this.ioManager = new IOManagerAsync();
	this.taskStateManager = Preconditions.checkNotNull(taskStateManager);
	this.aggregateManager = new TestGlobalAggregateManager();
	this.inputSplitProvider = inputSplitProvider;
	this.bufferSize = bufferSize;

	this.executionConfig = executionConfig;
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId());

	KvStateRegistry registry = new KvStateRegistry();
	this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
}
 

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this(jobConfig, taskConfig, new ExecutionConfig(), memorySize, inputSplitProvider, bufferSize, taskStateManager);
}
 

public StreamMockEnvironment(
	Configuration jobConfig,
	Configuration taskConfig,
	long memorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this(jobConfig, taskConfig, new ExecutionConfig(), memorySize, inputSplitProvider, bufferSize, taskStateManager);
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		taskStateManager.createLocalRecoveryConfig(),
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	LocalRecoveryConfig localRecoveryConfig = taskStateManager.createLocalRecoveryConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);

	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		localRecoveryConfig,
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public TaskStateManager getTaskStateManager() {
	return taskStateManager;
}
 

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

@Override
public TaskStateManager getTaskStateManager() {
	return taskStateManager;
}
 

public RuntimeEnvironment(
		JobID jobId,
		JobVertexID jobVertexId,
		ExecutionAttemptID executionId,
		ExecutionConfig executionConfig,
		TaskInfo taskInfo,
		Configuration jobConfiguration,
		Configuration taskConfiguration,
		ClassLoader userCodeClassLoader,
		MemoryManager memManager,
		IOManager ioManager,
		BroadcastVariableManager bcVarManager,
		TaskStateManager taskStateManager,
		GlobalAggregateManager aggregateManager,
		AccumulatorRegistry accumulatorRegistry,
		TaskKvStateRegistry kvStateRegistry,
		InputSplitProvider splitProvider,
		Map<String, Future<Path>> distCacheEntries,
		ResultPartitionWriter[] writers,
		IndexedInputGate[] inputGates,
		TaskEventDispatcher taskEventDispatcher,
		CheckpointResponder checkpointResponder,
		TaskOperatorEventGateway operatorEventGateway,
		TaskManagerRuntimeInfo taskManagerInfo,
		TaskMetricGroup metrics,
		Task containingTask,
		ExternalResourceInfoProvider externalResourceInfoProvider) {

	this.jobId = checkNotNull(jobId);
	this.jobVertexId = checkNotNull(jobVertexId);
	this.executionId = checkNotNull(executionId);
	this.taskInfo = checkNotNull(taskInfo);
	this.executionConfig = checkNotNull(executionConfig);
	this.jobConfiguration = checkNotNull(jobConfiguration);
	this.taskConfiguration = checkNotNull(taskConfiguration);
	this.userCodeClassLoader = checkNotNull(userCodeClassLoader);
	this.memManager = checkNotNull(memManager);
	this.ioManager = checkNotNull(ioManager);
	this.bcVarManager = checkNotNull(bcVarManager);
	this.taskStateManager = checkNotNull(taskStateManager);
	this.aggregateManager = checkNotNull(aggregateManager);
	this.accumulatorRegistry = checkNotNull(accumulatorRegistry);
	this.kvStateRegistry = checkNotNull(kvStateRegistry);
	this.splitProvider = checkNotNull(splitProvider);
	this.distCacheEntries = checkNotNull(distCacheEntries);
	this.writers = checkNotNull(writers);
	this.inputGates = checkNotNull(inputGates);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
	this.checkpointResponder = checkNotNull(checkpointResponder);
	this.operatorEventGateway = checkNotNull(operatorEventGateway);
	this.taskManagerInfo = checkNotNull(taskManagerInfo);
	this.containingTask = containingTask;
	this.metrics = metrics;
	this.externalResourceInfoProvider = checkNotNull(externalResourceInfoProvider);
}
 

@Override
public TaskStateManager getTaskStateManager() {
	return taskStateManager;
}
 

public MockEnvironmentBuilder setTaskStateManager(TaskStateManager taskStateManager) {
	this.taskStateManager = taskStateManager;
	return this;
}
 

protected MockEnvironment(
		JobID jobID,
		JobVertexID jobVertexID,
		String taskName,
		MockInputSplitProvider inputSplitProvider,
		int bufferSize,
		Configuration taskConfiguration,
		ExecutionConfig executionConfig,
		IOManager ioManager,
		TaskStateManager taskStateManager,
		GlobalAggregateManager aggregateManager,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		ClassLoader userCodeClassLoader,
		TaskMetricGroup taskMetricGroup,
		TaskManagerRuntimeInfo taskManagerRuntimeInfo,
		MemoryManager memManager,
		ExternalResourceInfoProvider externalResourceInfoProvider) {

	this.jobID = jobID;
	this.jobVertexID = jobVertexID;

	this.taskInfo = new TaskInfo(taskName, maxParallelism, subtaskIndex, parallelism, 0);
	this.jobConfiguration = new Configuration();
	this.taskConfiguration = taskConfiguration;
	this.inputs = new LinkedList<>();
	this.outputs = new LinkedList<ResultPartitionWriter>();

	this.memManager = memManager;
	this.ioManager = ioManager;
	this.taskManagerRuntimeInfo = taskManagerRuntimeInfo;

	this.executionConfig = executionConfig;
	this.inputSplitProvider = inputSplitProvider;
	this.bufferSize = bufferSize;

	this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId());

	this.kvStateRegistry = new KvStateRegistry();
	this.taskKvStateRegistry = kvStateRegistry.createTaskRegistry(jobID, getJobVertexId());

	this.userCodeClassLoader = Preconditions.checkNotNull(userCodeClassLoader);
	this.taskStateManager = Preconditions.checkNotNull(taskStateManager);
	this.aggregateManager = Preconditions.checkNotNull(aggregateManager);

	this.taskMetricGroup = taskMetricGroup;

	this.externalResourceInfoProvider = Preconditions.checkNotNull(externalResourceInfoProvider);
}
 

@Override
public TaskStateManager getTaskStateManager() {
	return taskStateManager;
}
 

public void setTaskStateManager(TaskStateManager taskStateManager) {
	this.taskStateManager = taskStateManager;
}
 

@Override
public TaskStateManager getTaskStateManager() {
	return taskStateManager;
}
 

public void setTaskStateManager(TaskStateManager taskStateManager) {
	this.taskStateManager = taskStateManager;
}