Java Code Examples for org.apache.flink.runtime.jobgraph.JobVertex#getID()

@Test
public void handleGlobalFailure() {
	final JobGraph jobGraph = singleNonParallelJobVertexJobGraph();
	final JobVertex onlyJobVertex = getOnlyJobVertex(jobGraph);

	final DefaultScheduler scheduler = createSchedulerAndStartScheduling(jobGraph);

	scheduler.handleGlobalFailure(new Exception("forced failure"));

	final ArchivedExecutionVertex onlyExecutionVertex = Iterables.getOnlyElement(scheduler.requestJob().getAllExecutionVertices());
	final ExecutionAttemptID attemptId = onlyExecutionVertex.getCurrentExecutionAttempt().getAttemptId();
	scheduler.updateTaskExecutionState(new TaskExecutionState(jobGraph.getJobID(), attemptId, ExecutionState.CANCELED));

	taskRestartExecutor.triggerScheduledTasks();

	final List<ExecutionVertexID> deployedExecutionVertices = testExecutionVertexOperations.getDeployedVertices();
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(onlyJobVertex.getID(), 0);
	assertThat(deployedExecutionVertices, contains(executionVertexId, executionVertexId));
}
 

/**
 * Tests a fix for FLINK-1930.
 *
 * <p>When consuming a pipelined result only partially, is is possible that local channels
 * release the buffer pool, which is associated with the result partition, too early. If the
 * producer is still producing data when this happens, it runs into an IllegalStateException,
 * because of the destroyed buffer pool.
 *
 * @see <a href="https://issues.apache.org/jira/browse/FLINK-1930">FLINK-1930</a>
 */
@Test
public void testPartialConsumePipelinedResultReceiver() throws Exception {
	final JobVertex sender = new JobVertex("Sender");
	sender.setInvokableClass(SlowBufferSender.class);
	sender.setParallelism(PARALLELISM);

	final JobVertex receiver = new JobVertex("Receiver");
	receiver.setInvokableClass(SingleBufferReceiver.class);
	receiver.setParallelism(PARALLELISM);

	// The partition needs to be pipelined, otherwise the original issue does not occur, because
	// the sender and receiver are not online at the same time.
	receiver.connectNewDataSetAsInput(
		sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final JobGraph jobGraph = new JobGraph("Partial Consume of Pipelined Result", sender, receiver);

	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup(
		sender.getID(), receiver.getID());

	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setSlotSharingGroup(slotSharingGroup);

	MINI_CLUSTER_RESOURCE.getMiniCluster().executeJobBlocking(jobGraph);
}
 

@Test
public void cancelJobWillIncrementVertexVersions() {
	final JobGraph jobGraph = singleNonParallelJobVertexJobGraph();
	final JobVertex onlyJobVertex = getOnlyJobVertex(jobGraph);
	final ExecutionVertexID onlyExecutionVertexId = new ExecutionVertexID(onlyJobVertex.getID(), 0);

	final DefaultScheduler scheduler = createSchedulerAndStartScheduling(jobGraph);
	final ExecutionVertexVersion executionVertexVersion = executionVertexVersioner.getExecutionVertexVersion(
		onlyExecutionVertexId);

	scheduler.cancel();

	assertTrue(executionVertexVersioner.isModified(executionVertexVersion));
}
 

@Test
public void scheduleWithLazyStrategy() {
	final JobGraph jobGraph = singleNonParallelJobVertexJobGraph();
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES);
	final JobVertex onlyJobVertex = getOnlyJobVertex(jobGraph);

	createSchedulerAndStartScheduling(jobGraph);

	final List<ExecutionVertexID> deployedExecutionVertices = testExecutionVertexOperations.getDeployedVertices();

	final ExecutionVertexID executionVertexId = new ExecutionVertexID(onlyJobVertex.getID(), 0);
	assertThat(deployedExecutionVertices, contains(executionVertexId));
}
 

/**
 * Tests that it will get empty task manager location if vertex is not scheduled.
 */
