org.apache.flink.runtime.io.network.partition.ResultPartitionConsumableNotifier Java Examples

public JobManagerConnection(
			JobID jobID,
			ResourceID resourceID,
			JobMasterGateway jobMasterGateway,
			TaskManagerActions taskManagerActions,
			CheckpointResponder checkpointResponder,
			GlobalAggregateManager aggregateManager,
			LibraryCacheManager libraryCacheManager,
			ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
			PartitionProducerStateChecker partitionStateChecker) {
	this.jobID = Preconditions.checkNotNull(jobID);
	this.resourceID = Preconditions.checkNotNull(resourceID);
	this.jobMasterGateway = Preconditions.checkNotNull(jobMasterGateway);
	this.taskManagerActions = Preconditions.checkNotNull(taskManagerActions);
	this.checkpointResponder = Preconditions.checkNotNull(checkpointResponder);
	this.aggregateManager = Preconditions.checkNotNull(aggregateManager);
	this.libraryCacheManager = Preconditions.checkNotNull(libraryCacheManager);
	this.resultPartitionConsumableNotifier = Preconditions.checkNotNull(resultPartitionConsumableNotifier);
	this.partitionStateChecker = Preconditions.checkNotNull(partitionStateChecker);
}
 

public static ResultPartitionWriter[] decorate(
		Collection<ResultPartitionDeploymentDescriptor> descs,
		ResultPartitionWriter[] partitionWriters,
		TaskActions taskActions,
		JobID jobId,
		ResultPartitionConsumableNotifier notifier) {

	ResultPartitionWriter[] consumableNotifyingPartitionWriters = new ResultPartitionWriter[partitionWriters.length];
	int counter = 0;
	for (ResultPartitionDeploymentDescriptor desc : descs) {
		if (desc.sendScheduleOrUpdateConsumersMessage() && desc.getPartitionType().isPipelined()) {
			consumableNotifyingPartitionWriters[counter] = new ConsumableNotifyingResultPartitionWriterDecorator(
				taskActions,
				jobId,
				partitionWriters[counter],
				notifier);
		} else {
			consumableNotifyingPartitionWriters[counter] = partitionWriters[counter];
		}
		counter++;
	}
	return consumableNotifyingPartitionWriters;
}
 

@Override
public JobTable.Connection connect(
		ResourceID resourceId,
		JobMasterGateway jobMasterGateway,
		TaskManagerActions taskManagerActions,
		CheckpointResponder checkpointResponder,
		GlobalAggregateManager aggregateManager,
		ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
		PartitionProducerStateChecker partitionStateChecker) {
	verifyJobIsNotClosed();
	Preconditions.checkState(connection == null);

	connection = new EstablishedConnection(
		resourceId,
		jobMasterGateway,
		taskManagerActions,
		checkpointResponder,
		aggregateManager,
		resultPartitionConsumableNotifier,
		partitionStateChecker);
	resourceIdJobIdIndex.put(resourceId, jobId);

	return this;
}
 

private EstablishedConnection(
	ResourceID resourceID,
	JobMasterGateway jobMasterGateway,
	TaskManagerActions taskManagerActions,
	CheckpointResponder checkpointResponder,
	GlobalAggregateManager globalAggregateManager,
	ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
	PartitionProducerStateChecker partitionStateChecker) {
	this.resourceID = Preconditions.checkNotNull(resourceID);
	this.jobMasterGateway = Preconditions.checkNotNull(jobMasterGateway);
	this.taskManagerActions = Preconditions.checkNotNull(taskManagerActions);
	this.checkpointResponder = Preconditions.checkNotNull(checkpointResponder);
	this.globalAggregateManager = Preconditions.checkNotNull(globalAggregateManager);
	this.resultPartitionConsumableNotifier = Preconditions.checkNotNull(resultPartitionConsumableNotifier);
	this.partitionStateChecker = Preconditions.checkNotNull(partitionStateChecker);
}
 

