org.apache.flink.runtime.clusterframework.types.SlotProfile Java Examples

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		boolean allowQueued,
		Time allocationTimeout) {
	final SlotContext slot;

	synchronized (lock) {
		if (slots.isEmpty()) {
			slot = null;
		} else {
			slot = slots.removeFirst();
		}
		if (slot != null) {
			SimpleSlot result = new SimpleSlot(slot, this, 0);
			allocatedSlots.put(slotRequestId, slot);
			return CompletableFuture.completedFuture(result);
		}
		else {
			return FutureUtils.completedExceptionally(new NoResourceAvailableException());
		}
	}
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		Time allocationTimeout) {
	final JobVertexID vertexId = task.getJobVertexId();
	final int subtask = task.getSubtaskIndex();

	CompletableFuture<LogicalSlot>[] forTask = slotFutures.get(vertexId);
	if (forTask != null) {
		CompletableFuture<LogicalSlot> future = forTask[subtask];
		if (future != null) {
			slotFutureRequested.get(vertexId)[subtask].complete(true);
			slotRequests.add(slotRequestId);

			return future;
		}
	}

	throw new IllegalArgumentException("No registered slot future for task " + vertexId + " (" + subtask + ')');
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		boolean allowQueued,
		Time allocationTimeout) {
	JobVertexID vertexId = task.getTaskToExecute().getVertex().getJobvertexId();
	int subtask = task.getTaskToExecute().getParallelSubtaskIndex();

	CompletableFuture<LogicalSlot>[] forTask = slotFutures.get(vertexId);
	if (forTask != null) {
		CompletableFuture<LogicalSlot> future = forTask[subtask];
		if (future != null) {
			slotFutureRequested.get(vertexId)[subtask].complete(true);
			slotRequests.add(slotRequestId);

			return future;
		}
	}

	throw new IllegalArgumentException("No registered slot future for task " + vertexId + " (" + subtask + ')');
}
 

@Nonnull
private CompletableFuture<LogicalSlot> allocateSlotInternal(
	SlotRequestId slotRequestId,
	ScheduledUnit scheduledUnit,
	SlotProfile slotProfile,
	@Nullable Time allocationTimeout) {
	log.debug("Received slot request [{}] for task: {}", slotRequestId, scheduledUnit.getJobVertexId());

	componentMainThreadExecutor.assertRunningInMainThread();

	final CompletableFuture<LogicalSlot> allocationResultFuture = new CompletableFuture<>();
	internalAllocateSlot(
			allocationResultFuture,
			slotRequestId,
			scheduledUnit,
			slotProfile,
			allocationTimeout);
	return allocationResultFuture;
}
 

private CompletableFuture<LogicalSlot> allocateSlot(
		final ExecutionVertexSchedulingRequirements schedulingRequirements,
		final SlotRequestId slotRequestId,
		final Set<AllocationID> allPreviousAllocationIds) {

	final ExecutionVertexID executionVertexId = schedulingRequirements.getExecutionVertexId();

	LOG.debug("Allocate slot with id {} for execution {}", slotRequestId, executionVertexId);

	final CompletableFuture<SlotProfile> slotProfileFuture = getSlotProfileFuture(
		schedulingRequirements,
		schedulingRequirements.getPhysicalSlotResourceProfile(),
		Collections.emptySet(),
		allPreviousAllocationIds);

	return slotProfileFuture.thenCompose(
		slotProfile -> slotProviderStrategy.allocateSlot(
			slotRequestId,
			new ScheduledUnit(
				executionVertexId,
				schedulingRequirements.getSlotSharingGroupId(),
				schedulingRequirements.getCoLocationConstraint()),
			slotProfile));
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		boolean allowQueued,
		Time allocationTimeout) {
	JobVertexID vertexId = task.getTaskToExecute().getVertex().getJobvertexId();
	int subtask = task.getTaskToExecute().getParallelSubtaskIndex();

	CompletableFuture<LogicalSlot>[] forTask = slotFutures.get(vertexId);
	if (forTask != null) {
		CompletableFuture<LogicalSlot> future = forTask[subtask];
		if (future != null) {
			slotFutureRequested.get(vertexId)[subtask].complete(true);
			slotRequests.add(slotRequestId);

			return future;
		}
	}

	throw new IllegalArgumentException("No registered slot future for task " + vertexId + " (" + subtask + ')');
}
 

