org.apache.flink.runtime.taskexecutor.slot.SlotOffer Java Examples

@Override
public Collection<SlotOffer> offerSlots(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		Collection<SlotOffer> offers) {

	ArrayList<SlotOffer> result = new ArrayList<>(offers.size());

	for (SlotOffer offer : offers) {
		if (offerSlot(
			taskManagerLocation,
			taskManagerGateway,
			offer)) {

			result.add(offer);
		}
	}

	return result;
}
 

@Override
public Collection<SlotOffer> offerSlots(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		Collection<SlotOffer> offers) {

	ArrayList<SlotOffer> result = new ArrayList<>(offers.size());

	for (SlotOffer offer : offers) {
		if (offerSlot(
			taskManagerLocation,
			taskManagerGateway,
			offer)) {

			result.add(offer);
		}
	}

	return result;
}
 

@Override
public CompletableFuture<Collection<SlotOffer>> offerSlots(
		final ResourceID taskManagerId,
		final Collection<SlotOffer> slots,
		final Time timeout) {

	Tuple2<TaskManagerLocation, TaskExecutorGateway> taskManager = registeredTaskManagers.get(taskManagerId);

	if (taskManager == null) {
		return FutureUtils.completedExceptionally(new Exception("Unknown TaskManager " + taskManagerId));
	}

	final TaskManagerLocation taskManagerLocation = taskManager.f0;
	final TaskExecutorGateway taskExecutorGateway = taskManager.f1;

	final RpcTaskManagerGateway rpcTaskManagerGateway = new RpcTaskManagerGateway(taskExecutorGateway, getFencingToken());

	return CompletableFuture.completedFuture(
		slotPool.offerSlots(
			taskManagerLocation,
			rpcTaskManagerGateway,
			slots));
}
 

private void registerTaskExecutorAtJobMaster(
		TestingRpcService rpcService,
		JobMasterGateway jobMasterGateway,
		TaskExecutorGateway taskExecutorGateway,
		SettableLeaderRetrievalService rmLeaderRetrievalService) throws ExecutionException, InterruptedException {

	final AllocationIdsResourceManagerGateway resourceManagerGateway = new AllocationIdsResourceManagerGateway();
	rpcService.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
	rmLeaderRetrievalService.notifyListener(resourceManagerGateway.getAddress(), resourceManagerGateway.getFencingToken().toUUID());

	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	jobMasterGateway.registerTaskManager(taskExecutorGateway.getAddress(), localTaskManagerUnresolvedLocation, testingTimeout).get();

	final AllocationID allocationId = resourceManagerGateway.takeAllocationId();
	Collection<SlotOffer> slotOffers = Collections.singleton(new SlotOffer(allocationId, 0, ResourceProfile.UNKNOWN));

	jobMasterGateway.offerSlots(localTaskManagerUnresolvedLocation.getResourceID(), slotOffers, testingTimeout).get();
}
 

private List<AllocationID> registerAndOfferSlots(TaskManagerLocation taskManagerLocation, SlotPoolImpl slotPool, int numberOfSlotsToRegister) {
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());
	final List<AllocationID> allocationIds = IntStream.range(0, numberOfSlotsToRegister)
		.mapToObj(ignored -> new AllocationID())
		.collect(Collectors.toList());

	Collection<SlotOffer> slotOffers = IntStream.range(0, numberOfSlotsToRegister)
		.mapToObj(index -> new SlotOffer(allocationIds.get(index), index, ResourceProfile.ANY))
		.collect(Collectors.toList());

	slotPool.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		slotOffers);

	return allocationIds;
}
 