public JobManagerConnection(
			JobID jobID,
			ResourceID resourceID,
			JobMasterGateway jobMasterGateway,
			TaskManagerActions taskManagerActions,
			CheckpointResponder checkpointResponder,
			GlobalAggregateManager aggregateManager,
			LibraryCacheManager libraryCacheManager,
			ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
			PartitionProducerStateChecker partitionStateChecker) {
	this.jobID = Preconditions.checkNotNull(jobID);
	this.resourceID = Preconditions.checkNotNull(resourceID);
	this.jobMasterGateway = Preconditions.checkNotNull(jobMasterGateway);
	this.taskManagerActions = Preconditions.checkNotNull(taskManagerActions);
	this.checkpointResponder = Preconditions.checkNotNull(checkpointResponder);
	this.aggregateManager = Preconditions.checkNotNull(aggregateManager);
	this.libraryCacheManager = Preconditions.checkNotNull(libraryCacheManager);
	this.resultPartitionConsumableNotifier = Preconditions.checkNotNull(resultPartitionConsumableNotifier);
	this.partitionStateChecker = Preconditions.checkNotNull(partitionStateChecker);
}
 

public static ResultPartitionWriter[] decorate(
		Collection<ResultPartitionDeploymentDescriptor> descs,
		ResultPartitionWriter[] partitionWriters,
		TaskActions taskActions,
		JobID jobId,
		ResultPartitionConsumableNotifier notifier) {

	ResultPartitionWriter[] consumableNotifyingPartitionWriters = new ResultPartitionWriter[partitionWriters.length];
	int counter = 0;
	for (ResultPartitionDeploymentDescriptor desc : descs) {
		if (desc.sendScheduleOrUpdateConsumersMessage() && desc.getPartitionType().isPipelined()) {
			consumableNotifyingPartitionWriters[counter] = new ConsumableNotifyingResultPartitionWriterDecorator(
				taskActions,
				jobId,
				partitionWriters[counter],
				notifier);
		} else {
			consumableNotifyingPartitionWriters[counter] = partitionWriters[counter];
		}
		counter++;
	}
	return consumableNotifyingPartitionWriters;
}
 

public ConsumableNotifyingResultPartitionWriterDecorator(
		TaskActions taskActions,
		JobID jobId,
		ResultPartitionWriter partitionWriter,
		ResultPartitionConsumableNotifier partitionConsumableNotifier) {
	this.taskActions = checkNotNull(taskActions);
	this.jobId = checkNotNull(jobId);
	this.partitionWriter = checkNotNull(partitionWriter);
	this.partitionConsumableNotifier = checkNotNull(partitionConsumableNotifier);
}
 

private JobManagerConnection associateWithJobManager(
		JobID jobID,
		ResourceID resourceID,
		JobMasterGateway jobMasterGateway) {
	checkNotNull(jobID);
	checkNotNull(resourceID);
	checkNotNull(jobMasterGateway);

	TaskManagerActions taskManagerActions = new TaskManagerActionsImpl(jobMasterGateway);

	CheckpointResponder checkpointResponder = new RpcCheckpointResponder(jobMasterGateway);
	GlobalAggregateManager aggregateManager = new RpcGlobalAggregateManager(jobMasterGateway);

	final LibraryCacheManager libraryCacheManager = new BlobLibraryCacheManager(
		blobCacheService.getPermanentBlobService(),
		taskManagerConfiguration.getClassLoaderResolveOrder(),
		taskManagerConfiguration.getAlwaysParentFirstLoaderPatterns());

	ResultPartitionConsumableNotifier resultPartitionConsumableNotifier = new RpcResultPartitionConsumableNotifier(
		jobMasterGateway,
		getRpcService().getExecutor(),
		taskManagerConfiguration.getTimeout());

	PartitionProducerStateChecker partitionStateChecker = new RpcPartitionStateChecker(jobMasterGateway);

	registerQueryableState(jobID, jobMasterGateway);

	return new JobManagerConnection(
		jobID,
		resourceID,
		jobMasterGateway,
		taskManagerActions,
		checkpointResponder,
		aggregateManager,
		libraryCacheManager,
		resultPartitionConsumableNotifier,
		partitionStateChecker);
}
 