private Optional<SlotAndLocality> tryAllocateFromAvailable(
	@Nonnull SlotRequestId slotRequestId,
	@Nonnull SlotProfile slotProfile) {

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfoList =
			slotPool.getAvailableSlotsInformation()
					.stream()
					.map(SlotSelectionStrategy.SlotInfoAndResources::new)
					.collect(Collectors.toList());

	Optional<SlotSelectionStrategy.SlotInfoAndLocality> selectedAvailableSlot =
		slotSelectionStrategy.selectBestSlotForProfile(slotInfoList, slotProfile);

	return selectedAvailableSlot.flatMap(slotInfoAndLocality -> {
		Optional<PhysicalSlot> optionalAllocatedSlot = slotPool.allocateAvailableSlot(
			slotRequestId,
			slotInfoAndLocality.getSlotInfo().getAllocationId());

		return optionalAllocatedSlot.map(
			allocatedSlot -> new SlotAndLocality(allocatedSlot, slotInfoAndLocality.getLocality()));
	});
}
 

private void internalAllocateSlot(
		CompletableFuture<LogicalSlot> allocationResultFuture,
		SlotRequestId slotRequestId,
		ScheduledUnit scheduledUnit,
		SlotProfile slotProfile,
		Time allocationTimeout) {
	CompletableFuture<LogicalSlot> allocationFuture = scheduledUnit.getSlotSharingGroupId() == null ?
		allocateSingleSlot(slotRequestId, slotProfile, allocationTimeout) :
		allocateSharedSlot(slotRequestId, scheduledUnit, slotProfile, allocationTimeout);

	allocationFuture.whenComplete((LogicalSlot slot, Throwable failure) -> {
		if (failure != null) {
			cancelSlotRequest(
				slotRequestId,
				scheduledUnit.getSlotSharingGroupId(),
				failure);
			allocationResultFuture.completeExceptionally(failure);
		} else {
			allocationResultFuture.complete(slot);
		}
	});
}
 

@Test
public void testNewPhysicalSlotAllocation() {
	final ResourceProfile taskResourceProfile = ResourceProfile.fromResources(0.5, 250);
	final ResourceProfile physicalSlotResourceProfile = ResourceProfile.fromResources(1.0, 300);

	testingSlotProvider.allocateSlot(
		new SlotRequestId(),
		new ScheduledUnit(new JobVertexID(), null, null),
		SlotProfile.priorAllocation(
			taskResourceProfile,
			physicalSlotResourceProfile,
			Collections.emptyList(),
			Collections.emptyList(),
			Collections.emptySet()),
		TestingUtils.infiniteTime());

	assertEquals(physicalSlotResourceProfile, testingSlotProvider.getSlotPool().getLastRequestedSlotResourceProfile());
}
 