private static Scheduler createSchedulerWithSlots(SlotPool slotPool, TaskManagerLocation taskManagerLocation, int numSlots) throws Exception {
	final TaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();
	setupSlotPool(slotPool);
	Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), slotPool);
	scheduler.start(mainThreadExecutor);
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	final List<SlotOffer> slotOffers = new ArrayList<>(NUM_TASKS);
	for (int i = 0; i < numSlots; i++) {
		final AllocationID allocationId = new AllocationID();
		final SlotOffer slotOffer = new SlotOffer(allocationId, 0, ResourceProfile.ANY);
		slotOffers.add(slotOffer);
	}

	slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

	return scheduler;
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

@Override
public CompletableFuture<Collection<SlotOffer>> offerSlots(
		final ResourceID taskManagerId,
		final Collection<SlotOffer> slots,
		final Time timeout) {

	Tuple2<TaskManagerLocation, TaskExecutorGateway> taskManager = registeredTaskManagers.get(taskManagerId);

	if (taskManager == null) {
		return FutureUtils.completedExceptionally(new Exception("Unknown TaskManager " + taskManagerId));
	}

	final TaskManagerLocation taskManagerLocation = taskManager.f0;
	final TaskExecutorGateway taskExecutorGateway = taskManager.f1;

	final RpcTaskManagerGateway rpcTaskManagerGateway = new RpcTaskManagerGateway(taskExecutorGateway, getFencingToken());

	return CompletableFuture.completedFuture(
		slotPool.offerSlots(
			taskManagerLocation,
			rpcTaskManagerGateway,
			slots));
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

private Collection<SlotOffer> registerSlotsAtJobMaster(
	int numberSlots,
	JobMasterGateway jobMasterGateway,
	TaskExecutorGateway taskExecutorGateway) throws ExecutionException, InterruptedException {
	return registerSlotsAtJobMaster(
		numberSlots,
		jobMasterGateway,
		taskExecutorGateway,
		new LocalUnresolvedTaskManagerLocation());
}
 

/**
 * Tests that create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = new SlotPoolImpl(jobId)) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

private Collection<SlotOffer> registerSlotsAtJobMaster(
	int numberSlots,
	JobMasterGateway jobMasterGateway,
	TaskExecutorGateway taskExecutorGateway) throws ExecutionException, InterruptedException {
	return registerSlotsAtJobMaster(
		numberSlots,
		jobMasterGateway,
		taskExecutorGateway,
		new LocalTaskManagerLocation());
}
 

@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 create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.ANY));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.ANY));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

private void offerSlotsToJobManager(final JobID jobId) {
	final JobManagerConnection jobManagerConnection = jobManagerTable.get(jobId);

	if (jobManagerConnection == null) {
		log.debug("There is no job manager connection to the leader of job {}.", jobId);
	} else {
		if (taskSlotTable.hasAllocatedSlots(jobId)) {
			log.info("Offer reserved slots to the leader of job {}.", jobId);

			final JobMasterGateway jobMasterGateway = jobManagerConnection.getJobManagerGateway();

			final Iterator<TaskSlot> reservedSlotsIterator = taskSlotTable.getAllocatedSlots(jobId);
			final JobMasterId jobMasterId = jobManagerConnection.getJobMasterId();

			final Collection<SlotOffer> reservedSlots = new HashSet<>(2);

			while (reservedSlotsIterator.hasNext()) {
				SlotOffer offer = reservedSlotsIterator.next().generateSlotOffer();
				reservedSlots.add(offer);
			}

			CompletableFuture<Collection<SlotOffer>> acceptedSlotsFuture = jobMasterGateway.offerSlots(
				getResourceID(),
				reservedSlots,
				taskManagerConfiguration.getTimeout());

			acceptedSlotsFuture.whenCompleteAsync(
				handleAcceptedSlotOffers(jobId, jobMasterGateway, jobMasterId, reservedSlots),
				getMainThreadExecutor());
		} else {
			log.debug("There are no unassigned slots for the job {}.", jobId);
		}
	}
}
 

/**
 * Tests that create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

@Test
public void testAllocateSimpleSlot() throws Exception {
	CompletableFuture<SlotRequest> slotRequestFuture = new CompletableFuture<>();
	resourceManagerGateway.setRequestSlotConsumer(slotRequestFuture::complete);

	try (SlotPoolImpl slotPool = new SlotPoolImpl(jobId)) {
		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);
		slotPool.registerTaskManager(taskManagerLocation.getResourceID());

		SlotRequestId requestId = new SlotRequestId();
		CompletableFuture<LogicalSlot> future = scheduler.allocateSlot(
			requestId,
			new DummyScheduledUnit(),
			SlotProfile.noLocality(DEFAULT_TESTING_PROFILE),
			true,
			timeout);
		assertFalse(future.isDone());

		final SlotRequest slotRequest = slotRequestFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		final SlotOffer slotOffer = new SlotOffer(
			slotRequest.getAllocationId(),
			0,
			DEFAULT_TESTING_PROFILE);

		assertTrue(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotOffer));

		LogicalSlot slot = future.get(1, TimeUnit.SECONDS);
		assertTrue(future.isDone());
		assertTrue(slot.isAlive());
		assertEquals(taskManagerLocation, slot.getTaskManagerLocation());
	}
}
 

@Test
public void testFailingAllocationFailsRemappedPendingSlotRequests() throws Exception {
	final List<AllocationID> allocations = new ArrayList<>();
	resourceManagerGateway.setRequestSlotConsumer(slotRequest -> allocations.add(slotRequest.getAllocationId()));

	try (SlotPoolImpl slotPool = setUpSlotPool()) {
		final CompletableFuture<PhysicalSlot> slotFuture1 = requestNewAllocatedSlot(slotPool, new SlotRequestId());
		final CompletableFuture<PhysicalSlot> slotFuture2 = requestNewAllocatedSlot(slotPool, new SlotRequestId());

		final AllocationID allocationId1 = allocations.get(0);
		final AllocationID allocationId2 = allocations.get(1);

		final TaskManagerLocation location = new LocalTaskManagerLocation();
		final SlotOffer slotOffer = new SlotOffer(allocationId2, 0, ResourceProfile.ANY);
		slotPool.registerTaskManager(location.getResourceID());
		slotPool.offerSlot(location, new SimpleAckingTaskManagerGateway(), slotOffer);

		assertThat(slotFuture1.isDone(), is(true));
		assertThat(slotFuture2.isDone(), is(false));

		final FlinkException cause = new FlinkException("Fail pending slot request failure.");
		final Optional<ResourceID> responseFuture = slotPool.failAllocation(allocationId1, cause);

		assertThat(responseFuture.isPresent(), is(false));

		try {
			slotFuture2.getNow(null);
			fail("Expected a slot allocation failure.");
		} catch (Throwable t) {
			assertThat(ExceptionUtils.stripCompletionException(t), equalTo(cause));
		}
	}
}
 

private void internalOfferSlotsToJobManager(JobTable.Connection jobManagerConnection) {
	final JobID jobId = jobManagerConnection.getJobId();

	if (taskSlotTable.hasAllocatedSlots(jobId)) {
		log.info("Offer reserved slots to the leader of job {}.", jobId);

		final JobMasterGateway jobMasterGateway = jobManagerConnection.getJobManagerGateway();

		final Iterator<TaskSlot<Task>> reservedSlotsIterator = taskSlotTable.getAllocatedSlots(jobId);
		final JobMasterId jobMasterId = jobManagerConnection.getJobMasterId();

		final Collection<SlotOffer> reservedSlots = new HashSet<>(2);

		while (reservedSlotsIterator.hasNext()) {
			SlotOffer offer = reservedSlotsIterator.next().generateSlotOffer();
			reservedSlots.add(offer);
		}

		CompletableFuture<Collection<SlotOffer>> acceptedSlotsFuture = jobMasterGateway.offerSlots(
			getResourceID(),
			reservedSlots,
			taskManagerConfiguration.getTimeout());

		acceptedSlotsFuture.whenCompleteAsync(
			handleAcceptedSlotOffers(jobId, jobMasterGateway, jobMasterId, reservedSlots),
			getMainThreadExecutor());
	} else {
		log.debug("There are no unassigned slots for the job {}.", jobId);
	}
}
 

public static ResourceID offerSlots(
		SlotPoolImpl slotPool,
		ComponentMainThreadExecutor mainThreadExecutor,
		List<ResourceProfile> resourceProfiles,
		TaskManagerGateway taskManagerGateway) {
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	CompletableFuture.runAsync(
		() -> {
			slotPool.registerTaskManager(taskManagerLocation.getResourceID());

			final Collection<SlotOffer> slotOffers = IntStream
				.range(0, resourceProfiles.size())
				.mapToObj(i -> new SlotOffer(new AllocationID(), i, resourceProfiles.get(i)))
				.collect(Collectors.toList());

			final Collection<SlotOffer> acceptedOffers = slotPool.offerSlots(
				taskManagerLocation,
				taskManagerGateway,
				slotOffers);

			assertThat(acceptedOffers, is(slotOffers));
		},
		mainThreadExecutor
	).join();

	return taskManagerLocation.getResourceID();
}
 

private Collection<SlotOffer> registerSlotsAtJobMaster(
	int numberSlots,
	JobMasterGateway jobMasterGateway,
	TaskExecutorGateway taskExecutorGateway) throws ExecutionException, InterruptedException {
	return registerSlotsAtJobMaster(
		numberSlots,
		jobMasterGateway,
		taskExecutorGateway,
		new LocalTaskManagerLocation());
}
 

public static ResourceID offerSlots(
		SlotPoolImpl slotPool,
		ComponentMainThreadExecutor mainThreadExecutor,
		List<ResourceProfile> resourceProfiles,
		TaskManagerGateway taskManagerGateway) {
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	CompletableFuture.runAsync(
		() -> {
			slotPool.registerTaskManager(taskManagerLocation.getResourceID());

			final Collection<SlotOffer> slotOffers = IntStream
				.range(0, resourceProfiles.size())
				.mapToObj(i -> new SlotOffer(new AllocationID(), i, resourceProfiles.get(i)))
				.collect(Collectors.toList());

			final Collection<SlotOffer> acceptedOffers = slotPool.offerSlots(
				taskManagerLocation,
				taskManagerGateway,
				slotOffers);

			assertThat(acceptedOffers, is(slotOffers));
		},
		mainThreadExecutor
	).join();

	return taskManagerLocation.getResourceID();
}
 

/**
 * Tests that idle slots which cannot be released will be discarded. See FLINK-11059.
 */
@Test
public void testDiscardIdleSlotIfReleasingFailed() throws Exception {
	final ManualClock clock = new ManualClock();

	try (TestingSlotPoolImpl slotPool = createSlotPoolImpl(clock)) {

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final AllocationID expiredAllocationId = new AllocationID();
		final SlotOffer slotToExpire = new SlotOffer(expiredAllocationId, 0, ResourceProfile.UNKNOWN);

		OneShotLatch freeSlotLatch = new OneShotLatch();
		taskManagerGateway.setFreeSlotFunction((AllocationID allocationId, Throwable cause) -> {
			freeSlotLatch.trigger();
			return FutureUtils.completedExceptionally(new TimeoutException("Test failure"));
		});

		assertThat(slotPool.registerTaskManager(taskManagerLocation.getResourceID()), Matchers.is(true));

		assertThat(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotToExpire), Matchers.is(true));

		clock.advanceTime(timeout.toMilliseconds() + 1, TimeUnit.MILLISECONDS);

		slotPool.triggerCheckIdleSlot();

		freeSlotLatch.await();

		CompletableFuture<LogicalSlot> allocatedSlotFuture = allocateSlot(scheduler, new SlotRequestId());

		try {
			// since the slot must have been discarded, we cannot fulfill the slot request
			allocatedSlotFuture.get(10L, TimeUnit.MILLISECONDS);
			fail("Expected to fail with a timeout.");
		} catch (TimeoutException ignored) {
			// expected
			assertEquals(0, slotPool.getAvailableSlots().size());
		}
	}
}
 