@Test
public void testExecutionFailsInNetworkRegistration() throws Exception {
	// mock a network manager that rejects registration
	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	final ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	final PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);

	final NetworkEnvironment network = mock(NetworkEnvironment.class);
	when(network.getResultPartitionManager()).thenReturn(partitionManager);
	when(network.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC);
	when(network.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);
	doThrow(new RuntimeException("buffers")).when(network).registerTask(any(Task.class));

	final QueuedNoOpTaskManagerActions taskManagerActions = new QueuedNoOpTaskManagerActions();
	final Task task = new TaskBuilder()
		.setTaskManagerActions(taskManagerActions)
		.setConsumableNotifier(consumableNotifier)
		.setPartitionProducerStateChecker(partitionProducerStateChecker)
		.setNetworkEnvironment(network)
		.build();

	// should fail
	task.run();

	// verify final state
	assertEquals(ExecutionState.FAILED, task.getExecutionState());
	assertTrue(task.isCanceledOrFailed());
	assertTrue(task.getFailureCause().getMessage().contains("buffers"));

	taskManagerActions.validateListenerMessage(
		ExecutionState.FAILED, task, new RuntimeException("buffers"));
}
 

private void testExecutionFailsInNetworkRegistration(
		List<ResultPartitionDeploymentDescriptor> resultPartitions,
		List<InputGateDeploymentDescriptor> inputGates) throws Exception {
	final String errorMessage = "Network buffer pool has already been destroyed.";

	final ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	final PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);

	final QueuedNoOpTaskManagerActions taskManagerActions = new QueuedNoOpTaskManagerActions();
	final Task task = new TestTaskBuilder(shuffleEnvironment)
		.setTaskManagerActions(taskManagerActions)
		.setConsumableNotifier(consumableNotifier)
		.setPartitionProducerStateChecker(partitionProducerStateChecker)
		.setResultPartitions(resultPartitions)
		.setInputGates(inputGates)
		.build();

	// shut down the network to make the following task registration failure
	shuffleEnvironment.close();

	// should fail
	task.run();

	// verify final state
	assertEquals(ExecutionState.FAILED, task.getExecutionState());
	assertTrue(task.isCanceledOrFailed());
	assertTrue(task.getFailureCause().getMessage().contains(errorMessage));

	taskManagerActions.validateListenerMessage(ExecutionState.FAILED, task, new IllegalStateException(errorMessage));
}
 

private JobManagerConnection associateWithJobManager(
		JobID jobID,
		ResourceID resourceID,
		JobMasterGateway jobMasterGateway) {
	checkNotNull(jobID);
	checkNotNull(resourceID);
	checkNotNull(jobMasterGateway);

	TaskManagerActions taskManagerActions = new TaskManagerActionsImpl(jobMasterGateway);

	CheckpointResponder checkpointResponder = new RpcCheckpointResponder(jobMasterGateway);
	GlobalAggregateManager aggregateManager = new RpcGlobalAggregateManager(jobMasterGateway);

	final LibraryCacheManager libraryCacheManager = new BlobLibraryCacheManager(
		blobCacheService.getPermanentBlobService(),
		taskManagerConfiguration.getClassLoaderResolveOrder(),
		taskManagerConfiguration.getAlwaysParentFirstLoaderPatterns());

	ResultPartitionConsumableNotifier resultPartitionConsumableNotifier = new RpcResultPartitionConsumableNotifier(
		jobMasterGateway,
		getRpcService().getExecutor(),
		taskManagerConfiguration.getTimeout());

	PartitionProducerStateChecker partitionStateChecker = new RpcPartitionStateChecker(jobMasterGateway);

	registerQueryableState(jobID, jobMasterGateway);

	return new JobManagerConnection(
		jobID,
		resourceID,
		jobMasterGateway,
		taskManagerActions,
		checkpointResponder,
		aggregateManager,
		libraryCacheManager,
		resultPartitionConsumableNotifier,
		partitionStateChecker);
}
 