@Test
public void testGetEmptyTaskManagerLocationIfVertexNotScheduled() throws Exception {
	final JobVertex jobVertex = ExecutionGraphTestUtils.createNoOpVertex(1);

	final ExecutionGraph eg = ExecutionGraphTestUtils.createSimpleTestGraph(new JobID(), jobVertex);
	final ExecutionGraphToInputsLocationsRetrieverAdapter inputsLocationsRetriever =
			new ExecutionGraphToInputsLocationsRetrieverAdapter(eg);

	ExecutionVertexID executionVertexId = new ExecutionVertexID(jobVertex.getID(), 0);
	Optional<CompletableFuture<TaskManagerLocation>> taskManagerLocation =
			inputsLocationsRetriever.getTaskManagerLocation(executionVertexId);

	assertFalse(taskManagerLocation.isPresent());
}
 

/**
 * Tests that can get the producers of consumed result partitions.
 */
@Test
public void testGetConsumedResultPartitionsProducers() throws Exception {
	final JobVertex producer1 = ExecutionGraphTestUtils.createNoOpVertex(1);
	final JobVertex producer2 = ExecutionGraphTestUtils.createNoOpVertex(1);
	final JobVertex consumer = ExecutionGraphTestUtils.createNoOpVertex(1);
	consumer.connectNewDataSetAsInput(producer1, DistributionPattern.ALL_TO_ALL, ResultPartitionType.PIPELINED);
	consumer.connectNewDataSetAsInput(producer2, DistributionPattern.ALL_TO_ALL, ResultPartitionType.PIPELINED);

	final ExecutionGraph eg = ExecutionGraphTestUtils.createSimpleTestGraph(new JobID(), producer1, producer2, consumer);
	final ExecutionGraphToInputsLocationsRetrieverAdapter inputsLocationsRetriever =
			new ExecutionGraphToInputsLocationsRetrieverAdapter(eg);

	ExecutionVertexID evIdOfProducer1 = new ExecutionVertexID(producer1.getID(), 0);
	ExecutionVertexID evIdOfProducer2 = new ExecutionVertexID(producer2.getID(), 0);
	ExecutionVertexID evIdOfConsumer = new ExecutionVertexID(consumer.getID(), 0);

	Collection<Collection<ExecutionVertexID>> producersOfProducer1 =
			inputsLocationsRetriever.getConsumedResultPartitionsProducers(evIdOfProducer1);
	Collection<Collection<ExecutionVertexID>> producersOfProducer2 =
			inputsLocationsRetriever.getConsumedResultPartitionsProducers(evIdOfProducer2);
	Collection<Collection<ExecutionVertexID>> producersOfConsumer =
			inputsLocationsRetriever.getConsumedResultPartitionsProducers(evIdOfConsumer);

	assertThat(producersOfProducer1, is(empty()));
	assertThat(producersOfProducer2, is(empty()));
	assertThat(producersOfConsumer, hasSize(2));
	assertThat(producersOfConsumer, hasItem(Collections.singletonList(evIdOfProducer1)));
	assertThat(producersOfConsumer, hasItem(Collections.singletonList(evIdOfProducer2)));
}
 

@Test
public void suspendJobWillIncrementVertexVersions() {
	final JobGraph jobGraph = singleNonParallelJobVertexJobGraph();
	final JobVertex onlyJobVertex = getOnlyJobVertex(jobGraph);
	final ExecutionVertexID onlyExecutionVertexId = new ExecutionVertexID(onlyJobVertex.getID(), 0);

	final DefaultScheduler scheduler = createSchedulerAndStartScheduling(jobGraph);
	final ExecutionVertexVersion executionVertexVersion = executionVertexVersioner.getExecutionVertexVersion(
		onlyExecutionVertexId);

	scheduler.suspend(new Exception("forced suspend"));

	assertTrue(executionVertexVersioner.isModified(executionVertexVersion));
}
 

/**
 * Tests that the task restore state is nulled after the {@link Execution} has been
 * deployed. See FLINK-9693.
 */