@Test
public void testRetryOnSharedSlotOverAllocated() throws InterruptedException, ExecutionException {
	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 firstAllocatedSlotRp = new ResourceProfile(3.0, 300);
	final ResourceProfile secondAllocatedSlotRp = new ResourceProfile(5.0, 500);

	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 SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final JobVertexID jobVertexId1 = new JobVertexID();
	final JobVertexID jobVertexId2 = new JobVertexID();
	final JobVertexID jobVertexId3 = new JobVertexID();

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

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

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

	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());
	assertFalse(logicalSlotFuture3.isDone());

	final AllocationID allocationId1 = allocationIds.take();

	// This should fulfill the first two requests.
	boolean offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId1,
					0,
					firstAllocatedSlotRp));

	assertTrue(offerFuture);

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

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

	// The third request will retry.
	assertFalse(logicalSlotFuture3.isDone());
	final AllocationID allocationId2 = allocationIds.take();

	offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId2,
					1,
					secondAllocatedSlotRp));

	assertTrue(offerFuture);

	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();
	assertEquals(allocationId2, logicalSlot3.getAllocationId());
}
 

@Test
public void testAllocationFulfilledByReturnedSlot() throws Exception {
	final ArrayBlockingQueue<SlotRequest> slotRequestQueue = new ArrayBlockingQueue<>(2);

	resourceManagerGateway.setRequestSlotConsumer(slotRequest -> {
		while (!slotRequestQueue.offer(slotRequest)) {
			// noop
		}
	});

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {
		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);
		slotPool.registerTaskManager(taskManagerLocation.getResourceID());

		CompletableFuture<LogicalSlot> future1 = scheduler.allocateSlot(
			new SlotRequestId(),
			new DummyScheduledUnit(),
			SlotProfile.noLocality(DEFAULT_TESTING_PROFILE),
			timeout);
		CompletableFuture<LogicalSlot> future2 = scheduler.allocateSlot(
			new SlotRequestId(),
			new DummyScheduledUnit(),
			SlotProfile.noLocality(DEFAULT_TESTING_PROFILE),
			timeout);

		assertFalse(future1.isDone());
		assertFalse(future2.isDone());

		final List<SlotRequest> slotRequests = new ArrayList<>(2);

		for (int i = 0; i < 2; i++) {
			slotRequests.add(slotRequestQueue.poll(timeout.toMilliseconds(), TimeUnit.MILLISECONDS));
		}

		final SlotOffer slotOffer = new SlotOffer(
			slotRequests.get(0).getAllocationId(),
			0,
			DEFAULT_TESTING_PROFILE);

		assertTrue(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotOffer));

		LogicalSlot slot1 = future1.get(1, TimeUnit.SECONDS);
		assertTrue(future1.isDone());
		assertFalse(future2.isDone());

		// return this slot to pool
		slot1.releaseSlot();

		// second allocation fulfilled by previous slot returning
		LogicalSlot slot2 = future2.get(1, TimeUnit.SECONDS);
		assertTrue(future2.isDone());

		assertNotEquals(slot1, slot2);
		assertFalse(slot1.isAlive());
		assertTrue(slot2.isAlive());
		assertEquals(slot1.getTaskManagerLocation(), slot2.getTaskManagerLocation());
		assertEquals(slot1.getPhysicalSlotNumber(), slot2.getPhysicalSlotNumber());
		assertEquals(slot1.getAllocationId(), slot2.getAllocationId());
	}
}
 

