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

private void validateConstraints(ExecutionGraph eg) {

		ExecutionJobVertex[] tasks = eg.getAllVertices().values().toArray(new ExecutionJobVertex[2]);

		for (int i = 0; i < NUM_TASKS; i++) {
			CoLocationConstraint constr1 = tasks[0].getTaskVertices()[i].getLocationConstraint();
			CoLocationConstraint constr2 = tasks[1].getTaskVertices()[i].getLocationConstraint();
			assertThat(constr1.isAssigned(), is(true));
			assertThat(constr1.getLocation(), equalTo(constr2.getLocation()));
		}
	}
 

@Test(expected = IllegalStateException.class)
public void testCoLocationConstraintThrowsException() {
	final ExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();

	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Collections.singletonList(
		new ExecutionVertexSchedulingRequirements.Builder()
			.withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 0))
			.withCoLocationConstraint(coLocationConstraint)
			.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
}
 

@Test
public void testImmediateReleaseTwoLevel() {
	try {
		JobVertexID vid = new JobVertexID();
		JobVertex vertex = new JobVertex("vertex", vid);
		
		SlotSharingGroup sharingGroup = new SlotSharingGroup(vid);
		SlotSharingGroupAssignment assignment = sharingGroup.getTaskAssignment();

		CoLocationGroup coLocationGroup = new CoLocationGroup(vertex);
		CoLocationConstraint constraint = coLocationGroup.getLocationConstraint(0);
		
		Instance instance = SchedulerTestUtils.getRandomInstance(1);
		
		SharedSlot sharedSlot = instance.allocateSharedSlot(assignment);

		SimpleSlot sub = assignment.addSharedSlotAndAllocateSubSlot(sharedSlot, Locality.UNCONSTRAINED, constraint);
		
		assertNull(sub.getGroupID());
		assertEquals(constraint.getSharedSlot(), sub.getParent());
		
		sub.releaseSlot();

		assertTrue(sub.isReleased());
		assertTrue(sharedSlot.isReleased());

		assertEquals(1, instance.getNumberOfAvailableSlots());
		assertEquals(0, instance.getNumberOfAllocatedSlots());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

private void validateConstraints(ExecutionGraph eg) {

		ExecutionJobVertex[] tasks = eg.getAllVertices().values().toArray(new ExecutionJobVertex[2]);

		for (int i = 0; i < NUM_TASKS; i++) {
			CoLocationConstraint constr1 = tasks[0].getTaskVertices()[i].getLocationConstraint();
			CoLocationConstraint constr2 = tasks[1].getTaskVertices()[i].getLocationConstraint();
			assertThat(constr1.isAssigned(), is(true));
			assertThat(constr1.getLocation(), equalTo(constr2.getLocation()));
		}
	}
 

private ExecutionVertexSchedulingRequirements(
		ExecutionVertexID executionVertexId,
		@Nullable AllocationID previousAllocationId,
		ResourceProfile resourceProfile,
		@Nullable SlotSharingGroupId slotSharingGroupId,
		@Nullable CoLocationConstraint coLocationConstraint,
		Collection<TaskManagerLocation> preferredLocations) {
	this.executionVertexId = checkNotNull(executionVertexId);
	this.previousAllocationId = previousAllocationId;
	this.resourceProfile = checkNotNull(resourceProfile);
	this.slotSharingGroupId = slotSharingGroupId;
	this.coLocationConstraint = coLocationConstraint;
	this.preferredLocations = checkNotNull(preferredLocations);
}
 

private ExecutionVertexSchedulingRequirements(
		ExecutionVertexID executionVertexId,
		@Nullable AllocationID previousAllocationId,
		ResourceProfile taskResourceProfile,
		ResourceProfile physicalSlotResourceProfile,
		@Nullable SlotSharingGroupId slotSharingGroupId,
		@Nullable CoLocationConstraint coLocationConstraint) {
	this.executionVertexId = checkNotNull(executionVertexId);
	this.previousAllocationId = previousAllocationId;
	this.taskResourceProfile = checkNotNull(taskResourceProfile);
	this.physicalSlotResourceProfile = checkNotNull(physicalSlotResourceProfile);
	this.slotSharingGroupId = slotSharingGroupId;
	this.coLocationConstraint = coLocationConstraint;
}
 

private void validateConstraints(ExecutionGraph eg) {

		ExecutionJobVertex[] tasks = eg.getAllVertices().values().toArray(new ExecutionJobVertex[2]);

		for (int i = 0; i < NUM_TASKS; i++) {
			CoLocationConstraint constr1 = tasks[0].getTaskVertices()[i].getLocationConstraint();
			CoLocationConstraint constr2 = tasks[1].getTaskVertices()[i].getLocationConstraint();
			assertThat(constr1.isAssigned(), is(true));
			assertThat(constr1.getLocation(), equalTo(constr2.getLocation()));
		}
	}
 

/**
 * Tests that validate the parameters when calling allocateSlot in SlotProvider.
 */
@Test
public void testAllocateSlotsParameters() {
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
	final AllocationID allocationId = new AllocationID();
	final SlotSharingGroupId sharingGroupId = new SlotSharingGroupId();
	final ResourceProfile resourceProfile = new ResourceProfile(0.5, 250);
	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final Collection<TaskManagerLocation> taskManagerLocations = Collections.singleton(new LocalTaskManagerLocation());

	final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();

	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
			new ExecutionVertexSchedulingRequirements.Builder()
					.withExecutionVertexId(executionVertexId)
					.withPreviousAllocationId(allocationId)
					.withSlotSharingGroupId(sharingGroupId)
					.withPreferredLocations(taskManagerLocations)
					.withResourceProfile(resourceProfile)
					.withCoLocationConstraint(coLocationConstraint)
					.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
	assertThat(slotProvider.getSlotAllocationRequests(), hasSize(1));

	ScheduledUnit expectedTask = slotProvider.getSlotAllocationRequests().get(0).f1;
	SlotProfile expectedSlotProfile = slotProvider.getSlotAllocationRequests().get(0).f2;

	assertEquals(sharingGroupId, expectedTask.getSlotSharingGroupId());
	assertEquals(coLocationConstraint, expectedTask.getCoLocationConstraint());
	assertThat(expectedSlotProfile.getPreferredAllocations(), contains(allocationId));
	assertThat(expectedSlotProfile.getPreviousExecutionGraphAllocations(), contains(allocationId));
	assertEquals(resourceProfile, expectedSlotProfile.getResourceProfile());
	assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocations.toArray()));
}
 