@Test
public void testTaskRestoreStateIsNulledAfterDeployment() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();
	final ProgrammedSlotProvider slotProvider = createProgrammedSlotProvider(
		1,
		Collections.singleton(jobVertexId),
		slotOwner);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	ExecutionVertex executionVertex = executionJobVertex.getTaskVertices()[0];

	final Execution execution = executionVertex.getCurrentExecutionAttempt();

	final JobManagerTaskRestore taskRestoreState = new JobManagerTaskRestore(1L, new TaskStateSnapshot());
	execution.setInitialState(taskRestoreState);

	assertThat(execution.getTaskRestore(), is(notNullValue()));

	// schedule the execution vertex and wait for its deployment
	executionVertex.scheduleForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ANY,
		Collections.emptySet())
		.get();

	assertThat(execution.getTaskRestore(), is(nullValue()));
}
 

@Test
public void failGlobalWhenRestoringStateFails() throws Exception {
	final JobGraph jobGraph = singleNonParallelJobVertexJobGraph();
	final JobVertex onlyJobVertex = getOnlyJobVertex(jobGraph);
	enableCheckpointing(jobGraph);

	final CountDownLatch checkpointTriggeredLatch = getCheckpointTriggeredLatch();

	final DefaultScheduler scheduler = createSchedulerAndStartScheduling(jobGraph);

	final ArchivedExecutionVertex onlyExecutionVertex = Iterables.getOnlyElement(scheduler.requestJob().getAllExecutionVertices());
	final ExecutionAttemptID attemptId = onlyExecutionVertex.getCurrentExecutionAttempt().getAttemptId();
	scheduler.updateTaskExecutionState(new TaskExecutionState(jobGraph.getJobID(), attemptId, ExecutionState.RUNNING));

	final CheckpointCoordinator checkpointCoordinator = getCheckpointCoordinator(scheduler);

	// register a master hook to fail state restore
	final TestMasterHook masterHook = TestMasterHook.fromId("testHook");
	masterHook.enableFailOnRestore();
	checkpointCoordinator.addMasterHook(masterHook);

	// complete one checkpoint for state restore
	checkpointCoordinator.triggerCheckpoint(false);
	checkpointTriggeredLatch.await();
	final long checkpointId = checkpointCoordinator.getPendingCheckpoints().keySet().iterator().next();
	acknowledgePendingCheckpoint(scheduler, checkpointId);

	scheduler.updateTaskExecutionState(new TaskExecutionState(jobGraph.getJobID(), attemptId, ExecutionState.FAILED));
	taskRestartExecutor.triggerScheduledTasks();
	final List<ExecutionVertexID> deployedExecutionVertices = testExecutionVertexOperations.getDeployedVertices();

	// the first task failover should be skipped on state restore failure
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(onlyJobVertex.getID(), 0);
	assertThat(deployedExecutionVertices, contains(executionVertexId));

	// a global failure should be triggered on state restore failure
	masterHook.disableFailOnRestore();
	taskRestartExecutor.triggerScheduledTasks();
	assertThat(deployedExecutionVertices, contains(executionVertexId, executionVertexId));
}
 

/**
 * Tests that (un)chaining affects the node hash (for intermediate nodes).
 *
 * <pre>
 * A (chained): [ (src0) -> (map) -> (filter) -> (sink) ]
 * B (unchained): [ (src0) ] -> [ (map) -> (filter) -> (sink) ]
 * </pre>
 *
 * <p>The hashes for the single vertex in A and the source vertex in B need to be different.
 */
@Test
public void testNodeHashAfterIntermediateUnchaining() throws Exception {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment();
	env.setParallelism(4);

	env.addSource(new NoOpSourceFunction())
			.map(new NoOpMapFunction()).name("map")
			.startNewChain()
			.filter(new NoOpFilterFunction())
			.addSink(new DiscardingSink<>());

	JobGraph jobGraph = env.getStreamGraph().getJobGraph();

	JobVertex chainedMap = jobGraph.getVerticesSortedTopologicallyFromSources().get(1);
	assertTrue(chainedMap.getName().startsWith("map"));
	JobVertexID chainedMapId = chainedMap.getID();

	env = StreamExecutionEnvironment.createLocalEnvironment();
	env.setParallelism(4);

	env.addSource(new NoOpSourceFunction())
			.map(new NoOpMapFunction()).name("map")
			.startNewChain()
			.filter(new NoOpFilterFunction())
			.startNewChain()
			.addSink(new DiscardingSink<>());

	jobGraph = env.getStreamGraph().getJobGraph();

	JobVertex unchainedMap = jobGraph.getVerticesSortedTopologicallyFromSources().get(1);
	assertEquals("map", unchainedMap.getName());
	JobVertexID unchainedMapId = unchainedMap.getID();

	assertNotEquals(chainedMapId, unchainedMapId);
}
 

