org.apache.flink.runtime.jobmanager.scheduler.LocationPreferenceConstraint Java Examples

/**
 * Tests that any preferred locations are calculated.
 */
@Test
public void testAnyPreferredLocationCalculation() throws Exception {
	final TaskManagerLocation taskManagerLocation1 = new LocalTaskManagerLocation();
	final TaskManagerLocation taskManagerLocation3 = new LocalTaskManagerLocation();

	final CompletableFuture<TaskManagerLocation> locationFuture1 = CompletableFuture.completedFuture(taskManagerLocation1);
	final CompletableFuture<TaskManagerLocation> locationFuture2 = new CompletableFuture<>();
	final CompletableFuture<TaskManagerLocation> locationFuture3 = CompletableFuture.completedFuture(taskManagerLocation3);

	final Execution execution = getExecution(Arrays.asList(locationFuture1, locationFuture2, locationFuture3));

	CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = execution.calculatePreferredLocations(LocationPreferenceConstraint.ANY);

	assertTrue(preferredLocationsFuture.isDone());

	final Collection<TaskManagerLocation> preferredLocations = preferredLocationsFuture.get();

	assertThat(preferredLocations, containsInAnyOrder(taskManagerLocation1, taskManagerLocation3));
}
 

@Test
public void testScheduleToDeploying() {
	try {
		final ExecutionJobVertex ejv = getExecutionJobVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

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

		CompletableFuture<LogicalSlot> future = CompletableFuture.completedFuture(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());

		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider(ignore -> future)),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());
		assertEquals(ExecutionState.DEPLOYING, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Schedules all execution vertices of this ExecutionJobVertex.
 *
 * @param slotProvider to allocate the slots from
 * @param queued if the allocations can be queued
 * @param locationPreferenceConstraint constraint for the location preferences
 * @param allPreviousExecutionGraphAllocationIds set with all previous allocation ids in the job graph.
 *                                                 Can be empty if the allocation ids are not required for scheduling.
 * @return Future which is completed once all {@link Execution} could be deployed
 */
public CompletableFuture<Void> scheduleAll(
		SlotProvider slotProvider,
		boolean queued,
		LocationPreferenceConstraint locationPreferenceConstraint,
		@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds) {

	final ExecutionVertex[] vertices = this.taskVertices;

	final ArrayList<CompletableFuture<Void>> scheduleFutures = new ArrayList<>(vertices.length);

	// kick off the tasks
	for (ExecutionVertex ev : vertices) {
		scheduleFutures.add(ev.scheduleForExecution(
			slotProvider,
			queued,
			locationPreferenceConstraint,
			allPreviousExecutionGraphAllocationIds));
	}

	return FutureUtils.waitForAll(scheduleFutures);
}
 

/**
 * Tests that any preferred locations are calculated.
 */
@Test
public void testAnyPreferredLocationCalculation() throws ExecutionException, InterruptedException {
	final TaskManagerLocation taskManagerLocation1 = new LocalTaskManagerLocation();
	final TaskManagerLocation taskManagerLocation3 = new LocalTaskManagerLocation();

	final CompletableFuture<TaskManagerLocation> locationFuture1 = CompletableFuture.completedFuture(taskManagerLocation1);
	final CompletableFuture<TaskManagerLocation> locationFuture2 = new CompletableFuture<>();
	final CompletableFuture<TaskManagerLocation> locationFuture3 = CompletableFuture.completedFuture(taskManagerLocation3);

	final Execution execution = SchedulerTestUtils.getTestVertex(Arrays.asList(locationFuture1, locationFuture2, locationFuture3));

	CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = execution.calculatePreferredLocations(LocationPreferenceConstraint.ANY);

	assertTrue(preferredLocationsFuture.isDone());

	final Collection<TaskManagerLocation> preferredLocations = preferredLocationsFuture.get();

	assertThat(preferredLocations, containsInAnyOrder(taskManagerLocation1, taskManagerLocation3));
}
 

@Test
public void testScheduleToDeploying() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		final LogicalSlot slot = new TestingLogicalSlot();

		CompletableFuture<LogicalSlot> future = CompletableFuture.completedFuture(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());

		// try to deploy to the slot
		vertex.scheduleForExecution(
			new TestingSlotProvider(ignore -> future),
			false,
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());
		assertEquals(ExecutionState.DEPLOYING, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

private CompletableFuture<Void> scheduleLazy(SlotProvider slotProvider) {

		final ArrayList<CompletableFuture<Void>> schedulingFutures = new ArrayList<>(numVerticesTotal);
		// simply take the vertices without inputs.
		for (ExecutionJobVertex ejv : verticesInCreationOrder) {
			if (ejv.getJobVertex().isInputVertex()) {
				final CompletableFuture<Void> schedulingJobVertexFuture = ejv.scheduleAll(
					slotProvider,
					allowQueuedScheduling,
					LocationPreferenceConstraint.ALL, // since it is an input vertex, the input based location preferences should be empty
					Collections.emptySet());

				schedulingFutures.add(schedulingJobVertexFuture);
			}
		}

		return FutureUtils.waitForAll(schedulingFutures);
	}
 

@Test
public void testScheduleToDeploying() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

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

		CompletableFuture<LogicalSlot> future = CompletableFuture.completedFuture(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());

		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider(ignore -> future), false),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());
		assertEquals(ExecutionState.DEPLOYING, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests that any preferred locations are calculated.
 */
@Test
public void testAnyPreferredLocationCalculation() throws ExecutionException, InterruptedException {
	final TaskManagerLocation taskManagerLocation1 = new LocalTaskManagerLocation();
	final TaskManagerLocation taskManagerLocation3 = new LocalTaskManagerLocation();

	final CompletableFuture<TaskManagerLocation> locationFuture1 = CompletableFuture.completedFuture(taskManagerLocation1);
	final CompletableFuture<TaskManagerLocation> locationFuture2 = new CompletableFuture<>();
	final CompletableFuture<TaskManagerLocation> locationFuture3 = CompletableFuture.completedFuture(taskManagerLocation3);

	final Execution execution = SchedulerTestUtils.getTestVertex(Arrays.asList(locationFuture1, locationFuture2, locationFuture3));

	CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = execution.calculatePreferredLocations(LocationPreferenceConstraint.ANY);

	assertTrue(preferredLocationsFuture.isDone());

	final Collection<TaskManagerLocation> preferredLocations = preferredLocationsFuture.get();

	assertThat(preferredLocations, containsInAnyOrder(taskManagerLocation1, taskManagerLocation3));
}
 

public CompletableFuture<Void> scheduleForExecution() {
	final ExecutionGraph executionGraph = getVertex().getExecutionGraph();
	final SlotProviderStrategy resourceProvider = executionGraph.getSlotProviderStrategy();
	return scheduleForExecution(
		resourceProvider,
		LocationPreferenceConstraint.ANY,
		Collections.emptySet());
}
 

private void scheduleConsumer(ExecutionVertex consumerVertex) {
	try {
		final ExecutionGraph executionGraph = consumerVertex.getExecutionGraph();
		consumerVertex.scheduleForExecution(
			executionGraph.getSlotProviderStrategy(),
			LocationPreferenceConstraint.ANY, // there must be at least one known location
			Collections.emptySet());
	} catch (Throwable t) {
		consumerVertex.fail(new IllegalStateException("Could not schedule consumer " +
			"vertex " + consumerVertex, t));
	}
}
 

/**
 * Calculates the preferred locations based on the location preference constraint.
 *
 * @param locationPreferenceConstraint constraint for the location preference
 * @return Future containing the collection of preferred locations. This might not be completed if not all inputs
 * 		have been a resource assigned.
 */
@VisibleForTesting
public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) {
	final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations();
	final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture;

	switch(locationPreferenceConstraint) {
		case ALL:
			preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures);
			break;
		case ANY:
			final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size());

			for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) {
				if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) {
					final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null);

					if (taskManagerLocation == null) {
						throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug.");
					}

					completedTaskManagerLocations.add(taskManagerLocation);
				}
			}

			preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations);
			break;
		default:
			throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.');
	}

	return preferredLocationsFuture;
}
 