/**
 * Tests that the {@link AllocatedSlotReport} contains up to date information and not
 * stale information about the allocated slots on the {@link JobMaster}.
 *
 * <p>This is a probabilistic test case which only fails if executed repeatedly without
 * the fix for FLINK-12863.
 */
@Test
public void testAllocatedSlotReportDoesNotContainStaleInformation() throws Exception {
	final CompletableFuture<Void> assertionFuture = new CompletableFuture<>();
	final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	final AtomicBoolean terminateHeartbeatVerification = new AtomicBoolean(false);
	final OneShotLatch hasReceivedSlotOffers = new OneShotLatch();
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setHeartbeatJobManagerConsumer((taskManagerId, allocatedSlotReport) -> {
			try {
				if (hasReceivedSlotOffers.isTriggered()) {
					assertThat(allocatedSlotReport.getAllocatedSlotInfos(), hasSize(1));
				} else {
					assertThat(allocatedSlotReport.getAllocatedSlotInfos(), empty());
				}
			} catch (AssertionError e) {
				assertionFuture.completeExceptionally(e);
			}

			if (terminateHeartbeatVerification.get()) {
				assertionFuture.complete(null);
			}
		})
		.createTestingTaskExecutorGateway();

	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	final JobManagerSharedServices jobManagerSharedServices = new TestingJobManagerSharedServicesBuilder().build();

	final JobMaster jobMaster = new JobMasterBuilder(JobGraphTestUtils.createSingleVertexJobGraph(), rpcService)
		.withHeartbeatServices(new HeartbeatServices(5L, 1000L))
		.withSlotPoolFactory(new TestingSlotPoolFactory(hasReceivedSlotOffers))
		.createJobMaster();

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

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

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

		// register task manager will trigger monitor heartbeat target, schedule heartbeat request at interval time
		CompletableFuture<RegistrationResponse> registrationResponse = jobMasterGateway.registerTaskManager(
			taskExecutorGateway.getAddress(),
			unresolvedTaskManagerLocation,
			testingTimeout);

		// wait for the completion of the registration
		registrationResponse.get();

		final SlotOffer slotOffer = new SlotOffer(new AllocationID(), 0, ResourceProfile.ANY);

		final CompletableFuture<Collection<SlotOffer>> slotOfferFuture = jobMasterGateway.offerSlots(unresolvedTaskManagerLocation.getResourceID(), Collections.singleton(slotOffer), testingTimeout);

		assertThat(slotOfferFuture.get(), containsInAnyOrder(slotOffer));

		terminateHeartbeatVerification.set(true);

		// make sure that no assertion has been violated
		assertionFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
		jobManagerSharedServices.shutdown();
	}
}
 