/**
 * Tests that the slot is released in case of a execution cancellation when having
 * a slot assigned and being in state SCHEDULED.
 */
@Test
public void testSlotReleaseOnExecutionCancellationInScheduled() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();

	final LogicalSlot slot = createTestingLogicalSlot(slotOwner);

	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(1);
	slotProvider.addSlot(jobVertexId, 0, CompletableFuture.completedFuture(slot));

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	final Execution execution = executionJobVertex.getTaskVertices()[0].getCurrentExecutionAttempt();

	CompletableFuture<Execution> allocationFuture = execution.allocateResourcesForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ALL,
		Collections.emptySet());

	assertTrue(allocationFuture.isDone());

	assertEquals(ExecutionState.SCHEDULED, execution.getState());

	assertEquals(slot, execution.getAssignedResource());

	// cancelling the execution should move it into state CANCELED
	execution.cancel();
	assertEquals(ExecutionState.CANCELED, execution.getState());

	assertEquals(slot, slotOwner.getReturnedSlotFuture().get());
}
 

/**
 * Tests notifications of multiple receivers when a task produces both a pipelined and blocking
 * result.
 *
 * <pre>
 *                             +----------+
 *            +-- pipelined -> | Receiver |
 * +--------+ |                +----------+
 * | Sender |-|
 * +--------+ |                +----------+
 *            +-- blocking --> | Receiver |
 *                             +----------+
 * </pre>
 *
 * <p>The pipelined receiver gets deployed after the first buffer is available and the blocking
 * one after all subtasks are finished.
 */
@Test
public void testMixedPipelinedAndBlockingResults() throws Exception {
	final JobVertex sender = new JobVertex("Sender");
	sender.setInvokableClass(BinaryRoundRobinSubtaskIndexSender.class);
	sender.getConfiguration().setInteger(BinaryRoundRobinSubtaskIndexSender.CONFIG_KEY, PARALLELISM);
	sender.setParallelism(PARALLELISM);

	final JobVertex pipelinedReceiver = new JobVertex("Pipelined Receiver");
	pipelinedReceiver.setInvokableClass(SlotCountExceedingParallelismTest.SubtaskIndexReceiver.class);
	pipelinedReceiver.getConfiguration().setInteger(CONFIG_KEY, PARALLELISM);
	pipelinedReceiver.setParallelism(PARALLELISM);

	pipelinedReceiver.connectNewDataSetAsInput(
			sender,
			DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);

	final JobVertex blockingReceiver = new JobVertex("Blocking Receiver");
	blockingReceiver.setInvokableClass(SlotCountExceedingParallelismTest.SubtaskIndexReceiver.class);
	blockingReceiver.getConfiguration().setInteger(CONFIG_KEY, PARALLELISM);
	blockingReceiver.setParallelism(PARALLELISM);

	blockingReceiver.connectNewDataSetAsInput(sender,
			DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.BLOCKING);

	SlotSharingGroup slotSharingGroup = new SlotSharingGroup(
			sender.getID(), pipelinedReceiver.getID(), blockingReceiver.getID());

	sender.setSlotSharingGroup(slotSharingGroup);
	pipelinedReceiver.setSlotSharingGroup(slotSharingGroup);
	blockingReceiver.setSlotSharingGroup(slotSharingGroup);

	final JobGraph jobGraph = new JobGraph(
			"Mixed pipelined and blocking result",
			sender,
			pipelinedReceiver,
			blockingReceiver);

	MINI_CLUSTER_RESOURCE.getMiniCluster().executeJobBlocking(jobGraph);
}
 