/**
 * Schedule vertices lazy. That means only vertices satisfying its input constraint will be scheduled.
 *
 * @param vertices Topologically sorted vertices to schedule.
 * @param executionGraph The graph the given vertices belong to.
 */
public static CompletableFuture<Void> scheduleLazy(
		final Iterable<ExecutionVertex> vertices,
		final ExecutionGraph executionGraph) {

	executionGraph.assertRunningInJobMasterMainThread();

	final SlotProviderStrategy slotProviderStrategy = executionGraph.getSlotProviderStrategy();
	final Set<AllocationID> previousAllocations = computePriorAllocationIdsIfRequiredByScheduling(
		vertices, slotProviderStrategy.asSlotProvider());

	final ArrayList<CompletableFuture<Void>> schedulingFutures = new ArrayList<>();
	for (ExecutionVertex executionVertex : vertices) {
		// only schedule vertex when its input constraint is satisfied
		if (executionVertex.getJobVertex().getJobVertex().isInputVertex() ||
			executionVertex.checkInputDependencyConstraints()) {

			final CompletableFuture<Void> schedulingVertexFuture = executionVertex.scheduleForExecution(
				slotProviderStrategy,
				LocationPreferenceConstraint.ANY,
				previousAllocations);

			schedulingFutures.add(schedulingVertexFuture);
		}
	}

	return FutureUtils.waitForAll(schedulingFutures);
}
 