@Test
public void testCoLocatedSlotRequestsFailBeforeResolved() throws ExecutionException, InterruptedException {
	final ResourceProfile rp1 = new ResourceProfile(1.0, 100);
	final ResourceProfile rp2 = new ResourceProfile(2.0, 200);
	final ResourceProfile rp3 = new ResourceProfile(5.0, 500);

	final ResourceProfile allocatedSlotRp = new ResourceProfile(3.0, 300);

	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();
	JobVertexID jobVertexId3 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture1 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId1,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp1),
			TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId2,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp2),
			TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId3,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp3),
			TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			Collections.singletonList(new SlotOffer(
					allocationId1,
					0,
					allocatedSlotRp)));

	assertFalse(slotOfferFuture1.isEmpty());

	for (CompletableFuture<LogicalSlot> logicalSlotFuture : Arrays.asList(logicalSlotFuture1, logicalSlotFuture2, logicalSlotFuture3)) {
		assertTrue(logicalSlotFuture.isDone() && logicalSlotFuture.isCompletedExceptionally());
		logicalSlotFuture.whenComplete((LogicalSlot ignored, Throwable throwable) -> {
			assertTrue(throwable instanceof SharedSlotOversubscribedException);
			assertTrue(((SharedSlotOversubscribedException) throwable).canRetry());
		});
	}
}
 

/**
 * Tests that validate the parameters when calling allocateSlot in SlotProvider.
 */