/**
 * Checks that the {@link Execution} termination future is only completed after the
 * assigned slot has been released.
 *
 * <p>NOTE: This test only fails spuriously without the fix of this commit. Thus, one has
 * to execute this test multiple times to see the failure.
 */
@Test
public void testTerminationFutureIsCompletedAfterSlotRelease() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();
	final ProgrammedSlotProvider slotProvider = createProgrammedSlotProvider(
		1,
		Collections.singleton(jobVertexId),
		slotOwner);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		new JobID(),
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	ExecutionVertex executionVertex = executionJobVertex.getTaskVertices()[0];

	executionVertex.scheduleForExecution(slotProvider, false, LocationPreferenceConstraint.ANY, Collections.emptySet()).get();

	Execution currentExecutionAttempt = executionVertex.getCurrentExecutionAttempt();

	CompletableFuture<LogicalSlot> returnedSlotFuture = slotOwner.getReturnedSlotFuture();
	CompletableFuture<?> terminationFuture = executionVertex.cancel();

	currentExecutionAttempt.completeCancelling();

	CompletableFuture<Boolean> restartFuture = terminationFuture.thenApply(
		ignored -> {
			assertTrue(returnedSlotFuture.isDone());
			return true;
		});

	// check if the returned slot future was completed first
	restartFuture.get();
}
 

/**
 * Tests that slots are released if we cannot assign the allocated resource to the
 * Execution.
 */
@Test
public void testSlotReleaseOnFailedResourceAssignment() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final CompletableFuture<LogicalSlot> slotFuture = new CompletableFuture<>();
	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(1);
	slotProvider.addSlot(jobVertexId, 0, slotFuture);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		new JobID(),
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	final Execution execution = executionJobVertex.getTaskVertices()[0].getCurrentExecutionAttempt();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();

	final SimpleSlot slot = new SimpleSlot(
		slotOwner,
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());

	final LogicalSlot otherSlot = new TestingLogicalSlot();

	CompletableFuture<Execution> allocationFuture = execution.allocateAndAssignSlotForExecution(
		slotProvider,
		false,
		LocationPreferenceConstraint.ALL,
		Collections.emptySet(),
		TestingUtils.infiniteTime());

	assertFalse(allocationFuture.isDone());

	assertEquals(ExecutionState.SCHEDULED, execution.getState());

	// assign a different resource to the execution
	assertTrue(execution.tryAssignResource(otherSlot));

	// completing now the future should cause the slot to be released
	slotFuture.complete(slot);

	assertEquals(slot, slotOwner.getReturnedSlotFuture().get());
}
 

/**
 * Checks that the {@link Execution} termination future is only completed after the
 * assigned slot has been released.
 *
 * <p>NOTE: This test only fails spuriously without the fix of this commit. Thus, one has
 * to execute this test multiple times to see the failure.
 */
@Test
public void testTerminationFutureIsCompletedAfterSlotRelease() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();
	final ProgrammedSlotProvider slotProvider = createProgrammedSlotProvider(
		1,
		Collections.singleton(jobVertexId),
		slotOwner);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	ExecutionVertex executionVertex = executionJobVertex.getTaskVertices()[0];

	executionVertex.scheduleForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ANY,
		Collections.emptySet()).get();

	Execution currentExecutionAttempt = executionVertex.getCurrentExecutionAttempt();

	CompletableFuture<LogicalSlot> returnedSlotFuture = slotOwner.getReturnedSlotFuture();
	CompletableFuture<?> terminationFuture = executionVertex.cancel();

	currentExecutionAttempt.completeCancelling();

	CompletableFuture<Boolean> restartFuture = terminationFuture.thenApply(
		ignored -> {
			assertTrue(returnedSlotFuture.isDone());
			return true;
		});

	// check if the returned slot future was completed first
	restartFuture.get();
}
 

/**
 * Tests that the slot is released in case of a execution cancellation when having
 * a slot assigned and being in state SCHEDULED.
 */