@Test
public void testSlotReleasedWhenScheduledImmediately() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		CompletableFuture<LogicalSlot> future = new CompletableFuture<>();
		future.complete(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider((i) -> future), false),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * 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(
		new JobID(),
		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 testSlotReleasedWhenScheduledQueued() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		final CompletableFuture<LogicalSlot> future = new CompletableFuture<>();

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider(ignore -> future), true),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());

		// future has not yet a slot
		assertEquals(ExecutionState.SCHEDULED, vertex.getExecutionState());

		future.complete(slot);

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Calculates the preferred locations based on the location preference constraint.
 *
 * @param locationPreferenceConstraint constraint for the location preference
 * @return Future containing the collection of preferred locations. This might not be completed if not all inputs
 * 		have been a resource assigned.
 */
@VisibleForTesting
public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) {
	final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations();
	final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture;

	switch(locationPreferenceConstraint) {
		case ALL:
			preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures);
			break;
		case ANY:
			final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size());

			for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) {
				if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) {
					final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null);

					if (taskManagerLocation == null) {
						throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug.");
					}

					completedTaskManagerLocations.add(taskManagerLocation);
				}
			}

			preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations);
			break;
		default:
			throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.');
	}

	return preferredLocationsFuture;
}
 

/**
 * Schedule vertices lazy. That means only vertices satisfying its input constraint will be scheduled.
 *
 * @param vertices Topologically sorted vertices to schedule.
 * @param executionGraph The graph the given vertices belong to.
 */
public static CompletableFuture<Void> scheduleLazy(
		final Iterable<ExecutionVertex> vertices,
		final ExecutionGraph executionGraph) {

	executionGraph.assertRunningInJobMasterMainThread();

	final SlotProviderStrategy slotProviderStrategy = executionGraph.getSlotProviderStrategy();
	final Set<AllocationID> previousAllocations = computePriorAllocationIdsIfRequiredByScheduling(
		vertices, slotProviderStrategy.asSlotProvider());

	final ArrayList<CompletableFuture<Void>> schedulingFutures = new ArrayList<>();
	for (ExecutionVertex executionVertex : vertices) {
		// only schedule vertex when its input constraint is satisfied
		if (executionVertex.getJobVertex().getJobVertex().isInputVertex() ||
			executionVertex.checkInputDependencyConstraints()) {

			final CompletableFuture<Void> schedulingVertexFuture = executionVertex.scheduleForExecution(
				slotProviderStrategy,
				LocationPreferenceConstraint.ANY,
				previousAllocations);

			schedulingFutures.add(schedulingVertexFuture);
		}
	}

	return FutureUtils.waitForAll(schedulingFutures);
}
 