/**
 * Tests queued slot scheduling with multiple slot sharing groups.
 */
@Test
public void testQueuedMultipleSlotSharingGroups() throws Exception {
	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(4);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final SlotSharingGroupId slotSharingGroupId1 = new SlotSharingGroupId();
	final SlotSharingGroupId slotSharingGroupId2 = new SlotSharingGroupId();
	final JobVertexID jobVertexId1 = new JobVertexID();
	final JobVertexID jobVertexId2 = new JobVertexID();
	final JobVertexID jobVertexId3 = new JobVertexID();
	final JobVertexID jobVertexId4 = new JobVertexID();

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

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

	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId1,
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId3,
			slotSharingGroupId2,
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture4 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId4,
			slotSharingGroupId2,
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());
	assertFalse(logicalSlotFuture3.isDone());
	assertFalse(logicalSlotFuture4.isDone());

	// we expect two slot requests
	final AllocationID allocationId1 = allocationIds.take();
	final AllocationID allocationId2 = allocationIds.take();

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

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

	assertTrue(offerFuture1);
	assertTrue(offerFuture2);

	LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	LogicalSlot logicalSlot2 = logicalSlotFuture2.get();
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();
	LogicalSlot logicalSlot4 = logicalSlotFuture4.get();

	assertEquals(logicalSlot1.getTaskManagerLocation(), logicalSlot2.getTaskManagerLocation());
	assertEquals(logicalSlot3.getTaskManagerLocation(), logicalSlot4.getTaskManagerLocation());

	assertEquals(allocationId1, logicalSlot1.getAllocationId());
	assertEquals(allocationId2, logicalSlot3.getAllocationId());
}
 

