Java Code Examples for org.apache.flink.runtime.jobmaster.JobMasterId#generate()

@Before
public void setup() throws Exception {
	final JobID jobId = new JobID();
	final SlotPool slotPool = new SlotPoolImpl(jobId);
	final TestingScheduler testingScheduler = new TestingScheduler(
		new HashMap<>(16),
		LocationPreferenceSlotSelectionStrategy.INSTANCE,
		slotPool);

	testingSlotProvider = new TestingSlotPoolSlotProvider(slotPool, testingScheduler);

	final JobMasterId jobMasterId = JobMasterId.generate();
	final String jobManagerAddress = "localhost";
	ComponentMainThreadExecutor executor = TestingComponentMainThreadExecutorServiceAdapter.forMainThread();
	slotPool.start(jobMasterId, jobManagerAddress, executor);
	testingScheduler.start(executor);
}
 

@Before
public void setup() throws Exception {
	final JobID jobId = new JobID();
	final SlotPool slotPool = new TestingSlotPoolImpl(jobId);
	final TestingScheduler testingScheduler = new TestingScheduler(
		new HashMap<>(16),
		LocationPreferenceSlotSelectionStrategy.INSTANCE,
		slotPool);

	testingSlotProvider = new TestingSlotPoolSlotProvider(slotPool, testingScheduler);

	final JobMasterId jobMasterId = JobMasterId.generate();
	final String jobManagerAddress = "localhost";
	ComponentMainThreadExecutor executor = ComponentMainThreadExecutorServiceAdapter.forMainThread();
	slotPool.start(jobMasterId, jobManagerAddress, executor);
	testingScheduler.start(executor);
}
 

/**
 * Test receive registration with unmatched leadershipId from job master.
 */
@Test
public void testRegisterJobMasterWithUnmatchedLeaderSessionId2() throws Exception {
			// test throw exception when receive a registration from job master which takes unmatched leaderSessionId
	JobMasterId differentJobMasterId = JobMasterId.generate();
	CompletableFuture<RegistrationResponse> unMatchedLeaderFuture = resourceManagerGateway.registerJobManager(
		differentJobMasterId,
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);
	assertTrue(unMatchedLeaderFuture.get() instanceof RegistrationResponse.Decline);
}
 

@Before
public void setup() throws Exception {
	final JobID jobId = new JobID();
	slotPool = new TestingSlotPoolImpl(jobId);
	scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), slotPool);

	testingSlotProvider = new TestingSlotPoolSlotProvider();

	final JobMasterId jobMasterId = JobMasterId.generate();
	final String jobManagerAddress = "localhost";
	componentMainThreadExecutor = getComponentMainThreadExecutor();
	slotPool.start(jobMasterId, jobManagerAddress, componentMainThreadExecutor);
	scheduler.start(componentMainThreadExecutor);
}
 

/**
 * Test receive registration with unmatched leadershipId from job master.
 */
@Test
public void testRegisterJobMasterWithUnmatchedLeaderSessionId2() throws Exception {
			// test throw exception when receive a registration from job master which takes unmatched leaderSessionId
	JobMasterId differentJobMasterId = JobMasterId.generate();
	CompletableFuture<RegistrationResponse> unMatchedLeaderFuture = resourceManagerGateway.registerJobManager(
		differentJobMasterId,
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);
	assertTrue(unMatchedLeaderFuture.get() instanceof RegistrationResponse.Decline);
}
 

private TaskSubmissionTestEnvironment createExecutorWithRunningTask(
		JobID jobId,
		ExecutionAttemptID executionAttemptId,
		Class<? extends AbstractInvokable> invokableClass) throws Exception {

	final TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(
			jobId, executionAttemptId, invokableClass);

	final CompletableFuture<Void> taskRunningFuture = new CompletableFuture<>();

	final JobMasterId token = JobMasterId.generate();
	final TaskSubmissionTestEnvironment env = new TaskSubmissionTestEnvironment.Builder(jobId)
			.setJobMasterId(token)
			.setSlotSize(1)
			.addTaskManagerActionListener(executionAttemptId, ExecutionState.RUNNING, taskRunningFuture)
			.setMetricQueryServiceAddress(metricRegistry.getMetricQueryServiceGatewayRpcAddress())
			.setJobMasterGateway(new TestingJobMasterGatewayBuilder()
				.setFencingTokenSupplier(() -> token)
				.setOperatorEventSender((eio, oid, value) -> {
					throw new RuntimeException();
				})
				.build())
			.build();

	env.getTaskSlotTable().allocateSlot(0, jobId, tdd.getAllocationId(), Time.seconds(60));

	final TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
	tmGateway.submitTask(tdd, env.getJobMasterId(), Time.seconds(10)).get();
	taskRunningFuture.get();

	return env;
}
 