@Test
public void testSlotReleaseOnExecutionCancellationInScheduled() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();

	final LogicalSlot slot = createTestingLogicalSlot(slotOwner);

	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(1);
	slotProvider.addSlot(jobVertexId, 0, CompletableFuture.completedFuture(slot));

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		new JobID(),
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	final Execution execution = executionJobVertex.getTaskVertices()[0].getCurrentExecutionAttempt();

	CompletableFuture<Execution> allocationFuture = execution.allocateResourcesForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ALL,
		Collections.emptySet());

	assertTrue(allocationFuture.isDone());

	assertEquals(ExecutionState.SCHEDULED, execution.getState());

	assertEquals(slot, execution.getAssignedResource());

	// cancelling the execution should move it into state CANCELED
	execution.cancel();
	assertEquals(ExecutionState.CANCELED, execution.getState());

	assertEquals(slot, slotOwner.getReturnedSlotFuture().get());
}
 

/**
 * Tests that slots are released if we cannot assign the allocated resource to the
 * Execution.
 */
@Test
public void testSlotReleaseOnFailedResourceAssignment() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final CompletableFuture<LogicalSlot> slotFuture = new CompletableFuture<>();
	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(1);
	slotProvider.addSlot(jobVertexId, 0, slotFuture);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	final Execution execution = executionJobVertex.getTaskVertices()[0].getCurrentExecutionAttempt();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();

	final LogicalSlot slot = createTestingLogicalSlot(slotOwner);

	final LogicalSlot otherSlot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();

	CompletableFuture<Execution> allocationFuture = execution.allocateResourcesForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ALL,
		Collections.emptySet());

	assertFalse(allocationFuture.isDone());

	assertEquals(ExecutionState.SCHEDULED, execution.getState());

	// assign a different resource to the execution
	assertTrue(execution.tryAssignResource(otherSlot));

	// completing now the future should cause the slot to be released
	slotFuture.complete(slot);

	assertEquals(slot, slotOwner.getReturnedSlotFuture().get());
}
 

/**
 * Checks that the {@link Execution} termination future is only completed after the
 * assigned slot has been released.
 *
 * <p>NOTE: This test only fails spuriously without the fix of this commit. Thus, one has
 * to execute this test multiple times to see the failure.
 */
@Test
public void testTerminationFutureIsCompletedAfterSlotRelease() throws Exception {
	final JobVertex jobVertex = createNoOpJobVertex();
	final JobVertexID jobVertexId = jobVertex.getID();

	final SingleSlotTestingSlotOwner slotOwner = new SingleSlotTestingSlotOwner();
	final ProgrammedSlotProvider slotProvider = createProgrammedSlotProvider(
		1,
		Collections.singleton(jobVertexId),
		slotOwner);

	ExecutionGraph executionGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		new JobID(),
		slotProvider,
		new NoRestartStrategy(),
		jobVertex);

	executionGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	ExecutionJobVertex executionJobVertex = executionGraph.getJobVertex(jobVertexId);

	ExecutionVertex executionVertex = executionJobVertex.getTaskVertices()[0];

	executionVertex.scheduleForExecution(
		executionGraph.getSlotProviderStrategy(),
		LocationPreferenceConstraint.ANY,
		Collections.emptySet()).get();

	Execution currentExecutionAttempt = executionVertex.getCurrentExecutionAttempt();

	CompletableFuture<LogicalSlot> returnedSlotFuture = slotOwner.getReturnedSlotFuture();
	CompletableFuture<?> terminationFuture = executionVertex.cancel();

	currentExecutionAttempt.completeCancelling();

	CompletableFuture<Boolean> restartFuture = terminationFuture.thenApply(
		ignored -> {
			assertTrue(returnedSlotFuture.isDone());
			return true;
		});

	// check if the returned slot future was completed first
	restartFuture.get();
}
 