@Override
public Optional<SlotInfoAndLocality> selectBestSlotForProfile(
	@Nonnull Collection<SlotInfoAndResources> availableSlots,
	@Nonnull SlotProfile slotProfile) {

	Collection<AllocationID> priorAllocations = slotProfile.getPreferredAllocations();

	// First, if there was a prior allocation try to schedule to the same/old slot
	if (!priorAllocations.isEmpty()) {
		for (SlotInfoAndResources availableSlot : availableSlots) {
			if (priorAllocations.contains(availableSlot.getSlotInfo().getAllocationId())) {
				return Optional.of(
					SlotInfoAndLocality.of(availableSlot.getSlotInfo(), Locality.LOCAL));
			}
		}
	}

	// Second, select based on location preference, excluding blacklisted allocations
	Set<AllocationID> blackListedAllocations = slotProfile.getPreviousExecutionGraphAllocations();
	Collection<SlotInfoAndResources> availableAndAllowedSlots = computeWithoutBlacklistedSlots(availableSlots, blackListedAllocations);
	return fallbackSlotSelectionStrategy.selectBestSlotForProfile(availableAndAllowedSlots, slotProfile);
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		boolean allowQueued,
		Time allocationTimeout) {
	final SlotContext slot;

	synchronized (lock) {
		if (slots.isEmpty()) {
			slot = null;
		} else {
			slot = slots.removeFirst();
		}
		if (slot != null) {
			TestingLogicalSlot result = new TestingLogicalSlotBuilder()
				.setTaskManagerLocation(slot.getTaskManagerLocation())
				.setTaskManagerGateway(slot.getTaskManagerGateway())
				.setSlotNumber(slot.getPhysicalSlotNumber())
				.setAllocationId(slot.getAllocationId())
				.setSlotRequestId(slotRequestId)
				.setSlotOwner(this)
				.createTestingLogicalSlot();
			allocatedSlots.put(slotRequestId, slot);
			return CompletableFuture.completedFuture(result);
		}
		else {
			return FutureUtils.completedExceptionally(new NoResourceAvailableException());
		}
	}
}
 

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

/**
 * Tests that returned slot futures are failed if the allocation request is failed.
 */
@Test
public void testFailingQueuedSharedSlotScheduling() throws Exception {
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture = slotProvider.allocateSlot(
		new ScheduledUnit(
			new JobVertexID(),
			new SlotSharingGroupId(),
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	final AllocationID allocationId = allocationIdFuture.get();

	// this should fail the returned logical slot future
	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.failAllocation(allocationId, new FlinkException("Testing Exception"));

	try {
		logicalSlotFuture.get();
		fail("The slot future should have failed.");
	} catch (ExecutionException ee) {
		assertTrue(ExceptionUtils.findThrowable(ee, FlinkException.class).isPresent());
	}
}
 

/**
 * Tests that we can correctly retrieve resolved slots.
 */
@Test
public void testGetResolvedSlot() {
	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot rootSlot = createRootSlot(new LocalTaskManagerLocation(), slotSharingManager);

	AbstractID groupId = new AbstractID();

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	Assert.assertEquals(1, slotInfos.size());

	SlotSelectionStrategy.SlotInfoAndResources slotInfoAndRemainingResource = slotInfos.iterator().next();
	SlotSharingManager.MultiTaskSlot resolvedMultiTaskSlot =
		slotSharingManager.getResolvedRootSlot(slotInfoAndRemainingResource.getSlotInfo());

	final LocationPreferenceSlotSelectionStrategy locationPreferenceSlotSelectionStrategy = LocationPreferenceSlotSelectionStrategy.createDefault();
	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality = locationPreferenceSlotSelectionStrategy
		.selectBestSlotForProfile(slotInfos, SlotProfile.noRequirements()).get();

	assertNotNull(resolvedMultiTaskSlot);
	assertEquals(Locality.UNCONSTRAINED, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot.getSlotRequestId(), resolvedMultiTaskSlot.getSlotRequestId());

	// occupy the resolved root slot
	resolvedMultiTaskSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId,
		Locality.UNCONSTRAINED);

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	assertTrue(slotInfos.isEmpty());
}
 

/**
 * Allocating slot with specific requirement.
 *
 * @param scheduledUnit The task to allocate the slot for
 * @param allowQueued Whether allow the task be queued if we do not have enough resource
 * @param slotProfile profile of the requested slot
 * @param allocationTimeout after which the allocation fails with a timeout exception
 * @return The future of the allocation
 */
default CompletableFuture<LogicalSlot> allocateSlot(
	ScheduledUnit scheduledUnit,
	boolean allowQueued,
	SlotProfile slotProfile,
	Time allocationTimeout) {
	return allocateSlot(
		new SlotRequestId(),
		scheduledUnit,
		slotProfile,
		allowQueued,
		allocationTimeout);
}
 

@Test
public void testCancelSlotAllocationWithoutResourceManager() throws Exception {
	final JobID jid = new JobID();

	try (TestingSlotPool pool = new TestingSlotPool(
		jid,
		SystemClock.getInstance(),
		TestingUtils.infiniteTime(),
		TestingUtils.infiniteTime())) {

		final CompletableFuture<SlotRequestId> timeoutFuture = new CompletableFuture<>();
		pool.setTimeoutPendingSlotRequestConsumer(timeoutFuture::complete);
		pool.start(JobMasterId.generate(), "foobar", testMainThreadExecutor.getMainThreadExecutor());
		Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.INSTANCE, pool);
		scheduler.start(testMainThreadExecutor.getMainThreadExecutor());

		SlotRequestId requestId = new SlotRequestId();
		CompletableFuture<LogicalSlot> future = testMainThreadExecutor.execute(() -> scheduler.allocateSlot(
			requestId,
			new ScheduledUnit(SchedulerTestUtils.getDummyTask()),
			SlotProfile.noLocality(DEFAULT_TESTING_PROFILE),
			true,
			fastTimeout));

		try {
			future.get();
			fail("We expected a TimeoutException.");
		} catch (ExecutionException e) {
			assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
		}

		// wait for the timeout of the pending slot request
		timeoutFuture.get();

		assertEquals(0L, pool.getNumberOfWaitingForResourceRequests());
	}
}
 