MockJobMaster(JobID jobID) {
	this.jobID = jobID;
	this.resourceID = new ResourceID(jobID.toString());
	this.address = "/" + jobID;
	this.gateway = mock(JobMasterGateway.class);
	this.jobMasterId = JobMasterId.generate();
	this.leaderRetrievalService = new SettableLeaderRetrievalService(this.address, this.jobMasterId.toUUID());
}
 

MockJobMaster(JobID jobID) {
	this.jobID = jobID;
	this.resourceID = new ResourceID(jobID.toString());
	this.address = "/" + jobID;
	this.gateway = mock(JobMasterGateway.class);
	this.jobMasterId = JobMasterId.generate();
	this.leaderRetrievalService = new SettableLeaderRetrievalService(this.address, this.jobMasterId.toUUID());
}
 

MockJobMaster(JobID jobID) {
	this.jobID = jobID;
	this.resourceID = new ResourceID(jobID.toString());
	this.address = "/" + jobID;
	this.gateway = mock(JobMasterGateway.class);
	this.jobMasterId = JobMasterId.generate();
	this.leaderRetrievalService = new SettableLeaderRetrievalService(this.address, this.jobMasterId.toUUID());
}
 

/**
 * Test receive registration with unmatched leadershipId from job master.
 */
@Test
public void testRegisterJobMasterWithUnmatchedLeaderSessionId2() throws Exception {
			// test throw exception when receive a registration from job master which takes unmatched leaderSessionId
	JobMasterId differentJobMasterId = JobMasterId.generate();
	CompletableFuture<RegistrationResponse> unMatchedLeaderFuture = resourceManagerGateway.registerJobManager(
		differentJobMasterId,
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);
	assertTrue(unMatchedLeaderFuture.get() instanceof RegistrationResponse.Decline);
}
 

/**
 * Tests that we can submit a task to the TaskManager given that we've allocated a slot there.
 */