private void testExecutionFailsInNetworkRegistration(
		Collection<ResultPartitionDeploymentDescriptor> resultPartitions,
		Collection<InputGateDeploymentDescriptor> inputGates) throws Exception {
	final String errorMessage = "Network buffer pool has already been destroyed.";

	final ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	final PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);

	final QueuedNoOpTaskManagerActions taskManagerActions = new QueuedNoOpTaskManagerActions();
	final Task task = new TestTaskBuilder(shuffleEnvironment)
		.setTaskManagerActions(taskManagerActions)
		.setConsumableNotifier(consumableNotifier)
		.setPartitionProducerStateChecker(partitionProducerStateChecker)
		.setResultPartitions(resultPartitions)
		.setInputGates(inputGates)
		.build();

	// shut down the network to make the following task registration failure
	shuffleEnvironment.close();

	// should fail
	task.run();

	// verify final state
	assertEquals(ExecutionState.FAILED, task.getExecutionState());
	assertTrue(task.isCanceledOrFailed());
	assertTrue(task.getFailureCause().getMessage().contains(errorMessage));

	taskManagerActions.validateListenerMessage(ExecutionState.FAILED, task, new IllegalStateException(errorMessage));
}
 

private JobTable.Connection associateWithJobManager(
		JobTable.Job job,
		ResourceID resourceID,
		JobMasterGateway jobMasterGateway) {
	checkNotNull(resourceID);
	checkNotNull(jobMasterGateway);

	TaskManagerActions taskManagerActions = new TaskManagerActionsImpl(jobMasterGateway);

	CheckpointResponder checkpointResponder = new RpcCheckpointResponder(jobMasterGateway);
	GlobalAggregateManager aggregateManager = new RpcGlobalAggregateManager(jobMasterGateway);

	ResultPartitionConsumableNotifier resultPartitionConsumableNotifier = new RpcResultPartitionConsumableNotifier(
		jobMasterGateway,
		getRpcService().getExecutor(),
		taskManagerConfiguration.getTimeout());

	PartitionProducerStateChecker partitionStateChecker = new RpcPartitionStateChecker(jobMasterGateway);

	registerQueryableState(job.getJobId(), jobMasterGateway);

	return job.connect(
		resourceID,
		jobMasterGateway,
		taskManagerActions,
		checkpointResponder,
		aggregateManager,
		resultPartitionConsumableNotifier,
		partitionStateChecker);
}
 

public ConsumableNotifyingResultPartitionWriterDecorator(
		TaskActions taskActions,
		JobID jobId,
		ResultPartitionWriter partitionWriter,
		ResultPartitionConsumableNotifier partitionConsumableNotifier) {
	this.taskActions = checkNotNull(taskActions);
	this.jobId = checkNotNull(jobId);
	this.partitionWriter = checkNotNull(partitionWriter);
	this.partitionConsumableNotifier = checkNotNull(partitionConsumableNotifier);
}
 

public TestTaskBuilder setConsumableNotifier(ResultPartitionConsumableNotifier consumableNotifier) {
	this.consumableNotifier = consumableNotifier;
	return this;
}
 

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {
	final TestingClassLoaderLease classLoaderHandle = TestingClassLoaderLease.newBuilder()
		.setGetOrResolveClassLoaderFunction((permanentBlobKeys, urls) -> new TestUserCodeClassLoader())
		.build();

	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);
	TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();

	JobInformation jobInformation = new JobInformation(
		new JobID(),
		"Job Name",
		new SerializedValue<>(new ExecutionConfig()),
		new Configuration(),
		Collections.emptyList(),
		Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
		new JobVertexID(),
		"Test Task",
		1,
		1,
		invokableClass.getName(),
		new Configuration());

	return new Task(
		jobInformation,
		taskInformation,
		new ExecutionAttemptID(),
		new AllocationID(),
		0,
		0,
		Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
		Collections.<InputGateDeploymentDescriptor>emptyList(),
		0,
		mock(MemoryManager.class),
		mock(IOManager.class),
		shuffleEnvironment,
		new KvStateService(new KvStateRegistry(), null, null),
		mock(BroadcastVariableManager.class),
		new TaskEventDispatcher(),
		ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES,
		new TestTaskStateManager(),
		mock(TaskManagerActions.class),
		mock(InputSplitProvider.class),
		mock(CheckpointResponder.class),
		new NoOpTaskOperatorEventGateway(),
		new TestGlobalAggregateManager(),
		classLoaderHandle,
		mock(FileCache.class),
		new TestingTaskManagerRuntimeInfo(),
		taskMetricGroup,
		consumableNotifier,
		partitionProducerStateChecker,
		executor);
}
 

