Java Code Examples for org.apache.flink.runtime.jobgraph.JobGraph#setScheduleMode()

@Nonnull
private JobGraph createJobGraph(long delay, int parallelism) throws IOException {
	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(OneTimeFailingInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(slotSharingGroup);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(slotSharingGroup);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
	JobGraph jobGraph = new JobGraph(source, sink);

	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, delay));
	jobGraph.setExecutionConfig(executionConfig);

	return jobGraph;
}
 

@Nonnull
private JobGraph createJobGraph(long delay, int parallelism) throws IOException {
	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(OneTimeFailingInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(slotSharingGroup);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(slotSharingGroup);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
	JobGraph jobGraph = new JobGraph(source, sink);

	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, delay));
	jobGraph.setExecutionConfig(executionConfig);

	return jobGraph;
}
 

private static JobGraph nonParallelSourceSinkJobGraph() {
	final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Testjob");
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(NoOpInvokable.class);
	jobGraph.addVertex(source);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	jobGraph.addVertex(sink);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	return jobGraph;
}
 

private static JobGraph createJobGraph() {
	final SlotSharingGroup group1 = new SlotSharingGroup();
	final SlotSharingGroup group2 = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(TestSourceInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(group1);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(TestSinkInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(group2);

	sink.connectNewDataSetAsInput(source, DistributionPattern.ALL_TO_ALL, ResultPartitionType.BLOCKING);

	final JobGraph jobGraph = new JobGraph(source, sink);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES);

	return jobGraph;
}
 

@Nonnull
private JobGraph createJobGraph(long delay, int parallelism) throws IOException {
	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(OneTimeFailingInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(slotSharingGroup);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(slotSharingGroup);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
	JobGraph jobGraph = new JobGraph(source, sink);

	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, delay));
	jobGraph.setExecutionConfig(executionConfig);

	return jobGraph;
}
 

private static JobGraph createJobGraph() {
	final SlotSharingGroup group1 = new SlotSharingGroup();
	final SlotSharingGroup group2 = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(TestSourceInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(group1);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(TestSinkInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(group2);

	sink.connectNewDataSetAsInput(source, DistributionPattern.ALL_TO_ALL, ResultPartitionType.BLOCKING);

	final JobGraph jobGraph = new JobGraph(source, sink);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES);

	return jobGraph;
}
 

/**
 * Tests that an ongoing scheduling operation does not fail the {@link ExecutionGraph}
 * if it gets concurrently cancelled.
 */
@Test
public void testSchedulingOperationCancellationWhenCancel() throws Exception {
	final JobVertex jobVertex = new JobVertex("NoOp JobVertex");
	jobVertex.setInvokableClass(NoOpInvokable.class);
	jobVertex.setParallelism(2);
	final JobGraph jobGraph = new JobGraph(jobVertex);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final CompletableFuture<LogicalSlot> slotFuture1 = new CompletableFuture<>();
	final CompletableFuture<LogicalSlot> slotFuture2 = new CompletableFuture<>();
	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(2);
	slotProvider.addSlots(jobVertex.getID(), new CompletableFuture[]{slotFuture1, slotFuture2});
	final ExecutionGraph executionGraph = createExecutionGraph(jobGraph, slotProvider);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());
	executionGraph.scheduleForExecution();

	final TestingLogicalSlot slot = createTestingSlot();
	final CompletableFuture<?> releaseFuture = slot.getReleaseFuture();
	slotFuture1.complete(slot);

	// cancel should change the state of all executions to CANCELLED
	executionGraph.cancel();

	// complete the now CANCELLED execution --> this should cause a failure
	slotFuture2.complete(new TestingLogicalSlotBuilder().createTestingLogicalSlot());

	Thread.sleep(1L);
	// release the first slot to finish the cancellation
	releaseFuture.complete(null);

	// NOTE: This test will only occasionally fail without the fix since there is
	// a race between the releaseFuture and the slotFuture2
	assertThat(executionGraph.getTerminationFuture().get(), is(JobStatus.CANCELED));
}
 

private static JobGraph createStreamingJobGraph() {
	final JobVertex v1 = new JobVertex("vertex1");
	final JobVertex v2 = new JobVertex("vertex2");
	v1.setInvokableClass(AbstractInvokable.class);
	v2.setInvokableClass(AbstractInvokable.class);

	final JobGraph jobGraph = new JobGraph(v1, v2);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	return jobGraph;
}
 

