org.apache.flink.runtime.jobmanager.scheduler.ScheduledUnit 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());
		}
	}
}
 

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

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

	componentMainThreadExecutor.assertRunningInMainThread();

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

@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 + ')');
}
 

@Test
public void allocateSingleSlot_withInputPreference_inputPreferenceHasPrecedenceOverSpreadOut() {
	registerTaskExecutors(2, 2);

	final ScheduledUnit sourceSlotRequest = createSimpleSlotRequest();
	final ScheduledUnit sinkSlotRequest = createSimpleSlotRequest();

	final CompletableFuture<LogicalSlot> sourceSlotFuture = allocateSlot(sourceSlotRequest);
	final TaskManagerLocation sourceTaskManagerLocation = getTaskManagerLocation(sourceSlotFuture);

	Collection<TaskManagerLocation> preferredLocations = Collections.singleton(sourceTaskManagerLocation);
	final CompletableFuture<LogicalSlot> sinkSlotFuture = allocateSlotWithInputPreference(sinkSlotRequest, preferredLocations);
	final TaskManagerLocation sinkTaskManagerLocation = getTaskManagerLocation(sinkSlotFuture);

	// input preference should have precedence over task spread out
	assertThat(sinkTaskManagerLocation, is(equalTo(sourceTaskManagerLocation)));
}
 

@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;
}
 

@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 + ')');
}
 

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

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

@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),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

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

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

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

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

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

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

/**
 * 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 allocateSharedSlot_withNoRequirements_selectsSlotsSoThatWorkloadIsSpreadOut() {
	final int numberSlotsPerTaskExecutor = 2;
	final int numberTaskExecutors = 2;
	final int numberSlots = numberTaskExecutors * numberSlotsPerTaskExecutor;

	registerTaskExecutors(numberTaskExecutors, numberSlotsPerTaskExecutor);

	final JobVertexID sourceJobVertexId = new JobVertexID();
	final JobVertexID sinkJobVertexId = new JobVertexID();
	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	final List<ScheduledUnit> sourceScheduledUnits = IntStream.range(0, numberSlots)
		.mapToObj(ignored -> createSharedSlotRequest(sourceJobVertexId, slotSharingGroupId))
		.collect(Collectors.toList());

	final List<ScheduledUnit> sinkScheduledUnits = IntStream.range(0, numberTaskExecutors)
		.mapToObj(ignored -> createSharedSlotRequest(sinkJobVertexId, slotSharingGroupId))
		.collect(Collectors.toList());

	sourceScheduledUnits.forEach(this::allocateSlot);
	final Set<TaskManagerLocation> sinkLocations = sinkScheduledUnits.stream()
		.map(this::allocateSlot)
		.map(this::getTaskManagerLocation)
		.collect(Collectors.toSet());

	// verify that the sinks have been evenly spread across the available TaskExecutors
	assertThat(sinkLocations, hasSize(numberTaskExecutors));
}
 

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

@Test
public void allocateSingleSlot_withNoRequirements_selectsSlotSoThatWorkloadIsSpreadOut() {
	registerTaskExecutors(2, 4);

	final ScheduledUnit firstSlotRequest = createSimpleSlotRequest();
	final ScheduledUnit secondSlotRequest = createSimpleSlotRequest();

	final CompletableFuture<LogicalSlot> firstSlotFuture = allocateSlot(firstSlotRequest);
	final CompletableFuture<LogicalSlot> secondSlotFuture = allocateSlot(secondSlotRequest);

	final TaskManagerLocation firstTaskManagerLocation = getTaskManagerLocation(firstSlotFuture);
	final TaskManagerLocation secondTaskManagerLocation = getTaskManagerLocation(secondSlotFuture);

	assertThat(firstTaskManagerLocation, is(not(equalTo(secondTaskManagerLocation))));
}
 

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

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

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

	try (TestingSlotPool pool = createTestingSlotPool(jid)) {

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

		SlotRequestId requestId = new SlotRequestId();
		CompletableFuture<LogicalSlot> future = testMainThreadExecutor.execute(() -> scheduler.allocateSlot(
			requestId,
			new ScheduledUnit(getExecution()),
			SlotProfile.noLocality(DEFAULT_TESTING_PROFILE),
			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 testSlotAllocationNoResourceManager() throws Exception {
	final JobID jid = new JobID();

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

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

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

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

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

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

		pool.start(JobMasterId.generate(), "foobar", testMainThreadExecutor.getMainThreadExecutor());
		Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.INSTANCE, pool);
		scheduler.start(testMainThreadExecutor.getMainThreadExecutor());

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

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

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

	try (TestingSlotPool pool = createTestingSlotPool(jid)) {

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

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

@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),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

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

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

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

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

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

	slotAllocationRequests.add(Tuple3.of(slotRequestId, task, slotProfile));
	if (slotAllocationDisabled) {
		return new CompletableFuture<>();
	} else {
		return CompletableFuture.completedFuture(new TestingLogicalSlotBuilder().createTestingLogicalSlot());
	}
}
 

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

@Override
public CompletableFuture<LogicalSlot> allocateSlot(SlotRequestId slotRequestId, ScheduledUnit task, SlotProfile slotProfile, Time timeout) {
	Preconditions.checkState(!slotFutures.containsKey(slotRequestId));
	final CompletableFuture<LogicalSlot> slotFuture = slotFutureCreator.apply(slotRequestId);

	slotFutures.put(slotRequestId, slotFuture);

	return slotFuture;
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
	SlotRequestId slotRequestId,
	ScheduledUnit scheduledUnit,
	SlotProfile slotProfile,
	boolean allowQueuedScheduling,
	Time allocationTimeout) {
	log.debug("Received slot request [{}] for task: {}", slotRequestId, scheduledUnit.getTaskToExecute());

	componentMainThreadExecutor.assertRunningInMainThread();

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

	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) {
			cancelSlotRequest(
				slotRequestId,
				scheduledUnit.getSlotSharingGroupId(),
				failure);
			allocationResultFuture.completeExceptionally(failure);
		} else {
			allocationResultFuture.complete(slot);
		}
	});

	return allocationResultFuture;
}