@Test
public void testAllocateSlotsParameters() {
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
	final AllocationID allocationId = new AllocationID();
	final SlotSharingGroupId sharingGroupId = new SlotSharingGroupId();
	final ResourceProfile taskResourceProfile = ResourceProfile.fromResources(0.5, 250);
	final ResourceProfile physicalSlotResourceProfile = ResourceProfile.fromResources(1.0, 300);
	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final TestingStateLocationRetriever stateLocationRetriever = new TestingStateLocationRetriever();
	stateLocationRetriever.setStateLocation(executionVertexId, taskManagerLocation);

	final ExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator(
		stateLocationRetriever,
		new TestingInputsLocationsRetriever.Builder().build());

	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
		new ExecutionVertexSchedulingRequirements.Builder()
			.withExecutionVertexId(executionVertexId)
			.withPreviousAllocationId(allocationId)
			.withSlotSharingGroupId(sharingGroupId)
			.withPhysicalSlotResourceProfile(physicalSlotResourceProfile)
			.withTaskResourceProfile(taskResourceProfile)
			.withCoLocationConstraint(coLocationConstraint)
			.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
	assertThat(slotProvider.getSlotAllocationRequests(), hasSize(1));

	ScheduledUnit expectedTask = slotProvider.getSlotAllocationRequests().get(0).f1;
	SlotProfile expectedSlotProfile = slotProvider.getSlotAllocationRequests().get(0).f2;

	assertEquals(sharingGroupId, expectedTask.getSlotSharingGroupId());
	assertEquals(coLocationConstraint, expectedTask.getCoLocationConstraint());
	assertThat(expectedSlotProfile.getPreferredAllocations(), contains(allocationId));
	assertThat(expectedSlotProfile.getPreviousExecutionGraphAllocations(), contains(allocationId));
	assertEquals(taskResourceProfile, expectedSlotProfile.getTaskResourceProfile());
	assertEquals(physicalSlotResourceProfile, expectedSlotProfile.getPhysicalSlotResourceProfile());
	assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocation));
}
 

/**
 * Tests the scheduling of two tasks with a parallelism of 2 and a co-location constraint.
 */
@Test
public void testSimpleCoLocatedSlotScheduling() throws ExecutionException, InterruptedException {
	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);
	CoLocationConstraint coLocationConstraint2 = group.getLocationConstraint(1);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture11 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint1),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture22 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint2),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture12 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint1),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture21 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint2),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();
	final AllocationID allocationId2 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId1,
			0,
			ResourceProfile.ANY)));

	Collection<SlotOffer> slotOfferFuture2 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId2,
			0,
			ResourceProfile.ANY)));

	assertFalse(slotOfferFuture1.isEmpty());
	assertFalse(slotOfferFuture2.isEmpty());

	LogicalSlot logicalSlot11 = logicalSlotFuture11.get();
	LogicalSlot logicalSlot12 = logicalSlotFuture12.get();
	LogicalSlot logicalSlot21 = logicalSlotFuture21.get();
	LogicalSlot logicalSlot22 = logicalSlotFuture22.get();

	assertEquals(logicalSlot11.getAllocationId(), logicalSlot12.getAllocationId());
	assertEquals(logicalSlot21.getAllocationId(), logicalSlot22.getAllocationId());
	assertNotEquals(logicalSlot11.getAllocationId(), logicalSlot21.getAllocationId());
}
 