private void setupAndRunHandleJobsWhenNotEnoughSlots(ScheduleMode scheduleMode) throws Exception {
	final JobVertex vertex = new JobVertex("Test Vertex");
	vertex.setParallelism(2);
	vertex.setMaxParallelism(2);
	vertex.setInvokableClass(BlockingNoOpInvokable.class);

	final JobGraph jobGraph = new JobGraph("Test Job", vertex);
	jobGraph.setScheduleMode(scheduleMode);

	runHandleJobsWhenNotEnoughSlots(jobGraph);
}
 

private static JobGraph getSimpleJob(int parallelism) throws IOException {
	final JobVertex task = new JobVertex("Test task");
	task.setParallelism(parallelism);
	task.setMaxParallelism(parallelism);
	task.setInvokableClass(NoOpInvokable.class);

	final JobGraph jg = new JobGraph(new JobID(), "Test Job", task);
	jg.setScheduleMode(ScheduleMode.EAGER);

	final ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000));
	jg.setExecutionConfig(executionConfig);

	return jg;
}
 

/**
 * SlotPool#failAllocation should not fail with a {@link java.util.ConcurrentModificationException}
 * if there is a concurrent scheduling operation. See FLINK-13421.
 */
@Test
public void slotPoolExecutionGraph_ConcurrentSchedulingAndAllocationFailure_ShouldNotFailWithConcurrentModificationException() throws Exception {
	final SlotSharingGroup group = new SlotSharingGroup();
	final JobVertex vertex1 = createNoOpVertex("vertex1", 1);
	vertex1.setSlotSharingGroup(group);
	final JobVertex vertex2 = createNoOpVertex("vertex2", 3);
	vertex2.setSlotSharingGroup(group);
	final JobVertex vertex3 = createNoOpVertex("vertex3", 1);
	vertex3.setSlotSharingGroup(group);
	vertex3.connectNewDataSetAsInput(vertex2, DistributionPattern.ALL_TO_ALL, ResultPartitionType.PIPELINED);

	try (SlotPool slotPool = createSlotPoolImpl()) {
		final SlotProvider slots = createSchedulerWithSlots(slotPool, new LocalTaskManagerLocation(), 2);

		final AllocationID allocationId = slotPool.getAvailableSlotsInformation().iterator().next().getAllocationId();

		final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Test Job", vertex1, vertex2, vertex3);
		jobGraph.setScheduleMode(ScheduleMode.EAGER);
		final ExecutionGraph eg = TestingExecutionGraphBuilder
			.newBuilder()
			.setJobGraph(jobGraph)
			.setSlotProvider(slots)
			.setAllocationTimeout(Time.minutes(60))
			.build();

		startAndScheduleExecutionGraph(eg);

		slotPool.failAllocation(
			allocationId,
			new Exception("test exception"));

		eg.waitUntilTerminal();
	}
}
 

/**
 * Tests that a partially completed eager scheduling operation fails if a
 * completed slot is released. See FLINK-9099.
 */
@Test
public void testSlotReleasingFailsSchedulingOperation() throws Exception {
	final int parallelism = 2;

	final JobVertex jobVertex = new JobVertex("Testing job vertex");
	jobVertex.setInvokableClass(NoOpInvokable.class);
	jobVertex.setParallelism(parallelism);
	final JobGraph jobGraph = new JobGraph(jobVertex);
	jobGraph.setAllowQueuedScheduling(true);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(parallelism);

	final SimpleSlot slot = createSlot(new SimpleAckingTaskManagerGateway(), jobGraph.getJobID(), new DummySlotOwner());
	slotProvider.addSlot(jobVertex.getID(), 0, CompletableFuture.completedFuture(slot));

	final CompletableFuture<LogicalSlot> slotFuture = new CompletableFuture<>();
	slotProvider.addSlot(jobVertex.getID(), 1, slotFuture);

	final ExecutionGraph executionGraph = createExecutionGraph(jobGraph, slotProvider);

	executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());
	executionGraph.scheduleForExecution();

	assertThat(executionGraph.getState(), is(JobStatus.RUNNING));

	final ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertex.getID());
	final ExecutionVertex[] taskVertices = executionJobVertex.getTaskVertices();
	assertThat(taskVertices[0].getExecutionState(), is(ExecutionState.SCHEDULED));
	assertThat(taskVertices[1].getExecutionState(), is(ExecutionState.SCHEDULED));

	// fail the single allocated slot --> this should fail the scheduling operation
	slot.releaseSlot(new FlinkException("Test failure"));

	assertThat(executionGraph.getTerminationFuture().get(), is(JobStatus.FAILED));
}
 