@Test(timeout = 10000L)
public void testTaskSubmission() throws Exception {
	final AllocationID allocationId = new AllocationID();
	final JobMasterId jobMasterId = JobMasterId.generate();
	final JobVertexID jobVertexId = new JobVertexID();

	JobInformation jobInformation = new JobInformation(
			jobId,
			testName.getMethodName(),
			new SerializedValue<>(new ExecutionConfig()),
			new Configuration(),
			Collections.emptyList(),
			Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
			jobVertexId,
			"test task",
			1,
			1,
			TestInvokable.class.getName(),
			new Configuration());

	SerializedValue<JobInformation> serializedJobInformation = new SerializedValue<>(jobInformation);
	SerializedValue<TaskInformation> serializedJobVertexInformation = new SerializedValue<>(taskInformation);

	final TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
			jobId,
			new TaskDeploymentDescriptor.NonOffloaded<>(serializedJobInformation),
			new TaskDeploymentDescriptor.NonOffloaded<>(serializedJobVertexInformation),
			new ExecutionAttemptID(),
			allocationId,
			0,
			0,
			0,
			null,
			Collections.emptyList(),
			Collections.emptyList());

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

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
	when(jobMasterGateway.getFencingToken()).thenReturn(jobMasterId);

	final OneShotLatch taskInTerminalState = new OneShotLatch();
	final TaskManagerActions taskManagerActions = new NoOpTaskManagerActions() {
		@Override
		public void updateTaskExecutionState(TaskExecutionState taskExecutionState) {
			if (taskExecutionState.getExecutionState().isTerminal()) {
				taskInTerminalState.trigger();
			}
		}
	};
	final JobManagerConnection jobManagerConnection = new JobManagerConnection(
		jobId,
		ResourceID.generate(),
		jobMasterGateway,
		taskManagerActions,
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		libraryCacheManager,
		new NoOpResultPartitionConsumableNotifier(),
		mock(PartitionProducerStateChecker.class));

	final JobManagerTable jobManagerTable = new JobManagerTable();
	jobManagerTable.put(jobId, jobManagerConnection);

	final TaskSlotTable taskSlotTable = mock(TaskSlotTable.class);
	when(taskSlotTable.tryMarkSlotActive(eq(jobId), eq(allocationId))).thenReturn(true);
	when(taskSlotTable.addTask(any(Task.class))).thenReturn(true);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setShuffleEnvironment(nettyShuffleEnvironment)
		.setTaskSlotTable(taskSlotTable)
		.setJobManagerTable(jobManagerTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		final TaskExecutorGateway tmGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		tmGateway.submitTask(tdd, jobMasterId, timeout);

		CompletableFuture<Boolean> completionFuture = TestInvokable.COMPLETABLE_FUTURE;

		completionFuture.get();

		taskInTerminalState.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Test that a failing schedule or update consumers call leads to the failing of the respective
 * task.
 *
 * <p>IMPORTANT: We have to make sure that the invokable's cancel method is called, because only
 * then the future is completed. We do this by not eagerly deploying consumer tasks and requiring
 * the invokable to fill one memory segment. The completed memory segment will trigger the
 * scheduling of the downstream operator since it is in pipeline mode. After we've filled the
 * memory segment, we'll block the invokable and wait for the task failure due to the failed
 * schedule or update consumers call.
 */
@Test(timeout = 10000L)
public void testFailingScheduleOrUpdateConsumers() throws Exception {
	final Configuration configuration = new Configuration();

	// set the memory segment to the smallest size possible, because we have to fill one
	// memory buffer to trigger the schedule or update consumers message to the downstream
	// operators
	configuration.setString(TaskManagerOptions.MEMORY_SEGMENT_SIZE, "4096");

	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), ResourceID.generate());
	TaskDeploymentDescriptor tdd = createSender(sdd, TestingAbstractInvokables.TestInvokableRecordCancel.class);
	ExecutionAttemptID eid = tdd.getExecutionAttemptId();

	final CompletableFuture<Void> taskRunningFuture = new CompletableFuture<>();

	final Exception exception = new Exception("Failed schedule or update consumers");

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
			.setFencingTokenSupplier(() -> jobMasterId)
			.setUpdateTaskExecutionStateFunction(resultPartitionID -> FutureUtils.completedExceptionally(exception))
			.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setSlotSize(1)
			.setConfiguration(configuration)
			.addTaskManagerActionListener(eid, ExecutionState.RUNNING, taskRunningFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable taskSlotTable = env.getTaskSlotTable();

		TestingAbstractInvokables.TestInvokableRecordCancel.resetGotCanceledFuture();

		taskSlotTable.allocateSlot(0, jobId, tdd.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd, jobMasterId, timeout).get();
		taskRunningFuture.get();

		CompletableFuture<Boolean> cancelFuture = TestingAbstractInvokables.TestInvokableRecordCancel.gotCanceled();

		assertTrue(cancelFuture.get());
		assertTrue(ExceptionUtils.findThrowableWithMessage(taskSlotTable.getTask(eid).getFailureCause(), exception.getMessage()).isPresent());
	}
}
 

/**
 * This tests creates two tasks. The sender sends data but fails to send the
 * state update back to the job manager.
 * the second one blocks to be canceled
 */
@Test(timeout = 10000L)
public void testCancellingDependentAndStateUpdateFails() throws Exception {
	ResourceID producerLocation = ResourceID.generate();
	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), producerLocation);

	TaskDeploymentDescriptor tdd1 = createSender(sdd);
	TaskDeploymentDescriptor tdd2 = createReceiver(sdd);
	ExecutionAttemptID eid1 = tdd1.getExecutionAttemptId();
	ExecutionAttemptID eid2 = tdd2.getExecutionAttemptId();

	final CompletableFuture<Void> task1RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task1FailedFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2CanceledFuture = new CompletableFuture<>();

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
		.setFencingTokenSupplier(() -> jobMasterId)
		.setUpdateTaskExecutionStateFunction(taskExecutionState -> {
			if (taskExecutionState != null && taskExecutionState.getID().equals(eid1)) {
				return FutureUtils.completedExceptionally(
					new ExecutionGraphException("The execution attempt " + eid2 + " was not found."));
			} else {
				return CompletableFuture.completedFuture(Acknowledge.get());
			}
		})
		.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setResourceID(producerLocation)
			.setSlotSize(2)
			.addTaskManagerActionListener(eid1, ExecutionState.RUNNING, task1RunningFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.RUNNING, task2RunningFuture)
			.addTaskManagerActionListener(eid1, ExecutionState.FAILED, task1FailedFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.CANCELED, task2CanceledFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable taskSlotTable = env.getTaskSlotTable();

		taskSlotTable.allocateSlot(0, jobId, tdd1.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd1, jobMasterId, timeout).get();
		task1RunningFuture.get();

		taskSlotTable.allocateSlot(1, jobId, tdd2.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd2, jobMasterId, timeout).get();
		task2RunningFuture.get();

		task1FailedFuture.get();
		assertSame(taskSlotTable.getTask(eid1).getExecutionState(), ExecutionState.FAILED);

		tmGateway.cancelTask(eid2, timeout);

		task2CanceledFuture.get();
		assertSame(taskSlotTable.getTask(eid2).getExecutionState(), ExecutionState.CANCELED);
	}
}
 