/**
 * 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) -> {
					LOG.info("Allocating slot with SlotRequestID {} for the execution attempt {}.", slotRequestId, attemptId);
					return slotProviderStrategy.allocateSlot(
						slotRequestId,
						toSchedule,
						SlotProfile.priorAllocation(
							vertex.getResourceProfile(),
							getPhysicalSlotResourceProfile(vertex),
							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);
	}
}
 

public CoLocationConstraint getLocationConstraint() {
	return locationConstraint;
}
 

public Builder withCoLocationConstraint(final CoLocationConstraint coLocationConstraint) {
	this.coLocationConstraint = coLocationConstraint;
	return this;
}
 

@Nullable
public CoLocationConstraint getCoLocationConstraint() {
	return coLocationConstraint;
}
 

/**
 * Allocates a co-located {@link SlotSharingManager.MultiTaskSlot} for the given {@link CoLocationConstraint}.
 *
 * <p>The returned {@link SlotSharingManager.MultiTaskSlot} can be uncompleted.
 *
 * @param coLocationConstraint for which to allocate a {@link SlotSharingManager.MultiTaskSlot}
 * @param multiTaskSlotManager responsible for the slot sharing group for which to allocate the slot
 * @param slotProfile specifying the requirements for the requested slot
 * @param allocationTimeout timeout before the slot allocation times out
 * @return A {@link SlotAndLocality} which contains the allocated{@link SlotSharingManager.MultiTaskSlot}
 * 		and its locality wrt the given location preferences
 */
private SlotSharingManager.MultiTaskSlotLocality allocateCoLocatedMultiTaskSlot(
	CoLocationConstraint coLocationConstraint,
	SlotSharingManager multiTaskSlotManager,
	SlotProfile slotProfile,
	@Nullable Time allocationTimeout) throws NoResourceAvailableException {
	final SlotRequestId coLocationSlotRequestId = coLocationConstraint.getSlotRequestId();

	if (coLocationSlotRequestId != null) {
		// we have a slot assigned --> try to retrieve it
		final SlotSharingManager.TaskSlot taskSlot = multiTaskSlotManager.getTaskSlot(coLocationSlotRequestId);

		if (taskSlot != null) {
			Preconditions.checkState(taskSlot instanceof SlotSharingManager.MultiTaskSlot);

			SlotSharingManager.MultiTaskSlot multiTaskSlot = (SlotSharingManager.MultiTaskSlot) taskSlot;

			if (multiTaskSlot.mayHaveEnoughResourcesToFulfill(slotProfile.getTaskResourceProfile())) {
				return SlotSharingManager.MultiTaskSlotLocality.of(multiTaskSlot, Locality.LOCAL);
			}

			throw new NoResourceAvailableException("Not enough resources in the slot for all co-located tasks.");
		} else {
			// the slot may have been cancelled in the mean time
			coLocationConstraint.setSlotRequestId(null);
		}
	}

	if (coLocationConstraint.isAssigned()) {
		// refine the preferred locations of the slot profile
		slotProfile = SlotProfile.priorAllocation(
			slotProfile.getTaskResourceProfile(),
			slotProfile.getPhysicalSlotResourceProfile(),
			Collections.singleton(coLocationConstraint.getLocation()),
			slotProfile.getPreferredAllocations(),
			slotProfile.getPreviousExecutionGraphAllocations());
	}

	// get a new multi task slot
	SlotSharingManager.MultiTaskSlotLocality multiTaskSlotLocality = allocateMultiTaskSlot(
		coLocationConstraint.getGroupId(),
		multiTaskSlotManager,
		slotProfile,
		allocationTimeout);

	// check whether we fulfill the co-location constraint
	if (coLocationConstraint.isAssigned() && multiTaskSlotLocality.getLocality() != Locality.LOCAL) {
		multiTaskSlotLocality.getMultiTaskSlot().release(
			new FlinkException("Multi task slot is not local and, thus, does not fulfill the co-location constraint."));

		throw new NoResourceAvailableException("Could not allocate a local multi task slot for the " +
			"co location constraint " + coLocationConstraint + '.');
	}

	final SlotRequestId slotRequestId = new SlotRequestId();
	final SlotSharingManager.MultiTaskSlot coLocationSlot =
		multiTaskSlotLocality.getMultiTaskSlot().allocateMultiTaskSlot(
			slotRequestId,
			coLocationConstraint.getGroupId());

	// mark the requested slot as co-located slot for other co-located tasks
	coLocationConstraint.setSlotRequestId(slotRequestId);

	// lock the co-location constraint once we have obtained the allocated slot
	coLocationSlot.getSlotContextFuture().whenComplete(
		(SlotContext slotContext, Throwable throwable) -> {
			if (throwable == null) {
				// check whether we are still assigned to the co-location constraint
				if (Objects.equals(coLocationConstraint.getSlotRequestId(), slotRequestId)) {
					coLocationConstraint.lockLocation(slotContext.getTaskManagerLocation());
				} else {
					log.debug("Failed to lock colocation constraint {} because assigned slot " +
							"request {} differs from fulfilled slot request {}.",
						coLocationConstraint.getGroupId(),
						coLocationConstraint.getSlotRequestId(),
						slotRequestId);
				}
			} else {
				log.debug("Failed to lock colocation constraint {} because the slot " +
						"allocation for slot request {} failed.",
					coLocationConstraint.getGroupId(),
					coLocationConstraint.getSlotRequestId(),
					throwable);
			}
		});

	return SlotSharingManager.MultiTaskSlotLocality.of(coLocationSlot, multiTaskSlotLocality.getLocality());
}
 