private static JobGraph singleJobVertexJobGraph(final int parallelism) {
	final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Testjob");
	jobGraph.setScheduleMode(ScheduleMode.EAGER);
	final JobVertex vertex = new JobVertex("source");
	vertex.setInvokableClass(NoOpInvokable.class);
	vertex.setParallelism(parallelism);
	jobGraph.addVertex(vertex);
	return jobGraph;
}
 

private static JobGraph getSimpleJob(int parallelism) throws IOException {
	final JobVertex task = new JobVertex("Test task");
	task.setParallelism(parallelism);
	task.setMaxParallelism(parallelism);
	task.setInvokableClass(NoOpInvokable.class);

	final JobGraph jg = new JobGraph(new JobID(), "Test Job", task);
	jg.setScheduleMode(ScheduleMode.EAGER);

	final ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000));
	jg.setExecutionConfig(executionConfig);

	return jg;
}
 

/**
 * Tests that a batch job can be executed with fewer slots than its parallelism.
 * See FLINK-13187 for more information.
 */
@Test
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception {
	final int parallelism = 5;
	final Time batchSlotTimeout = Time.milliseconds(5L);
	final JobGraph jobGraph = createJobGraph(parallelism);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST);

	try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) {
		final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism);
		TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setSubmitTaskConsumer(
				(tdd, ignored) -> {
					submittedTasksQueue.offer(tdd.getExecutionAttemptId());
					return CompletableFuture.completedFuture(Acknowledge.get());
				})
			.createTestingTaskExecutorGateway();

		// register a single slot at the slot pool
		SlotPoolUtils.offerSlots(
			slotPool,
			mainThreadExecutor,
			Collections.singletonList(ResourceProfile.ANY),
			new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate()));

		final SlotProvider slotProvider = createSlotProvider(slotPool, mainThreadExecutor);
		final SchedulerNG scheduler = createScheduler(jobGraph, slotProvider, batchSlotTimeout);

		final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener();
		scheduler.registerJobStatusListener(jobStatusListener);
		startScheduling(scheduler, mainThreadExecutor);

		// wait until the batch slot timeout has been reached
		Thread.sleep(batchSlotTimeout.toMilliseconds());

		final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture();

		for (int i = 0; i < parallelism; i++) {
			final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take));

			// wait until one of them is completed
			CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join();

			if (submittedTaskFuture.isDone()) {
				finishExecution(submittedTaskFuture.get(), scheduler, mainThreadExecutor);
			} else {
				fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get()));
			}
		}

		assertThat(terminationFuture.get(), is(JobStatus.FINISHED));
	}
}
 

@Test
public void testLocalFailAndRestart() throws Exception {
	final int parallelism = 10;
	final TestRestartStrategy triggeredRestartStrategy = TestRestartStrategy.manuallyTriggered();

	final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Test Job", createNoOpVertex(parallelism));
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final ExecutionGraph eg = TestingExecutionGraphBuilder
		.newBuilder()
		.setJobGraph(jobGraph)
		.setRestartStrategy(triggeredRestartStrategy)
		.build();

	startAndScheduleExecutionGraph(eg);

	switchToRunning(eg);

	final ExecutionJobVertex vertex = eg.getVerticesTopologically().iterator().next();
	final Execution first = vertex.getTaskVertices()[0].getCurrentExecutionAttempt();
	final Execution last = vertex.getTaskVertices()[vertex.getParallelism() - 1].getCurrentExecutionAttempt();

	// Have two executions fail
	first.fail(new Exception("intended test failure 1"));
	last.fail(new Exception("intended test failure 2"));

	assertEquals(JobStatus.FAILING, eg.getState());

	completeCancellingForAllVertices(eg);

	// Now trigger the restart
	assertEquals(1, triggeredRestartStrategy.getNumberOfQueuedActions());
	triggeredRestartStrategy.triggerAll().join();

	assertEquals(JobStatus.RUNNING, eg.getState());

	switchToRunning(eg);
	finishAllVertices(eg);

	eg.waitUntilTerminal();
	assertEquals(JobStatus.FINISHED, eg.getState());
}
 

/**
 * Tests that all slots are being returned to the {@link SlotOwner} if the
 * {@link ExecutionGraph} is being cancelled. See FLINK-9908
 */