@Test(timeout = TEST_TIMEOUT)
public void testRunJobWithForwardChannel() throws Exception {
	ResourceID producerLocation = ResourceID.generate();
	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), producerLocation);

	TaskDeploymentDescriptor tdd1 = createSender(sdd);
	TaskDeploymentDescriptor tdd2 = createReceiver(sdd);
	ExecutionAttemptID eid1 = tdd1.getExecutionAttemptId();
	ExecutionAttemptID eid2 = tdd2.getExecutionAttemptId();

	final CompletableFuture<Void> task1RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task1FinishedFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2FinishedFuture = new CompletableFuture<>();

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
		.setFencingTokenSupplier(() -> jobMasterId)
		.setScheduleOrUpdateConsumersFunction(
			resultPartitionID -> CompletableFuture.completedFuture(Acknowledge.get()))
		.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setResourceID(producerLocation)
			.setSlotSize(2)
			.addTaskManagerActionListener(eid1, ExecutionState.RUNNING, task1RunningFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.RUNNING, task2RunningFuture)
			.addTaskManagerActionListener(eid1, ExecutionState.FINISHED, task1FinishedFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.FINISHED, task2FinishedFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable<Task> taskSlotTable = env.getTaskSlotTable();

		taskSlotTable.allocateSlot(0, jobId, tdd1.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd1, jobMasterId, timeout).get();
		task1RunningFuture.get();

		taskSlotTable.allocateSlot(1, jobId, tdd2.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd2, jobMasterId, timeout).get();
		task2RunningFuture.get();

		task1FinishedFuture.get();
		task2FinishedFuture.get();

		assertSame(taskSlotTable.getTask(eid1).getExecutionState(), ExecutionState.FINISHED);
		assertSame(taskSlotTable.getTask(eid2).getExecutionState(), ExecutionState.FINISHED);
	}
}
 

/**
 * This tests creates two tasks. The sender sends data but fails to send the
 * state update back to the job manager.
 * the second one blocks to be canceled
 */
@Test(timeout = TEST_TIMEOUT)
public void testCancellingDependentAndStateUpdateFails() throws Exception {
	ResourceID producerLocation = ResourceID.generate();
	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), producerLocation);

	TaskDeploymentDescriptor tdd1 = createSender(sdd);
	TaskDeploymentDescriptor tdd2 = createReceiver(sdd);
	ExecutionAttemptID eid1 = tdd1.getExecutionAttemptId();
	ExecutionAttemptID eid2 = tdd2.getExecutionAttemptId();

	final CompletableFuture<Void> task1RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task1FailedFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2CanceledFuture = new CompletableFuture<>();

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
		.setFencingTokenSupplier(() -> jobMasterId)
		.setUpdateTaskExecutionStateFunction(taskExecutionState -> {
			if (taskExecutionState != null && taskExecutionState.getID().equals(eid1)) {
				return FutureUtils.completedExceptionally(
					new ExecutionGraphException("The execution attempt " + eid2 + " was not found."));
			} else {
				return CompletableFuture.completedFuture(Acknowledge.get());
			}
		})
		.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setResourceID(producerLocation)
			.setSlotSize(2)
			.addTaskManagerActionListener(eid1, ExecutionState.RUNNING, task1RunningFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.RUNNING, task2RunningFuture)
			.addTaskManagerActionListener(eid1, ExecutionState.FAILED, task1FailedFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.CANCELED, task2CanceledFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable<Task> taskSlotTable = env.getTaskSlotTable();

		taskSlotTable.allocateSlot(0, jobId, tdd1.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd1, jobMasterId, timeout).get();
		task1RunningFuture.get();

		taskSlotTable.allocateSlot(1, jobId, tdd2.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd2, jobMasterId, timeout).get();
		task2RunningFuture.get();

		task1FailedFuture.get();
		assertSame(taskSlotTable.getTask(eid1).getExecutionState(), ExecutionState.FAILED);

		tmGateway.cancelTask(eid2, timeout);

		task2CanceledFuture.get();
		assertSame(taskSlotTable.getTask(eid2).getExecutionState(), ExecutionState.CANCELED);
	}
}
 