/**
 * Tests queued slot scheduling with a single slot sharing group
 */
@Test
public void testQueuedSharedSlotScheduling() throws Exception {
	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 SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final JobVertexID jobVertexId1 = new JobVertexID();
	final JobVertexID jobVertexId2 = new JobVertexID();

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

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

	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	final AllocationID allocationId1 = allocationIds.take();

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

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

	assertFalse(logicalSlotFuture3.isDone());
	assertFalse(logicalSlotFuture4.isDone());

	allocationIds.take();

	// this should fulfill the first two slot futures
	boolean offerFuture = slotPool.offerSlot(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		new SlotOffer(
			allocationId1,
			0,
			ResourceProfile.UNKNOWN));

	assertTrue(offerFuture);

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

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

	assertFalse(logicalSlotFuture3.isDone());
	assertFalse(logicalSlotFuture4.isDone());

	// release the shared slot by releasing the individual tasks
	logicalSlot1.releaseSlot(null);
	logicalSlot2.releaseSlot(null);

	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();
	LogicalSlot logicalSlot4 = logicalSlotFuture4.get();

	assertEquals(logicalSlot3.getTaskManagerLocation(), logicalSlot4.getTaskManagerLocation());
	assertEquals(allocationId1, logicalSlot3.getAllocationId());
	assertEquals(allocationId1, logicalSlot4.getAllocationId());
}
 