@Test
public void testSlotReleasedWhenScheduledImmediately() {
	try {
		final ExecutionJobVertex ejv = getExecutionJobVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		CompletableFuture<LogicalSlot> future = new CompletableFuture<>();
		future.complete(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider((i) -> future)),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testSlotReleasedWhenScheduledQueued() {
	try {
		final ExecutionJobVertex ejv = getExecutionJobVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		final CompletableFuture<LogicalSlot> future = new CompletableFuture<>();

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(
			TestingSlotProviderStrategy.from(new TestingSlotProvider(ignore -> future)),
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());

		// future has not yet a slot
		assertEquals(ExecutionState.SCHEDULED, vertex.getExecutionState());

		future.complete(slot);

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

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

/**
 * Schedules the current execution of this ExecutionVertex.
 *
 * @param slotProvider to allocate the slots from
 * @param queued if the allocation can be queued
 * @param locationPreferenceConstraint constraint for the location preferences
 * @param allPreviousExecutionGraphAllocationIds set with all previous allocation ids in the job graph.
 *                                                 Can be empty if the allocation ids are not required for scheduling.
 * @return Future which is completed once the execution is deployed. The future
 * can also completed exceptionally.
 */
public CompletableFuture<Void> scheduleForExecution(
		SlotProvider slotProvider,
		boolean queued,
		LocationPreferenceConstraint locationPreferenceConstraint,
		@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds) {
	return this.currentExecution.scheduleForExecution(
		slotProvider,
		queued,
		locationPreferenceConstraint,
		allPreviousExecutionGraphAllocationIds);
}
 

public CompletableFuture<Void> scheduleForExecution() {
	final ExecutionGraph executionGraph = getVertex().getExecutionGraph();
	final SlotProviderStrategy resourceProvider = executionGraph.getSlotProviderStrategy();
	return scheduleForExecution(
		resourceProvider,
		LocationPreferenceConstraint.ANY,
		Collections.emptySet());
}
 

/**
 * 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(
		new JobID(),
		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(slotProvider, false, LocationPreferenceConstraint.ANY, Collections.emptySet()).get();

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

public CompletableFuture<Void> scheduleForExecution() {
	final ExecutionGraph executionGraph = getVertex().getExecutionGraph();
	final SlotProvider resourceProvider = executionGraph.getSlotProvider();
	final boolean allowQueued = executionGraph.isQueuedSchedulingAllowed();
	return scheduleForExecution(
		resourceProvider,
		allowQueued,
		LocationPreferenceConstraint.ANY,
		Collections.emptySet());
}
 

private void scheduleConsumer(ExecutionVertex consumerVertex) {
	try {
		final ExecutionGraph executionGraph = consumerVertex.getExecutionGraph();
		consumerVertex.scheduleForExecution(
			executionGraph.getSlotProvider(),
			executionGraph.isQueuedSchedulingAllowed(),
			LocationPreferenceConstraint.ANY, // there must be at least one known location
			Collections.emptySet());
	} catch (Throwable t) {
		consumerVertex.fail(new IllegalStateException("Could not schedule consumer " +
			"vertex " + consumerVertex, t));
	}
}
 

/**
 * Calculates the preferred locations based on the location preference constraint.
 *
 * @param locationPreferenceConstraint constraint for the location preference
 * @return Future containing the collection of preferred locations. This might not be completed if not all inputs
 * 		have been a resource assigned.
 */
@VisibleForTesting
public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) {
	final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations();
	final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture;

	switch(locationPreferenceConstraint) {
		case ALL:
			preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures);
			break;
		case ANY:
			final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size());

			for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) {
				if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) {
					final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null);

					if (taskManagerLocation == null) {
						throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug.");
					}

					completedTaskManagerLocations.add(taskManagerLocation);
				}
			}

			preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations);
			break;
		default:
			throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.');
	}

	return preferredLocationsFuture;
}
 

/**
 * Acquires a slot for all the execution vertices of this ExecutionJobVertex. The method returns
 * pairs of the slots and execution attempts, to ease correlation between vertices and execution
 * attempts.
 *
 * <p>If this method throws an exception, it makes sure to release all so far requested slots.
 *
 * @param resourceProvider The resource provider from whom the slots are requested.
 * @param queued if the allocation can be queued
 * @param locationPreferenceConstraint constraint for the location preferences
 * @param allPreviousExecutionGraphAllocationIds the allocation ids of all previous executions in the execution job graph.
 * @param allocationTimeout timeout for allocating the individual slots
 */
public Collection<CompletableFuture<Execution>> allocateResourcesForAll(
		SlotProvider resourceProvider,
		boolean queued,
		LocationPreferenceConstraint locationPreferenceConstraint,
		@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds,
		Time allocationTimeout) {
	final ExecutionVertex[] vertices = this.taskVertices;

	@SuppressWarnings("unchecked")
	final CompletableFuture<Execution>[] slots = new CompletableFuture[vertices.length];

	// try to acquire a slot future for each execution.
	// we store the execution with the future just to be on the safe side
	for (int i = 0; i < vertices.length; i++) {
		// allocate the next slot (future)
		final Execution exec = vertices[i].getCurrentExecutionAttempt();
		final CompletableFuture<Execution> allocationFuture = exec.allocateAndAssignSlotForExecution(
			resourceProvider,
			queued,
			locationPreferenceConstraint,
			allPreviousExecutionGraphAllocationIds,
			allocationTimeout);
		slots[i] = allocationFuture;
	}

	// all good, we acquired all slots
	return Arrays.asList(slots);
}
 