@Test
public void testCoLocatedSlotRequestsFailAfterResolved() throws ExecutionException, InterruptedException {
	final ResourceProfile rp1 = new ResourceProfile(1.0, 100);
	final ResourceProfile rp2 = new ResourceProfile(2.0, 200);
	final ResourceProfile rp3 = new ResourceProfile(5.0, 500);

	final ResourceProfile allocatedSlotRp = new ResourceProfile(3.0, 300);

	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();
	JobVertexID jobVertexId3 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture1 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId1,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp1),
			TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId2,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp2),
			TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			Collections.singletonList(new SlotOffer(
					allocationId1,
					0,
					allocatedSlotRp)));

	assertFalse(slotOfferFuture1.isEmpty());

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId3,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp3),
			TestingUtils.infiniteTime());

	LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(allocationId1, logicalSlot1.getAllocationId());
	assertEquals(allocationId1, logicalSlot2.getAllocationId());

	assertTrue(logicalSlotFuture3.isDone() && logicalSlotFuture3.isCompletedExceptionally());
	logicalSlotFuture3.whenComplete((LogicalSlot ignored, Throwable throwable) -> {
		assertTrue(throwable instanceof SharedSlotOversubscribedException);
		assertTrue(((SharedSlotOversubscribedException) throwable).canRetry());
	});
}
 

public SimpleSlot addSharedSlotAndAllocateSubSlot(
		SharedSlot sharedSlot, Locality locality, CoLocationConstraint constraint)
{
	return addSharedSlotAndAllocateSubSlot(sharedSlot, locality, null, constraint);
}
 

/**
 * Tests the scheduling of two tasks with a parallelism of 2 and a co-location constraint.
 */
@Test
public void testSimpleCoLocatedSlotScheduling() throws ExecutionException, InterruptedException {
	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);
	CoLocationConstraint coLocationConstraint2 = group.getLocationConstraint(1);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture11 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint1),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture22 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint2),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture12 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint1),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture21 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint2),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();
	final AllocationID allocationId2 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId1,
			0,
			ResourceProfile.UNKNOWN)));

	Collection<SlotOffer> slotOfferFuture2 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId2,
			0,
			ResourceProfile.UNKNOWN)));

	assertFalse(slotOfferFuture1.isEmpty());
	assertFalse(slotOfferFuture2.isEmpty());

	LogicalSlot logicalSlot11 = logicalSlotFuture11.get();
	LogicalSlot logicalSlot12 = logicalSlotFuture12.get();
	LogicalSlot logicalSlot21 = logicalSlotFuture21.get();
	LogicalSlot logicalSlot22 = logicalSlotFuture22.get();

	assertEquals(logicalSlot11.getAllocationId(), logicalSlot12.getAllocationId());
	assertEquals(logicalSlot21.getAllocationId(), logicalSlot22.getAllocationId());
	assertNotEquals(logicalSlot11.getAllocationId(), logicalSlot21.getAllocationId());
}
 

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

public CoLocationConstraint getLocationConstraint() {
	return locationConstraint;
}
 