@Test
public void testCancellationOfIncompleteScheduling() throws Exception {
	final int parallelism = 10;

	final JobVertex jobVertex = new JobVertex("Test job vertex");
	jobVertex.setInvokableClass(NoOpInvokable.class);
	jobVertex.setParallelism(parallelism);

	final JobGraph jobGraph = new JobGraph(jobVertex);
	jobGraph.setAllowQueuedScheduling(true);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final TestingSlotOwner slotOwner = new TestingSlotOwner();
	final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

	final ConcurrentMap<SlotRequestId, Integer> slotRequestIds = new ConcurrentHashMap<>(parallelism);

	final TestingSlotProvider slotProvider = new TestingSlotProvider(
		(SlotRequestId slotRequestId) -> {
			slotRequestIds.put(slotRequestId, 1);
			// return 50/50 fulfilled and unfulfilled requests
			return slotRequestIds.size() % 2 == 0 ?
				CompletableFuture.completedFuture(
					createSingleLogicalSlot(slotOwner, taskManagerGateway, slotRequestId)) :
				new CompletableFuture<>();
		});

	final ExecutionGraph executionGraph = createExecutionGraph(jobGraph, slotProvider);

	executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());
	final Set<SlotRequestId> slotRequestIdsToReturn = ConcurrentHashMap.newKeySet(slotRequestIds.size());

	executionGraph.scheduleForExecution();

	slotRequestIdsToReturn.addAll(slotRequestIds.keySet());

	slotOwner.setReturnAllocatedSlotConsumer(logicalSlot -> {
		slotRequestIdsToReturn.remove(logicalSlot.getSlotRequestId());
	});

	slotProvider.setSlotCanceller(slotRequestIdsToReturn::remove);

	// make sure that we complete cancellations of deployed tasks
	taskManagerGateway.setCancelConsumer(
		(ExecutionAttemptID executionAttemptId) -> {
			final Execution execution = executionGraph.getRegisteredExecutions().get(executionAttemptId);

			// if the execution was cancelled in state SCHEDULING, then it might already have been removed
			if (execution != null) {
				execution.completeCancelling();
			}
		}
	);

	executionGraph.cancel();
	assertThat(slotRequestIdsToReturn, is(empty()));
}
 

/**
 * Tests that with scheduling futures and pipelined deployment, the target vertex will
 * not deploy its task before the source vertex does.
 */
@Test
public void testScheduleSourceBeforeTarget() throws Exception {

	//                                            [pipelined]
	//  we construct a simple graph    (source) ----------------> (target)

	final int parallelism = 1;

	final JobVertex sourceVertex = new JobVertex("source");
	sourceVertex.setParallelism(parallelism);
	sourceVertex.setInvokableClass(NoOpInvokable.class);

	final JobVertex targetVertex = new JobVertex("target");
	targetVertex.setParallelism(parallelism);
	targetVertex.setInvokableClass(NoOpInvokable.class);

	targetVertex.connectNewDataSetAsInput(sourceVertex, DistributionPattern.ALL_TO_ALL, ResultPartitionType.PIPELINED);

	final JobID jobId = new JobID();
	final JobGraph jobGraph = new JobGraph(jobId, "test", sourceVertex, targetVertex);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);
	jobGraph.setAllowQueuedScheduling(true);

	final CompletableFuture<LogicalSlot> sourceFuture = new CompletableFuture<>();
	final CompletableFuture<LogicalSlot> targetFuture = new CompletableFuture<>();

	ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(parallelism);
	slotProvider.addSlot(sourceVertex.getID(), 0, sourceFuture);
	slotProvider.addSlot(targetVertex.getID(), 0, targetFuture);

	final ExecutionGraph eg = createExecutionGraph(jobGraph, slotProvider);
	eg.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	//  set up two TaskManager gateways and slots

	final InteractionsCountingTaskManagerGateway gatewaySource = createTaskManager();
	final InteractionsCountingTaskManagerGateway gatewayTarget = createTaskManager();

	final LogicalSlot sourceSlot = createTestingLogicalSlot(gatewaySource);
	final LogicalSlot targetSlot = createTestingLogicalSlot(gatewayTarget);

	eg.scheduleForExecution();

	// job should be running
	assertEquals(JobStatus.RUNNING, eg.getState());

	// we fulfill the target slot before the source slot
	// that should not cause a deployment or deployment related failure
	targetFuture.complete(targetSlot);

	assertThat(gatewayTarget.getSubmitTaskCount(), is(0));
	assertEquals(JobStatus.RUNNING, eg.getState());

	// now supply the source slot
	sourceFuture.complete(sourceSlot);

	// by now, all deployments should have happened
	assertThat(gatewaySource.getSubmitTaskCount(), is(1));
	assertThat(gatewayTarget.getSubmitTaskCount(), is(1));

	assertEquals(JobStatus.RUNNING, eg.getState());
}
 

/**
 * Tests that all slots are being returned to the {@link SlotOwner} if the
 * {@link ExecutionGraph} is being cancelled. See FLINK-9908
 */