/**
 * Test that a failing schedule or update consumers call leads to the failing of the respective
 * task.
 *
 * <p>IMPORTANT: We have to make sure that the invokable's cancel method is called, because only
 * then the future is completed. We do this by not eagerly deploying consumer tasks and requiring
 * the invokable to fill one memory segment. The completed memory segment will trigger the
 * scheduling of the downstream operator since it is in pipeline mode. After we've filled the
 * memory segment, we'll block the invokable and wait for the task failure due to the failed
 * schedule or update consumers call.
 */
@Test(timeout = TEST_TIMEOUT)
public void testFailingScheduleOrUpdateConsumers() throws Exception {
	final Configuration configuration = new Configuration();

	// set the memory segment to the smallest size possible, because we have to fill one
	// memory buffer to trigger the schedule or update consumers message to the downstream
	// operators
	configuration.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("4096"));

	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), ResourceID.generate());
	TaskDeploymentDescriptor tdd = createSender(sdd, TestingAbstractInvokables.TestInvokableRecordCancel.class);
	ExecutionAttemptID eid = tdd.getExecutionAttemptId();

	final CompletableFuture<Void> taskRunningFuture = new CompletableFuture<>();

	final Exception exception = new Exception("Failed schedule or update consumers");

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
			.setFencingTokenSupplier(() -> jobMasterId)
			.setUpdateTaskExecutionStateFunction(resultPartitionID -> FutureUtils.completedExceptionally(exception))
			.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setSlotSize(1)
			.setConfiguration(configuration)
			.addTaskManagerActionListener(eid, ExecutionState.RUNNING, taskRunningFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable<Task> taskSlotTable = env.getTaskSlotTable();

		TestingAbstractInvokables.TestInvokableRecordCancel.resetGotCanceledFuture();

		taskSlotTable.allocateSlot(0, jobId, tdd.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd, jobMasterId, timeout).get();
		taskRunningFuture.get();

		CompletableFuture<Boolean> cancelFuture = TestingAbstractInvokables.TestInvokableRecordCancel.gotCanceled();

		assertTrue(cancelFuture.get());
		assertTrue(ExceptionUtils.findThrowableWithMessage(taskSlotTable.getTask(eid).getFailureCause(), exception.getMessage()).isPresent());
	}
}
 

@Test(timeout = 10000L)
public void testRunJobWithForwardChannel() throws Exception {
	ResourceID producerLocation = ResourceID.generate();
	NettyShuffleDescriptor sdd =
		createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), producerLocation);

	TaskDeploymentDescriptor tdd1 = createSender(sdd);
	TaskDeploymentDescriptor tdd2 = createReceiver(sdd);
	ExecutionAttemptID eid1 = tdd1.getExecutionAttemptId();
	ExecutionAttemptID eid2 = tdd2.getExecutionAttemptId();

	final CompletableFuture<Void> task1RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2RunningFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task1FinishedFuture = new CompletableFuture<>();
	final CompletableFuture<Void> task2FinishedFuture = new CompletableFuture<>();

	final JobMasterId jobMasterId = JobMasterId.generate();
	TestingJobMasterGateway testingJobMasterGateway =
		new TestingJobMasterGatewayBuilder()
		.setFencingTokenSupplier(() -> jobMasterId)
		.setScheduleOrUpdateConsumersFunction(
			resultPartitionID -> CompletableFuture.completedFuture(Acknowledge.get()))
		.build();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setResourceID(producerLocation)
			.setSlotSize(2)
			.addTaskManagerActionListener(eid1, ExecutionState.RUNNING, task1RunningFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.RUNNING, task2RunningFuture)
			.addTaskManagerActionListener(eid1, ExecutionState.FINISHED, task1FinishedFuture)
			.addTaskManagerActionListener(eid2, ExecutionState.FINISHED, task2FinishedFuture)
			.setJobMasterId(jobMasterId)
			.setJobMasterGateway(testingJobMasterGateway)
			.useRealNonMockShuffleEnvironment()
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable taskSlotTable = env.getTaskSlotTable();

		taskSlotTable.allocateSlot(0, jobId, tdd1.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd1, jobMasterId, timeout).get();
		task1RunningFuture.get();

		taskSlotTable.allocateSlot(1, jobId, tdd2.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd2, jobMasterId, timeout).get();
		task2RunningFuture.get();

		task1FinishedFuture.get();
		task2FinishedFuture.get();

		assertSame(taskSlotTable.getTask(eid1).getExecutionState(), ExecutionState.FINISHED);
		assertSame(taskSlotTable.getTask(eid2).getExecutionState(), ExecutionState.FINISHED);
	}
}