/**
 * Tests that when matching from the allocated slot, the remaining resources of the slot
 * will be used instead of the total resource.
 */
@Test
public void testSlotSharingRespectsRemainingResource() throws Exception {
	final ResourceProfile allocatedSlotRp = new ResourceProfile(3.0, 300);
	final ResourceProfile largeRequestResource = new ResourceProfile(2.0, 200);
	final ResourceProfile smallRequestResource = new ResourceProfile(1.0, 100);

	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 SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final JobVertexID jobVertexId1 = new JobVertexID();
	final JobVertexID jobVertexId2 = new JobVertexID();
	final JobVertexID jobVertexId3 = new JobVertexID();

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

	final AllocationID allocationId1 = allocationIds.take();

	// This should fulfill the first request.
	boolean offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId1,
					0,
					allocatedSlotRp));

	assertTrue(offerFuture);
	assertTrue(logicalSlotFuture1.isDone());
	assertEquals(allocationId1, logicalSlotFuture1.get().getAllocationId());

	// The second request should not share the same slot with the first request since it is large.
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId2,
					slotSharingGroupId,
					null),
			true,
			SlotProfile.noLocality(largeRequestResource),
			TestingUtils.infiniteTime());
	assertFalse(logicalSlotFuture2.isDone());

	// The third request should be able to share the same slot with the first request since it is small.
	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId3,
					slotSharingGroupId,
					null),
			true,
			SlotProfile.noLocality(smallRequestResource),
			TestingUtils.infiniteTime());
	assertTrue(logicalSlotFuture3.isDone());
	assertEquals(allocationId1, logicalSlotFuture1.get().getAllocationId());

	// The second request should be finally fulfilled by a new slot.
	final AllocationID allocationId2 = allocationIds.take();
	// This should fulfill the first two requests.
	offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId2,
					0,
					allocatedSlotRp));

	assertTrue(offerFuture);
	assertTrue(logicalSlotFuture2.isDone());
	assertEquals(allocationId2, logicalSlotFuture2.get().getAllocationId());
}