public Builder withCoLocationConstraint(final CoLocationConstraint coLocationConstraint) {
	this.coLocationConstraint = coLocationConstraint;
	return this;
}
 

@Nullable
public CoLocationConstraint getCoLocationConstraint() {
	return coLocationConstraint;
}
 

/**
 * We allocate and the structure below and release it from the root.
 *
 * <pre>
 *     Shared(0)(root)
 *        |
 *        +-- Simple(2)(sink)
 *        |
 *        +-- Shared(1)(co-location-group)
 *        |      |
 *        |      +-- Simple(0)(tail)
 *        |      +-- Simple(1)(head)
 *        |
 *        +-- Simple(0)(source)
 * </pre>
 */
@Test
public void testReleaseTwoLevelsFromRoot() {
	try {
		JobVertexID sourceId = new JobVertexID();
		JobVertexID headId = new JobVertexID();
		JobVertexID tailId = new JobVertexID();
		JobVertexID sinkId = new JobVertexID();

		JobVertex headVertex = new JobVertex("head", headId);
		JobVertex tailVertex = new JobVertex("tail", tailId);

		SlotSharingGroup sharingGroup = new SlotSharingGroup(sourceId, headId, tailId, sinkId);
		SlotSharingGroupAssignment assignment = sharingGroup.getTaskAssignment();
		assertEquals(0, assignment.getNumberOfSlots());

		CoLocationGroup coLocationGroup = new CoLocationGroup(headVertex, tailVertex);
		CoLocationConstraint constraint = coLocationGroup.getLocationConstraint(0);
		assertFalse(constraint.isAssigned());

		Instance instance = SchedulerTestUtils.getRandomInstance(1);

		// allocate a shared slot
		SharedSlot sharedSlot = instance.allocateSharedSlot(assignment);

		// get the first simple slot
		SimpleSlot sourceSlot = assignment.addSharedSlotAndAllocateSubSlot(sharedSlot, Locality.LOCAL, sourceId);
		
		SimpleSlot headSlot = assignment.getSlotForTask(constraint, NO_LOCATION);
		constraint.lockLocation();
		SimpleSlot tailSlot = assignment.getSlotForTask(constraint, NO_LOCATION);
		
		SimpleSlot sinkSlot = assignment.getSlotForTask(sinkId, NO_LOCATION);
		
		assertEquals(4, sharedSlot.getNumberLeaves());

		// release all
		sourceSlot.releaseSlot();
		headSlot.releaseSlot();
		tailSlot.releaseSlot();
		sinkSlot.releaseSlot();

		assertTrue(sharedSlot.isReleased());
		assertTrue(sourceSlot.isReleased());
		assertTrue(headSlot.isReleased());
		assertTrue(tailSlot.isReleased());
		assertTrue(sinkSlot.isReleased());
		assertTrue(constraint.getSharedSlot().isReleased());

		assertTrue(constraint.isAssigned());
		assertFalse(constraint.isAssignedAndAlive());

		assertEquals(1, instance.getNumberOfAvailableSlots());
		assertEquals(0, instance.getNumberOfAllocatedSlots());
		
		assertEquals(0, assignment.getNumberOfSlots());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests the scheduling of two tasks with a parallelism of 2 and a co-location constraint.
 */
@Test
public void testSimpleCoLocatedSlotScheduling() throws ExecutionException, InterruptedException {
	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);
	CoLocationConstraint coLocationConstraint2 = group.getLocationConstraint(1);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture11 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint1),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture22 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint2),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture12 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint1),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture21 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint2),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();
	final AllocationID allocationId2 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId1,
			0,
			ResourceProfile.UNKNOWN)));

	Collection<SlotOffer> slotOfferFuture2 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId2,
			0,
			ResourceProfile.UNKNOWN)));

	assertFalse(slotOfferFuture1.isEmpty());
	assertFalse(slotOfferFuture2.isEmpty());

	LogicalSlot logicalSlot11 = logicalSlotFuture11.get();
	LogicalSlot logicalSlot12 = logicalSlotFuture12.get();
	LogicalSlot logicalSlot21 = logicalSlotFuture21.get();
	LogicalSlot logicalSlot22 = logicalSlotFuture22.get();

	assertEquals(logicalSlot11.getAllocationId(), logicalSlot12.getAllocationId());
	assertEquals(logicalSlot21.getAllocationId(), logicalSlot22.getAllocationId());
	assertNotEquals(logicalSlot11.getAllocationId(), logicalSlot21.getAllocationId());
}
 