@Test
public void testSlotReleasedWhenScheduledImmediately() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		CompletableFuture<LogicalSlot> future = new CompletableFuture<>();
		future.complete(slot);

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(new TestingSlotProvider((i) -> future), false, LocationPreferenceConstraint.ALL, Collections.emptySet());

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testSlotReleasedWhenScheduledQueued() {
	try {
		final ExecutionJobVertex ejv = getExecutionVertex(new JobVertexID());
		final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
				AkkaUtils.getDefaultTimeout());

		// a slot than cannot be deployed to
		final LogicalSlot slot = new TestingLogicalSlot();
		slot.releaseSlot(new Exception("Test Exception"));

		assertFalse(slot.isAlive());

		final CompletableFuture<LogicalSlot> future = new CompletableFuture<>();

		assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
		// try to deploy to the slot
		vertex.scheduleForExecution(
			new TestingSlotProvider(ignore -> future),
			true,
			LocationPreferenceConstraint.ALL,
			Collections.emptySet());

		// future has not yet a slot
		assertEquals(ExecutionState.SCHEDULED, vertex.getExecutionState());

		future.complete(slot);

		// will have failed
		assertEquals(ExecutionState.FAILED, vertex.getExecutionState());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Allocates and assigns a slot obtained from the slot provider to the execution.
 *
 * @param slotProviderStrategy to obtain a new slot from
 * @param locationPreferenceConstraint constraint for the location preferences
 * @param allPreviousExecutionGraphAllocationIds set with all previous allocation ids in the job graph.
 *                                                 Can be empty if the allocation ids are not required for scheduling.
 * @return Future which is completed with the allocated slot once it has been assigned
 * 			or with an exception if an error occurred.
 */
private CompletableFuture<LogicalSlot> allocateAndAssignSlotForExecution(
		SlotProviderStrategy slotProviderStrategy,
		LocationPreferenceConstraint locationPreferenceConstraint,
		@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds) {

	checkNotNull(slotProviderStrategy);

	assertRunningInJobMasterMainThread();

	final SlotSharingGroup sharingGroup = vertex.getJobVertex().getSlotSharingGroup();
	final CoLocationConstraint locationConstraint = vertex.getLocationConstraint();

	// sanity check
	if (locationConstraint != null && sharingGroup == null) {
		throw new IllegalStateException(
				"Trying to schedule with co-location constraint but without slot sharing allowed.");
	}

	// this method only works if the execution is in the state 'CREATED'
	if (transitionState(CREATED, SCHEDULED)) {

		final SlotSharingGroupId slotSharingGroupId = sharingGroup != null ? sharingGroup.getSlotSharingGroupId() : null;

		ScheduledUnit toSchedule = locationConstraint == null ?
				new ScheduledUnit(this, slotSharingGroupId) :
				new ScheduledUnit(this, slotSharingGroupId, locationConstraint);

		// try to extract previous allocation ids, if applicable, so that we can reschedule to the same slot
		ExecutionVertex executionVertex = getVertex();
		AllocationID lastAllocation = executionVertex.getLatestPriorAllocation();

		Collection<AllocationID> previousAllocationIDs =
			lastAllocation != null ? Collections.singletonList(lastAllocation) : Collections.emptyList();

		// calculate the preferred locations
		final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture =
			calculatePreferredLocations(locationPreferenceConstraint);

		final SlotRequestId slotRequestId = new SlotRequestId();

		final CompletableFuture<LogicalSlot> logicalSlotFuture =
			preferredLocationsFuture.thenCompose(
				(Collection<TaskManagerLocation> preferredLocations) ->
					slotProviderStrategy.allocateSlot(
						slotRequestId,
						toSchedule,
						new SlotProfile(
							vertex.getResourceProfile(),
							preferredLocations,
							previousAllocationIDs,
							allPreviousExecutionGraphAllocationIds)));

		// register call back to cancel slot request in case that the execution gets canceled
		releaseFuture.whenComplete(
			(Object ignored, Throwable throwable) -> {
				if (logicalSlotFuture.cancel(false)) {
					slotProviderStrategy.cancelSlotRequest(
						slotRequestId,
						slotSharingGroupId,
						new FlinkException("Execution " + this + " was released."));
				}
			});

		// This forces calls to the slot pool back into the main thread, for normal and exceptional completion
		return logicalSlotFuture.handle(
			(LogicalSlot logicalSlot, Throwable failure) -> {

				if (failure != null) {
					throw new CompletionException(failure);
				}

				if (tryAssignResource(logicalSlot)) {
					return logicalSlot;
				} else {
					// release the slot
					logicalSlot.releaseSlot(new FlinkException("Could not assign logical slot to execution " + this + '.'));
					throw new CompletionException(
						new FlinkException(
							"Could not assign slot " + logicalSlot + " to execution " + this + " because it has already been assigned "));
				}
			});
	} else {
		// call race, already deployed, or already done
		throw new IllegalExecutionStateException(this, CREATED, state);
	}
}