public ResultPartitionConsumableNotifier getResultPartitionConsumableNotifier() {
	return resultPartitionConsumableNotifier;
}
 

@Override
public ResultPartitionConsumableNotifier getResultPartitionConsumableNotifier() {
	return verifyContainsEstablishedConnection().getResultPartitionConsumableNotifier();
}
 

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {

		ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
		PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
		Executor executor = mock(Executor.class);
		ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();

		TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();

		JobInformation jobInformation = new JobInformation(
				new JobID(),
				"Job Name",
				new SerializedValue<>(new ExecutionConfig()),
				new Configuration(),
				Collections.emptyList(),
				Collections.emptyList());

		TaskInformation taskInformation = new TaskInformation(
				new JobVertexID(),
				"Test Task",
				1,
				1,
				invokableClass.getName(),
				new Configuration());

		return new Task(
				jobInformation,
				taskInformation,
				new ExecutionAttemptID(),
				new AllocationID(),
				0,
				0,
				Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
				Collections.<InputGateDeploymentDescriptor>emptyList(),
				0,
				mock(MemoryManager.class),
				mock(IOManager.class),
				shuffleEnvironment,
				new KvStateService(new KvStateRegistry(), null, null),
				mock(BroadcastVariableManager.class),
				new TaskEventDispatcher(),
				ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES,
				new TestTaskStateManager(),
				mock(TaskManagerActions.class),
				mock(InputSplitProvider.class),
				mock(CheckpointResponder.class),
				new NoOpTaskOperatorEventGateway(),
				new TestGlobalAggregateManager(),
				TestingClassLoaderLease.newBuilder().build(),
				mock(FileCache.class),
				new TestingTaskManagerRuntimeInfo(),
				taskMetricGroup,
				consumableNotifier,
				partitionProducerStateChecker,
				executor);
	}
 