public CoLocationConstraint getLocationConstraint() {
	return locationConstraint;
}
 

private SimpleSlot addSharedSlotAndAllocateSubSlot(
		SharedSlot sharedSlot, Locality locality, JobVertexID groupId, CoLocationConstraint constraint) {

	// sanity checks
	if (!sharedSlot.isRootAndEmpty()) {
		throw new IllegalArgumentException("The given slot is not an empty root slot.");
	}

	final ResourceID location = sharedSlot.getTaskManagerID();

	synchronized (lock) {
		// early out in case that the slot died (instance disappeared)
		if (!sharedSlot.isAlive()) {
			return null;
		}
		
		// add to the total bookkeeping
		if (!allSlots.add(sharedSlot)) {
			throw new IllegalArgumentException("Slot was already contained in the assignment group");
		}
		
		SimpleSlot subSlot;
		AbstractID groupIdForMap;
				
		if (constraint == null) {
			// allocate us a sub slot to return
			subSlot = sharedSlot.allocateSubSlot(groupId);
			groupIdForMap = groupId;
		}
		else {
			// sanity check
			if (constraint.isAssignedAndAlive()) {
				throw new IllegalStateException(
						"Trying to add a shared slot to a co-location constraint that has a life slot.");
			}
			
			// we need a co-location slot --> a SimpleSlot nested in a SharedSlot to
			//                                host other co-located tasks
			SharedSlot constraintGroupSlot = sharedSlot.allocateSharedSlot(constraint.getGroupId());
			groupIdForMap = constraint.getGroupId();
			
			if (constraintGroupSlot != null) {
				// the sub-slots in the co-location constraint slot have no own group IDs
				subSlot = constraintGroupSlot.allocateSubSlot(null);
				if (subSlot != null) {
					// all went well, we can give the constraint its slot
					constraint.setSharedSlot(constraintGroupSlot);
					
					// NOTE: Do not lock the location constraint, because we don't yet know whether we will
					// take the slot here
				}
				else {
					// if we could not create a sub slot, release the co-location slot
					// note that this does implicitly release the slot we have just added
					// as well, because we release its last child slot. That is expected
					// and desired.
					constraintGroupSlot.releaseSlot(new FlinkException("Could not create a sub slot in this shared slot."));
				}
			}
			else {
				// this should not happen, as we are under the lock that also
				// guards slot disposals. Keep the check to be on the safe side
				subSlot = null;
			}
		}
		
		if (subSlot != null) {
			// preserve the locality information
			subSlot.setLocality(locality);
			
			// let the other groups know that this slot exists and that they
			// can place a task into this slot.
			boolean entryForNewJidExists = false;
			
			for (Map.Entry<AbstractID, Map<ResourceID, List<SharedSlot>>> entry : availableSlotsPerJid.entrySet()) {
				// there is already an entry for this groupID
				if (entry.getKey().equals(groupIdForMap)) {
					entryForNewJidExists = true;
					continue;
				}

				Map<ResourceID, List<SharedSlot>> available = entry.getValue();
				putIntoMultiMap(available, location, sharedSlot);
			}

			// make sure an empty entry exists for this group, if no other entry exists
			if (!entryForNewJidExists) {
				availableSlotsPerJid.put(groupIdForMap, new LinkedHashMap<ResourceID, List<SharedSlot>>());
			}

			return subSlot;
		}
		else {
			// if sharedSlot is releases, abort.
			// This should be a rare case, since this method is called with a fresh slot.
			return null;
		}
	}
	// end synchronized (lock)
}