@Test
public void testCancellationOfIncompleteScheduling() throws Exception {
	final int parallelism = 10;

	final JobVertex jobVertex = new JobVertex("Test job vertex");
	jobVertex.setInvokableClass(NoOpInvokable.class);
	jobVertex.setParallelism(parallelism);

	final JobGraph jobGraph = new JobGraph(jobVertex);
	jobGraph.setAllowQueuedScheduling(true);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final TestingSlotOwner slotOwner = new TestingSlotOwner();
	final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

	final ConcurrentMap<SlotRequestId, Integer> slotRequestIds = new ConcurrentHashMap<>(parallelism);

	final TestingSlotProvider slotProvider = new TestingSlotProvider(
		(SlotRequestId slotRequestId) -> {
			slotRequestIds.put(slotRequestId, 1);
			// return 50/50 fulfilled and unfulfilled requests
			return slotRequestIds.size() % 2 == 0 ?
				CompletableFuture.completedFuture(
					createSingleLogicalSlot(slotOwner, taskManagerGateway, slotRequestId)) :
				new CompletableFuture<>();
		});

	final ExecutionGraph executionGraph = createExecutionGraph(jobGraph, slotProvider);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());
	final Set<SlotRequestId> slotRequestIdsToReturn = ConcurrentHashMap.newKeySet(slotRequestIds.size());

	executionGraph.scheduleForExecution();

	slotRequestIdsToReturn.addAll(slotRequestIds.keySet());

	slotOwner.setReturnAllocatedSlotConsumer(logicalSlot -> {
		slotRequestIdsToReturn.remove(logicalSlot.getSlotRequestId());
	});

	slotProvider.setSlotCanceller(slotRequestIdsToReturn::remove);

	// make sure that we complete cancellations of deployed tasks
	taskManagerGateway.setCancelConsumer(
		(ExecutionAttemptID executionAttemptId) -> {
			final Execution execution = executionGraph.getRegisteredExecutions().get(executionAttemptId);

			// if the execution was cancelled in state SCHEDULING, then it might already have been removed
			if (execution != null) {
				execution.completeCancelling();
			}
		}
	);

	executionGraph.cancel();
	assertThat(slotRequestIdsToReturn, is(empty()));
}
 

/**
 * Tests that all slots are being returned to the {@link SlotOwner} if the
 * {@link ExecutionGraph} is being cancelled. See FLINK-9908
 */
@Test
public void testCancellationOfIncompleteScheduling() throws Exception {
	final int parallelism = 10;

	final JobVertex jobVertex = new JobVertex("Test job vertex");
	jobVertex.setInvokableClass(NoOpInvokable.class);
	jobVertex.setParallelism(parallelism);

	final JobGraph jobGraph = new JobGraph(jobVertex);
	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	final TestingSlotOwner slotOwner = new TestingSlotOwner();
	final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

	final ConcurrentMap<SlotRequestId, Integer> slotRequestIds = new ConcurrentHashMap<>(parallelism);

	final TestingSlotProvider slotProvider = new TestingSlotProvider(
		(SlotRequestId slotRequestId) -> {
			slotRequestIds.put(slotRequestId, 1);
			// return 50/50 fulfilled and unfulfilled requests
			return slotRequestIds.size() % 2 == 0 ?
				CompletableFuture.completedFuture(
					createSingleLogicalSlot(slotOwner, taskManagerGateway, slotRequestId)) :
				new CompletableFuture<>();
		});

	final ExecutionGraph executionGraph = createExecutionGraph(jobGraph, slotProvider);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());
	final Set<SlotRequestId> slotRequestIdsToReturn = ConcurrentHashMap.newKeySet(slotRequestIds.size());

	executionGraph.scheduleForExecution();

	slotRequestIdsToReturn.addAll(slotRequestIds.keySet());

	slotOwner.setReturnAllocatedSlotConsumer(logicalSlot -> {
		slotRequestIdsToReturn.remove(logicalSlot.getSlotRequestId());
	});

	slotProvider.setSlotCanceller(slotRequestIdsToReturn::remove);

	// make sure that we complete cancellations of deployed tasks
	taskManagerGateway.setCancelConsumer(
		(ExecutionAttemptID executionAttemptId) -> {
			final Execution execution = executionGraph.getRegisteredExecutions().get(executionAttemptId);

			// if the execution was cancelled in state SCHEDULING, then it might already have been removed
			if (execution != null) {
				execution.completeCancelling();
			}
		}
	);

	executionGraph.cancel();
	assertThat(slotRequestIdsToReturn, is(empty()));
}