public static Task createTask(
		Class<? extends AbstractInvokable> invokable,
		StreamConfig taskConfig,
		Configuration taskManagerConfig,
		TestTaskStateManager taskStateManager,
		TaskManagerActions taskManagerActions) throws Exception {

	BlobCacheService blobService =
		new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

	LibraryCacheManager libCache = mock(LibraryCacheManager.class);
	when(libCache.getClassLoader(any(JobID.class))).thenReturn(StreamTaskTest.class.getClassLoader());

	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);

	ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();

	JobInformation jobInformation = new JobInformation(
		new JobID(),
		"Job Name",
		new SerializedValue<>(new ExecutionConfig()),
		new Configuration(),
		Collections.emptyList(),
		Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
		new JobVertexID(),
		"Test Task",
		1,
		1,
		invokable.getName(),
		taskConfig.getConfiguration());

	return new Task(
		jobInformation,
		taskInformation,
		new ExecutionAttemptID(),
		new AllocationID(),
		0,
		0,
		Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
		Collections.<InputGateDeploymentDescriptor>emptyList(),
		0,
		mock(MemoryManager.class),
		mock(IOManager.class),
		shuffleEnvironment,
		new KvStateService(new KvStateRegistry(), null, null),
		mock(BroadcastVariableManager.class),
		new TaskEventDispatcher(),
		taskStateManager,
		taskManagerActions,
		mock(InputSplitProvider.class),
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		blobService,
		libCache,
		mock(FileCache.class),
		new TestingTaskManagerRuntimeInfo(taskManagerConfig, new String[] {System.getProperty("java.io.tmpdir")}),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
		consumableNotifier,
		partitionProducerStateChecker,
		executor);
}
 

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {
	BlobCacheService blobService =
			new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

	LibraryCacheManager libCache = mock(LibraryCacheManager.class);
	when(libCache.getClassLoader(any(JobID.class))).thenReturn(ClassLoader.getSystemClassLoader());

	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);
	ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();

	TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();

	JobInformation jobInformation = new JobInformation(
			new JobID(),
			"Job Name",
			new SerializedValue<>(new ExecutionConfig()),
			new Configuration(),
			Collections.emptyList(),
			Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
			new JobVertexID(),
			"Test Task",
			1,
			1,
			invokableClass.getName(),
			new Configuration());

	return new Task(
			jobInformation,
			taskInformation,
			new ExecutionAttemptID(),
			new AllocationID(),
			0,
			0,
			Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
			Collections.<InputGateDeploymentDescriptor>emptyList(),
			0,
			mock(MemoryManager.class),
			mock(IOManager.class),
			shuffleEnvironment,
			new KvStateService(new KvStateRegistry(), null, null),
			mock(BroadcastVariableManager.class),
			new TaskEventDispatcher(),
			new TestTaskStateManager(),
			mock(TaskManagerActions.class),
			mock(InputSplitProvider.class),
			mock(CheckpointResponder.class),
			new TestGlobalAggregateManager(),
			blobService,
			libCache,
			mock(FileCache.class),
			new TestingTaskManagerRuntimeInfo(),
			taskMetricGroup,
			consumableNotifier,
			partitionProducerStateChecker,
			executor);
}
 

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {
	BlobCacheService blobService =
		new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

	LibraryCacheManager libCache = mock(LibraryCacheManager.class);
	when(libCache.getClassLoader(any(JobID.class))).thenReturn(new TestUserCodeClassLoader());

	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);
	TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();

	JobInformation jobInformation = new JobInformation(
		new JobID(),
		"Job Name",
		new SerializedValue<>(new ExecutionConfig()),
		new Configuration(),
		Collections.emptyList(),
		Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
		new JobVertexID(),
		"Test Task",
		1,
		1,
		invokableClass.getName(),
		new Configuration());

	return new Task(
		jobInformation,
		taskInformation,
		new ExecutionAttemptID(),
		new AllocationID(),
		0,
		0,
		Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
		Collections.<InputGateDeploymentDescriptor>emptyList(),
		0,
		mock(MemoryManager.class),
		mock(IOManager.class),
		shuffleEnvironment,
		new KvStateService(new KvStateRegistry(), null, null),
		mock(BroadcastVariableManager.class),
		new TaskEventDispatcher(),
		new TestTaskStateManager(),
		mock(TaskManagerActions.class),
		mock(InputSplitProvider.class),
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		blobService,
		libCache,
		mock(FileCache.class),
		new TestingTaskManagerRuntimeInfo(),
		taskMetricGroup,
		consumableNotifier,
		partitionProducerStateChecker,
		executor);
}
 

public TestTaskBuilder setConsumableNotifier(ResultPartitionConsumableNotifier consumableNotifier) {
	this.consumableNotifier = consumableNotifier;
	return this;
}
 

public ResultPartitionConsumableNotifier getResultPartitionConsumableNotifier() {
	return resultPartitionConsumableNotifier;
}
 

public static Task createTask(
		Class<? extends AbstractInvokable> invokable,
		StreamConfig taskConfig,
		Configuration taskManagerConfig,
		TestTaskStateManager taskStateManager,
		TaskManagerActions taskManagerActions) throws Exception {

	BlobCacheService blobService =
		new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

	LibraryCacheManager libCache = mock(LibraryCacheManager.class);
	when(libCache.getClassLoader(any(JobID.class))).thenReturn(StreamTaskTest.class.getClassLoader());

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);
	TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();

	NetworkEnvironment network = mock(NetworkEnvironment.class);
	when(network.getResultPartitionManager()).thenReturn(partitionManager);
	when(network.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC);
	when(network.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class)))
			.thenReturn(mock(TaskKvStateRegistry.class));
	when(network.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);

	JobInformation jobInformation = new JobInformation(
		new JobID(),
		"Job Name",
		new SerializedValue<>(new ExecutionConfig()),
		new Configuration(),
		Collections.emptyList(),
		Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
		new JobVertexID(),
		"Test Task",
		1,
		1,
		invokable.getName(),
		taskConfig.getConfiguration());

	return new Task(
		jobInformation,
		taskInformation,
		new ExecutionAttemptID(),
		new AllocationID(),
		0,
		0,
		Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
		Collections.<InputGateDeploymentDescriptor>emptyList(),
		0,
		mock(MemoryManager.class),
		mock(IOManager.class),
		network,
		mock(BroadcastVariableManager.class),
		taskStateManager,
		taskManagerActions,
		mock(InputSplitProvider.class),
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		blobService,
		libCache,
		mock(FileCache.class),
		new TestingTaskManagerRuntimeInfo(taskManagerConfig, new String[] {System.getProperty("java.io.tmpdir")}),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
		consumableNotifier,
		partitionProducerStateChecker,
		executor);
}
 