@Test
public void scheduleWithIntermediateRelease() throws Exception {
	JobVertexID jid1 = new JobVertexID();
	JobVertexID jid2 = new JobVertexID();
	JobVertexID jid3 = new JobVertexID();
	JobVertexID jid4 = new JobVertexID();

	testingSlotProvider.addTaskManager(1);
	testingSlotProvider.addTaskManager(1);

	assertEquals(2, testingSlotProvider.getNumberOfAvailableSlots());

	SlotSharingGroup sharingGroup = new SlotSharingGroup();
	CoLocationConstraint c1 = new CoLocationConstraint(new CoLocationGroup());

	LogicalSlot s1 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getExecution(jid1, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	LogicalSlot s2 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getExecution(jid2, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();

	LogicalSlot sSolo = testingSlotProvider.allocateSlot(new ScheduledUnit(getExecution(jid4, 0, 1, null)), SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();

	ResourceID taskManager = s1.getTaskManagerLocation().getResourceID();

	s1.releaseSlot();
	s2.releaseSlot();
	sSolo.releaseSlot();

	LogicalSlot sNew = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getExecution(jid3, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	assertEquals(taskManager, sNew.getTaskManagerLocation().getResourceID());

	assertEquals(2, testingSlotProvider.getNumberOfLocalizedAssignments());
	assertEquals(0, testingSlotProvider.getNumberOfNonLocalizedAssignments());
	assertEquals(2, testingSlotProvider.getNumberOfUnconstrainedAssignments());
}
 

/**
 * This case make sure when allocateSlot in ProviderAndOwner timeout,
 * it will automatically call cancelSlotAllocation as will inject future.whenComplete in ProviderAndOwner.
 */
@Test
public void testProviderAndOwnerSlotAllocationTimeout() throws Exception {
	final JobID jid = new JobID();

	try (TestingSlotPool pool = createTestingSlotPool(jid)) {

		final CompletableFuture<SlotRequestId> releaseSlotFuture = new CompletableFuture<>();

		pool.setReleaseSlotConsumer(releaseSlotFuture::complete);

		pool.start(JobMasterId.generate(), "foobar", testMainThreadExecutor.getMainThreadExecutor());
		ResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
		pool.connectToResourceManager(resourceManagerGateway);

		Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), pool);
		scheduler.start(testMainThreadExecutor.getMainThreadExecutor());

		// test the pending request is clear when timed out
		CompletableFuture<LogicalSlot> future = testMainThreadExecutor.execute(() -> scheduler.allocateSlot(
			new DummyScheduledUnit(),
			SlotProfile.noRequirements(),
			fastTimeout));
		try {
			future.get();
			fail("We expected a TimeoutException.");
		} catch (ExecutionException e) {
			assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
		}

		// wait for the cancel call on the SlotPoolImpl
		releaseSlotFuture.get();

		assertEquals(0L, pool.getNumberOfPendingRequests());
	}
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		boolean allowQueued,
		Time allocationTimeout) {

	try {
		final Object ret = scheduleTask(task, allowQueued, slotProfile.getPreferredLocations());

		if (ret instanceof SimpleSlot) {
			return CompletableFuture.completedFuture((SimpleSlot) ret);
		}
		else if (ret instanceof CompletableFuture) {
			@SuppressWarnings("unchecked")
			CompletableFuture<LogicalSlot> typed = (CompletableFuture<LogicalSlot>) ret;
			return FutureUtils.orTimeout(typed, allocationTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		}
		else {
			// this should never happen, simply guard this case with an exception
			throw new RuntimeException();
		}
	} catch (NoResourceAvailableException e) {
		return FutureUtils.completedExceptionally(e);
	}
}
 

@Test
public void scheduleWithIntermediateRelease() throws Exception {
	JobVertexID jid1 = new JobVertexID();
	JobVertexID jid2 = new JobVertexID();
	JobVertexID jid3 = new JobVertexID();
	JobVertexID jid4 = new JobVertexID();

	testingSlotProvider.addTaskManager(1);
	testingSlotProvider.addTaskManager(1);

	assertEquals(2, testingSlotProvider.getNumberOfAvailableSlots());

	SlotSharingGroup sharingGroup = new SlotSharingGroup();
	CoLocationConstraint c1 = new CoLocationConstraint(new CoLocationGroup());

	LogicalSlot s1 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertex(jid1, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	LogicalSlot s2 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertex(jid2, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();

	LogicalSlot sSolo = testingSlotProvider.allocateSlot(new ScheduledUnit(getTestVertex(jid4, 0, 1, null)), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();

	ResourceID taskManager = s1.getTaskManagerLocation().getResourceID();

	s1.releaseSlot();
	s2.releaseSlot();
	sSolo.releaseSlot();

	LogicalSlot sNew = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertex(jid3, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	assertEquals(taskManager, sNew.getTaskManagerLocation().getResourceID());

	assertEquals(2, testingSlotProvider.getNumberOfLocalizedAssignments());
	assertEquals(0, testingSlotProvider.getNumberOfNonLocalizedAssignments());
	assertEquals(2, testingSlotProvider.getNumberOfUnconstrainedAssignments());
}
 

@Test
public void scheduleWithReleaseNoResource() throws Exception {
	JobVertexID jid1 = new JobVertexID();
	JobVertexID jid2 = new JobVertexID();
	JobVertexID jid3 = new JobVertexID();

	testingSlotProvider.addTaskManager(1);
	testingSlotProvider.addTaskManager(1);

	assertEquals(2, testingSlotProvider.getNumberOfAvailableSlots());

	SlotSharingGroup sharingGroup = new SlotSharingGroup();
	CoLocationConstraint c1 = new CoLocationConstraint(new CoLocationGroup());

	LogicalSlot s1 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertex(jid1, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	s1.releaseSlot();

	testingSlotProvider.allocateSlot(new ScheduledUnit(getTestVertex(jid2, 0, 1, null)), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
	testingSlotProvider.allocateSlot(new ScheduledUnit(getTestVertex(jid2, 1, 2, null)), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();

	try {
		testingSlotProvider.allocateSlot(new ScheduledUnit(getTestVertex(jid3, 0, 1, sharingGroup), sharingGroup.getSlotSharingGroupId(), c1), false, SlotProfile.noRequirements(), TestingUtils.infiniteTime()).get();
		fail("Scheduled even though no resource was available.");
	} catch (ExecutionException e) {
		assertTrue(e.getCause() instanceof NoResourceAvailableException);
	}

	assertEquals(0, testingSlotProvider.getNumberOfLocalizedAssignments());
	assertEquals(0, testingSlotProvider.getNumberOfNonLocalizedAssignments());
	assertEquals(3, testingSlotProvider.getNumberOfUnconstrainedAssignments());
}
 

@Nonnull
private CompletableFuture<PhysicalSlot> requestNewAllocatedSlot(
		SlotRequestId slotRequestId,
		SlotProfile slotProfile,
		@Nullable Time allocationTimeout) {
	if (allocationTimeout == null) {
		return slotPool.requestNewAllocatedBatchSlot(slotRequestId, slotProfile.getPhysicalSlotResourceProfile());
	} else {
		return slotPool.requestNewAllocatedSlot(slotRequestId, slotProfile.getPhysicalSlotResourceProfile(), allocationTimeout);
	}
}
 

@Test
public void testSingleQueuedSharedSlotScheduling() throws Exception {
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture = slotProvider.allocateSlot(
		new ScheduledUnit(
			new JobVertexID(),
			slotSharingGroupId,
			null),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

	boolean booleanCompletableFuture = slotPool.offerSlot(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		new SlotOffer(
			allocationId,
			0,
			ResourceProfile.ANY));

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

	assertEquals(slotSharingGroupId, logicalSlot.getSlotSharingGroupId());
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
	SlotRequestId slotRequestId,
	ScheduledUnit task,
	SlotProfile slotProfile,
	boolean allowQueued,
	Time allocationTimeout) {
	return scheduler.allocateSlot(task, allowQueued, slotProfile, allocationTimeout).thenApply(
		(LogicalSlot logicalSlot) -> {
			switch (logicalSlot.getLocality()) {
				case LOCAL:
					numberOfLocalizedAssignments.incrementAndGet();
					break;
				case UNCONSTRAINED:
					numberOfUnconstrainedAssignments.incrementAndGet();
					break;
				case NON_LOCAL:
					numberOfNonLocalizedAssignments.incrementAndGet();
					break;
				case HOST_LOCAL:
					numberOfHostLocalizedAssignments.incrementAndGet();
					break;
				default:
					// ignore
			}

			return logicalSlot;
		});
}
 

/**
 * Tests that the location preferences are honoured when looking for a resolved slot.
 */
@Test
public void testGetResolvedSlotWithLocationPreferences() {
	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot rootSlot1 = createRootSlot(new LocalTaskManagerLocation(), slotSharingManager);

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	SlotSharingManager.MultiTaskSlot rootSlot2 = createRootSlot(taskManagerLocation, slotSharingManager);

	AbstractID groupId = new AbstractID();

	SlotProfile slotProfile = SlotProfile.preferredLocality(ResourceProfile.UNKNOWN, Collections.singleton(taskManagerLocation));

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	final LocationPreferenceSlotSelectionStrategy locationPreferenceSlotSelectionStrategy = LocationPreferenceSlotSelectionStrategy.createDefault();
	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		locationPreferenceSlotSelectionStrategy.selectBestSlotForProfile(slotInfos, slotProfile).get();
	SlotSharingManager.MultiTaskSlot resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());

	assertNotNull(resolvedRootSlot);
	assertEquals(Locality.LOCAL, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot2.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());

	// occupy the slot
	resolvedRootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId,
		slotInfoAndLocality.getLocality());

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	slotInfoAndLocality = locationPreferenceSlotSelectionStrategy.selectBestSlotForProfile(slotInfos, slotProfile).get();
	resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());
	assertNotNull(resolvedRootSlot);
	assertNotSame(Locality.LOCAL, (slotInfoAndLocality.getLocality()));
	assertEquals(rootSlot1.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit scheduledUnit,
		SlotProfile slotProfile,
		boolean allowQueuedScheduling,
		Time allocationTimeout) {
	return allocateSlotInternal(
		slotRequestId,
		scheduledUnit,
		slotProfile,
		allowQueuedScheduling,
		allocationTimeout);
}
 

@Override
public CompletableFuture<LogicalSlot> allocateBatchSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit scheduledUnit,
		SlotProfile slotProfile,
		boolean allowQueuedScheduling) {
	return allocateSlotInternal(
		slotRequestId,
		scheduledUnit,
		slotProfile,
		allowQueuedScheduling,
		null);
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		Time allocationTimeout) {
	final SlotContext slot;

	synchronized (lock) {
		if (slots.isEmpty()) {
			slot = null;
		} else {
			slot = slots.removeFirst();
		}
		if (slot != null) {
			TestingLogicalSlot result = new TestingLogicalSlotBuilder()
				.setTaskManagerLocation(slot.getTaskManagerLocation())
				.setTaskManagerGateway(slot.getTaskManagerGateway())
				.setSlotNumber(slot.getPhysicalSlotNumber())
				.setAllocationId(slot.getAllocationId())
				.setSlotRequestId(slotRequestId)
				.setSlotSharingGroupId(task.getSlotSharingGroupId())
				.setSlotOwner(this)
				.createTestingLogicalSlot();
			allocatedSlots.put(slotRequestId, slot);
			return CompletableFuture.completedFuture(result);
		} else {
			return FutureUtils.completedExceptionally(new NoResourceAvailableException());
		}
	}
}
 

private void internalAllocateSlot(
		CompletableFuture<LogicalSlot> allocationResultFuture,
		SlotRequestId slotRequestId,
		ScheduledUnit scheduledUnit,
		SlotProfile slotProfile,
		boolean allowQueuedScheduling,
		Time allocationTimeout) {
	CompletableFuture<LogicalSlot> allocationFuture = scheduledUnit.getSlotSharingGroupId() == null ?
		allocateSingleSlot(slotRequestId, slotProfile, allowQueuedScheduling, allocationTimeout) :
		allocateSharedSlot(slotRequestId, scheduledUnit, slotProfile, allowQueuedScheduling, allocationTimeout);

	allocationFuture.whenComplete((LogicalSlot slot, Throwable failure) -> {
		if (failure != null) {
			Optional<SharedSlotOversubscribedException> sharedSlotOverAllocatedException =
					ExceptionUtils.findThrowable(failure, SharedSlotOversubscribedException.class);
			if (sharedSlotOverAllocatedException.isPresent() &&
					sharedSlotOverAllocatedException.get().canRetry()) {

				// Retry the allocation
				internalAllocateSlot(
						allocationResultFuture,
						slotRequestId,
						scheduledUnit,
						slotProfile,
						allowQueuedScheduling,
						allocationTimeout);
			} else {
				cancelSlotRequest(
						slotRequestId,
						scheduledUnit.getSlotSharingGroupId(),
						failure);
				allocationResultFuture.completeExceptionally(failure);
			}
		} else {
			allocationResultFuture.complete(slot);
		}
	});
}
 

/**
 * Allocating slot with specific requirement.
 *
 * @param scheduledUnit The task to allocate the slot for
 * @param slotProfile profile of the requested slot
 * @param allocationTimeout after which the allocation fails with a timeout exception
 * @return The future of the allocation
 */
default CompletableFuture<LogicalSlot> allocateSlot(
	ScheduledUnit scheduledUnit,
	SlotProfile slotProfile,
	Time allocationTimeout) {
	return allocateSlot(
		new SlotRequestId(),
		scheduledUnit,
		slotProfile,
		allocationTimeout);
}