private void runRequestNextInputSplitTest(Function<List<List<InputSplit>>, Collection<InputSplit>> expectedRemainingInputSplits) throws Exception {
	final int parallelism = 2;
	final int splitsPerTask = 2;

	final int totalSplits = parallelism * splitsPerTask;
	final List<TestingInputSplit> allInputSplits = new ArrayList<>(totalSplits);

	for (int i = 0; i < totalSplits; i++) {
		allInputSplits.add(new TestingInputSplit(i));
	}

	final InputSplitSource<TestingInputSplit> inputSplitSource = new TestingInputSplitSource(allInputSplits);

	JobVertex source = new JobVertex("source");
	source.setParallelism(parallelism);
	source.setInputSplitSource(inputSplitSource);
	source.setInvokableClass(AbstractInvokable.class);

	final JobGraph inputSplitJobGraph = new JobGraph(source);

	final ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(100, 0));
	inputSplitJobGraph.setExecutionConfig(executionConfig);

	final JobMaster jobMaster = createJobMaster(
		configuration,
		inputSplitJobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

		final JobVertexID sourceId = source.getID();

		final List<AccessExecution> executions = getExecutions(jobMasterGateway, sourceId);
		final ExecutionAttemptID initialAttemptId = executions.get(0).getAttemptId();
		final List<List<InputSplit>> inputSplitsPerTask = new ArrayList<>(parallelism);

		// request all input splits
		for (AccessExecution execution : executions) {
			inputSplitsPerTask.add(getInputSplits(splitsPerTask, getInputSplitSupplier(sourceId, jobMasterGateway, execution.getAttemptId())));
		}

		final List<InputSplit> allRequestedInputSplits = flattenCollection(inputSplitsPerTask);
		assertThat(allRequestedInputSplits, containsInAnyOrder(allInputSplits.toArray(EMPTY_TESTING_INPUT_SPLITS)));

		waitUntilAllExecutionsAreScheduled(jobMasterGateway);

		// fail the first execution to trigger a failover
		jobMasterGateway.updateTaskExecutionState(new TaskExecutionState(inputSplitJobGraph.getJobID(), initialAttemptId, ExecutionState.FAILED)).get();

		// wait until the job has been recovered
		waitUntilAllExecutionsAreScheduled(jobMasterGateway);

		final ExecutionAttemptID restartedAttemptId = getFirstExecution(jobMasterGateway, sourceId).getAttemptId();

		final List<InputSplit> inputSplits = getRemainingInputSplits(getInputSplitSupplier(sourceId, jobMasterGateway, restartedAttemptId));

		assertThat(inputSplits, containsInAnyOrder(expectedRemainingInputSplits.apply(inputSplitsPerTask).toArray(EMPTY_TESTING_INPUT_SPLITS)));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@Test
public void testSchedulingAllAtOnce() throws Exception {
	final int parallelism = 11;

	final MiniClusterConfiguration cfg = new MiniClusterConfiguration.Builder()
		.setNumTaskManagers(1)
		.setNumSlotsPerTaskManager(parallelism)
		.setConfiguration(getDefaultConfiguration())
		.build();

	try (final MiniCluster miniCluster = new MiniCluster(cfg)) {
		miniCluster.start();

		final JobVertex sender = new JobVertex("Sender");
		sender.setInvokableClass(Sender.class);
		sender.setParallelism(parallelism);

		final JobVertex forwarder = new JobVertex("Forwarder");
		forwarder.setInvokableClass(Forwarder.class);
		forwarder.setParallelism(parallelism);

		final JobVertex receiver = new JobVertex("Receiver");
		receiver.setInvokableClass(AgnosticReceiver.class);
		receiver.setParallelism(parallelism);

		final SlotSharingGroup sharingGroup = new SlotSharingGroup(sender.getID(), receiver.getID());
		sender.setSlotSharingGroup(sharingGroup);
		forwarder.setSlotSharingGroup(sharingGroup);
		receiver.setSlotSharingGroup(sharingGroup);

		forwarder.connectNewDataSetAsInput(sender, DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);
		receiver.connectNewDataSetAsInput(forwarder, DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);

		final JobGraph jobGraph = new JobGraph("Forwarding Job", sender, forwarder, receiver);

		jobGraph.setScheduleMode(ScheduleMode.EAGER);

		miniCluster.executeJobBlocking(jobGraph);
	}
}