TaskBuilder setConsumableNotifier(ResultPartitionConsumableNotifier consumableNotifier) {
	this.consumableNotifier = consumableNotifier;
	return this;
}
 

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {
	BlobCacheService blobService =
		new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

	LibraryCacheManager libCache = mock(LibraryCacheManager.class);
	when(libCache.getClassLoader(any(JobID.class))).thenReturn(new TestUserCodeClassLoader());

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	Executor executor = mock(Executor.class);
	TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);
	NetworkEnvironment networkEnvironment = mock(NetworkEnvironment.class);
	when(networkEnvironment.getResultPartitionManager()).thenReturn(partitionManager);
	when(networkEnvironment.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC);
	when(networkEnvironment.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class)))
			.thenReturn(mock(TaskKvStateRegistry.class));
	when(networkEnvironment.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);

	TaskMetricGroup taskMetricGroup = mock(TaskMetricGroup.class);
	when(taskMetricGroup.getIOMetricGroup()).thenReturn(mock(TaskIOMetricGroup.class));

	JobInformation jobInformation = new JobInformation(
		new JobID(),
		"Job Name",
		new SerializedValue<>(new ExecutionConfig()),
		new Configuration(),
		Collections.emptyList(),
		Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
		new JobVertexID(),
		"Test Task",
		1,
		1,
		invokableClass.getName(),
		new Configuration());

	return new Task(
		jobInformation,
		taskInformation,
		new ExecutionAttemptID(),
		new AllocationID(),
		0,
		0,
		Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
		Collections.<InputGateDeploymentDescriptor>emptyList(),
		0,
		mock(MemoryManager.class),
		mock(IOManager.class),
		networkEnvironment,
		mock(BroadcastVariableManager.class),
		new TestTaskStateManager(),
		mock(TaskManagerActions.class),
		mock(InputSplitProvider.class),
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		blobService,
		libCache,
		mock(FileCache.class),
		new TestingTaskManagerRuntimeInfo(),
		taskMetricGroup,
		consumableNotifier,
		partitionProducerStateChecker,
		executor);
}
 

public ResultPartitionConsumableNotifier getResultPartitionConsumableNotifier() {
	return resultPartitionConsumableNotifier;
}
 

/**
 * Connects the job to a JobManager and associates the provided services with this connection.
 *
 * <p>A job can only be connected iff {@code Job#isConnected() == false}.
 *
 * @param resourceId resourceId of the JobManager to connect to
 * @param jobMasterGateway jobMasterGateway of the JobManager to connect to
 * @param taskManagerActions taskManagerActions associated with this connection
 * @param checkpointResponder checkpointResponder associated with this connection
 * @param aggregateManager aggregateManager associated with this connection
 * @param resultPartitionConsumableNotifier resultPartitionConsumableNotifier associated with this connection
 * @param partitionStateChecker partitionStateChecker associated with this connection
 * @return the established {@link Connection}
 * @throws IllegalStateException if the job is already connected
 */
Connection connect(
	ResourceID resourceId,
	JobMasterGateway jobMasterGateway,
	TaskManagerActions taskManagerActions,
	CheckpointResponder checkpointResponder,
	GlobalAggregateManager aggregateManager,
	ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
	PartitionProducerStateChecker partitionStateChecker);
 

ResultPartitionConsumableNotifier getResultPartitionConsumableNotifier();