org.apache.flink.runtime.taskexecutor.SlotReport Java Examples

/**
 * Tests that we can register task manager and their slots at the slot manager.
 */
@Test
public void testTaskManagerRegistration() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final ResourceID resourceId = ResourceID.generate();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceId, taskExecutorGateway);

	final SlotID slotId1 = new SlotID(resourceId, 0);
	final SlotID slotId2 = new SlotID(resourceId, 1);
	final ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, resourceProfile);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, resourceProfile);
	final SlotReport slotReport = new SlotReport(Arrays.asList(slotStatus1, slotStatus2));

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		assertTrue("The number registered slots does not equal the expected number.", 2 == slotManager.getNumberRegisteredSlots());

		assertNotNull(slotManager.getSlot(slotId1));
		assertNotNull(slotManager.getSlot(slotId2));
	}
}
 

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

/**
 * Tests the completion of pending task manager slots by registering a TaskExecutor.
 */
@Test
public void testPendingTaskManagerSlotCompletion() throws Exception {
	final int numberSlots = 3;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
		assertThat(slotManager.getNumberRegisteredSlots(), is(0));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberSlots - 1);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberSlots - 1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(1));
	}
}
 

/**
 * Tests that only free slots can fulfill/complete a pending task manager slot.
 */
@Test
public void testOnlyFreeSlotsCanFulfillPendingTaskManagerSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotID slotId = new SlotID(taskExecutorConnection.getResourceID(), 0);
		final SlotStatus slotStatus = new SlotStatus(slotId, ResourceProfile.UNKNOWN, jobId, new AllocationID());
		final SlotReport slotReport = new SlotReport(slotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
	}
}
 

/**
 * Reports the current slot allocations for a task manager identified by the given instance id.
 *
 * @param instanceId identifying the task manager for which to report the slot status
 * @param slotReport containing the status for all of its slots
 * @return true if the slot status has been updated successfully, otherwise false
 */
public boolean reportSlotStatus(InstanceID instanceId, SlotReport slotReport) {
	checkInit();

	LOG.debug("Received slot report from instance {}: {}.", instanceId, slotReport);

	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.get(instanceId);

	if (null != taskManagerRegistration) {

		for (SlotStatus slotStatus : slotReport) {
			updateSlot(slotStatus.getSlotID(), slotStatus.getAllocationID(), slotStatus.getJobID());
		}

		return true;
	} else {
		LOG.debug("Received slot report for unknown task manager with instance id {}. Ignoring this report.", instanceId);

		return false;
	}
}
 

public SlotReport createSlotReport(ResourceID resourceId) {
	final int numberSlots = taskSlots.size();

	List<SlotStatus> slotStatuses = Arrays.asList(new SlotStatus[numberSlots]);

	for (int i = 0; i < numberSlots; i++) {
		TaskSlot taskSlot = taskSlots.get(i);
		SlotID slotId = new SlotID(resourceId, taskSlot.getIndex());

		SlotStatus slotStatus = new SlotStatus(
			slotId,
			taskSlot.getResourceProfile(),
			taskSlot.getJobId(),
			taskSlot.getAllocationId());

		slotStatuses.set(i, slotStatus);
	}

	final SlotReport slotReport = new SlotReport(slotStatuses);

	return slotReport;
}
 

/**
 * Tests that if we have received a slot report with some allocated slots, then we don't accept
 * slot requests with allocated allocation ids.
 */
@Test
public void testDuplicatePendingSlotRequestAfterSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);

	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile, jobId, allocationId);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		assertFalse(slotManager.registerSlotRequest(slotRequest));
	}
}
 

/**
 * Tests that the slot manager ignores slot reports of unknown origin (not registered
 * task managers).
 */
@Test
public void testReceivingUnknownSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus unknownSlotStatus = new SlotStatus(unknownSlotId, unknownResourceProfile);
	final SlotReport unknownSlotReport = new SlotReport(unknownSlotStatus);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		// check that we don't have any slots registered
		assertTrue(0 == slotManager.getNumberRegisteredSlots());

		// this should not update anything since the instance id is not known to the slot manager
		assertFalse(slotManager.reportSlotStatus(unknownInstanceID, unknownSlotReport));

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that the slot manager ignores slot reports of unknown origin (not registered
 * task managers).
 */
@Test
public void testReceivingUnknownSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus unknownSlotStatus = new SlotStatus(unknownSlotId, unknownResourceProfile);
	final SlotReport unknownSlotReport = new SlotReport(unknownSlotStatus);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		// check that we don't have any slots registered
		assertTrue(0 == slotManager.getNumberRegisteredSlots());

		// this should not update anything since the instance id is not known to the slot manager
		assertFalse(slotManager.reportSlotStatus(unknownInstanceID, unknownSlotReport));

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that if we have received a slot report with some allocated slots, then we don't accept
 * slot requests with allocated allocation ids.
 */
@Test
public void testDuplicatePendingSlotRequestAfterSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);

	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile, jobId, allocationId);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		assertFalse(slotManager.registerSlotRequest(slotRequest));
	}
}
 

/**
 * Tests the completion of pending task manager slots by registering a TaskExecutor.
 */
@Test
public void testPendingTaskManagerSlotCompletion() throws Exception {
	final int numberSlots = 3;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
		assertThat(slotManager.getNumberRegisteredSlots(), is(0));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberSlots - 1);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberSlots - 1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(1));
	}
}
 

/**
 * Tests that only free slots can fulfill/complete a pending task manager slot.
 */
@Test
public void testOnlyFreeSlotsCanFulfillPendingTaskManagerSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotID slotId = new SlotID(taskExecutorConnection.getResourceID(), 0);
		final SlotStatus slotStatus = new SlotStatus(slotId, ResourceProfile.UNKNOWN, jobId, new AllocationID());
		final SlotReport slotReport = new SlotReport(slotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
	}
}
 

/**
 * Tests that we can register task manager and their slots at the slot manager.
 */
@Test
public void testTaskManagerRegistration() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final ResourceID resourceId = ResourceID.generate();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceId, taskExecutorGateway);

	final SlotID slotId1 = new SlotID(resourceId, 0);
	final SlotID slotId2 = new SlotID(resourceId, 1);
	final ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, resourceProfile);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, resourceProfile);
	final SlotReport slotReport = new SlotReport(Arrays.asList(slotStatus1, slotStatus2));

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		assertTrue("The number registered slots does not equal the expected number.", 2 == slotManager.getNumberRegisteredSlots());

		assertNotNull(slotManager.getSlot(slotId1));
		assertNotNull(slotManager.getSlot(slotId2));
	}
}
 

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

/**
 * Reports the current slot allocations for a task manager identified by the given instance id.
 *
 * @param instanceId identifying the task manager for which to report the slot status
 * @param slotReport containing the status for all of its slots
 * @return true if the slot status has been updated successfully, otherwise false
 */
@Override
public boolean reportSlotStatus(InstanceID instanceId, SlotReport slotReport) {
	checkInit();

	LOG.debug("Received slot report from instance {}: {}.", instanceId, slotReport);

	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.get(instanceId);

	if (null != taskManagerRegistration) {

		for (SlotStatus slotStatus : slotReport) {
			updateSlot(slotStatus.getSlotID(), slotStatus.getAllocationID(), slotStatus.getJobID());
		}

		return true;
	} else {
		LOG.debug("Received slot report for unknown task manager with instance id {}. Ignoring this report.", instanceId);

		return false;
	}
}
 

@Override
public void heartbeatFromTaskManager(ResourceID heartbeatOrigin, SlotReport slotReport) {
	final BiConsumer<ResourceID, SlotReport> currentTaskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;

	if (currentTaskExecutorHeartbeatConsumer != null) {
		currentTaskExecutorHeartbeatConsumer.accept(heartbeatOrigin, slotReport);
	}
}
 

/**
 * Tests that a task manager timeout does not remove the slots from the SlotManager.
 * A timeout should only trigger the {@link ResourceActions#releaseResource(InstanceID, Exception)}
 * callback. The receiver of the callback can then decide what to do with the TaskManager.
 *
 * <p>See FLINK-7793
 */
@Test
public void testTaskManagerTimeoutDoesNotRemoveSlots() throws Exception {
	final Time taskManagerTimeout = Time.milliseconds(10L);
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final ResourceActions resourceActions = mock(ResourceActions.class);
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

	when(taskExecutorGateway.canBeReleased()).thenReturn(CompletableFuture.completedFuture(true));

	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
	final SlotStatus slotStatus = new SlotStatus(
		new SlotID(resourceID, 0),
		new ResourceProfile(1.0, 1));
	final SlotReport initialSlotReport = new SlotReport(slotStatus);

	try (final SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.build()) {

		slotManager.start(resourceManagerId, Executors.directExecutor(), resourceActions);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		verify(resourceActions, timeout(taskManagerTimeout.toMilliseconds() * 20L).atLeast(1)).releaseResource(eq(taskExecutorConnection.getInstanceID()), any(Exception.class));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID());

		assertEquals(0, slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that a task manager timeout does not remove the slots from the SlotManager.
 * A timeout should only trigger the {@link ResourceActions#releaseResource(InstanceID, Exception)}
 * callback. The receiver of the callback can then decide what to do with the TaskManager.
 *
 * <p>See FLINK-7793
 */
@Test
public void testTaskManagerTimeoutDoesNotRemoveSlots() throws Exception {
	final Time taskManagerTimeout = Time.milliseconds(10L);
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final ResourceActions resourceActions = mock(ResourceActions.class);
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

	when(taskExecutorGateway.canBeReleased()).thenReturn(CompletableFuture.completedFuture(true));

	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
	final SlotStatus slotStatus = new SlotStatus(
		new SlotID(resourceID, 0),
		new ResourceProfile(1.0, 1));
	final SlotReport initialSlotReport = new SlotReport(slotStatus);

	try (final SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.build()) {

		slotManager.start(resourceManagerId, Executors.directExecutor(), resourceActions);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		verify(resourceActions, timeout(taskManagerTimeout.toMilliseconds() * 20L).atLeast(1)).releaseResource(eq(taskExecutorConnection.getInstanceID()), any(Exception.class));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID());

		assertEquals(0, slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that duplicate slot requests (requests with an already registered allocation id) are
 * also detected after a pending slot request has been fulfilled but not yet freed.
 */
@Test
public void testDuplicatePendingSlotRequestAfterSuccessfulAllocation() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final AtomicInteger allocateResourceCalls = new AtomicInteger(0);
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceConsumer(resourceProfile -> allocateResourceCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = new ResourceProfile(1.0, 2);
	final ResourceProfile resourceProfile2 = new ResourceProfile(2.0, 1);
	final SlotRequest slotRequest1 = new SlotRequest(new JobID(), allocationId, resourceProfile1, "foobar");
	final SlotRequest slotRequest2 = new SlotRequest(new JobID(), allocationId, resourceProfile2, "barfoo");

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();

	final ResourceID resourceID = ResourceID.generate();

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile1);
	final SlotReport slotReport = new SlotReport(slotStatus);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);
		assertTrue(slotManager.registerSlotRequest(slotRequest1));

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());

		assertFalse(slotManager.registerSlotRequest(slotRequest2));
	}

	// check that we have only called the resource allocation only for the first slot request,
	// since the second request is a duplicate
	assertThat(allocateResourceCalls.get(), is(0));
}
 

@Override
public void reportPayload(final ResourceID resourceID, final SlotReport payload) {
	validateRunsInMainThread();
	final WorkerRegistration<WorkerType> workerRegistration = taskExecutors.get(resourceID);

	if (workerRegistration == null) {
		log.debug("Received slot report from TaskManager {} which is no longer registered.", resourceID);
	} else {
		InstanceID instanceId = workerRegistration.getInstanceID();

		slotManager.reportSlotStatus(instanceId, payload);
	}
}
 

@Nonnull
private SlotReport createSlotReport(ResourceID taskExecutorResourceId, int numberSlots, ResourceProfile resourceProfile) {
	final Set<SlotStatus> slotStatusSet = new HashSet<>(numberSlots);
	for (int i = 0; i < numberSlots; i++) {
		slotStatusSet.add(new SlotStatus(new SlotID(taskExecutorResourceId, i), resourceProfile));
	}

	return new SlotReport(slotStatusSet);
}
 

/**
 * Tests that a different slot can fulfill a pending slot request. If the
 * pending slot request has a pending task manager slot assigned, it should
 * be freed.
 */
@Test
public void testRegistrationOfDifferentSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		final ResourceProfile requestedSlotProfile = new ResourceProfile(1.0, 1);

		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId, requestedSlotProfile)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));

		final int numberOfferedSlots = 1;
		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final ResourceProfile offeredSlotProfile = new ResourceProfile(2.0, 2);
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberOfferedSlots, offeredSlotProfile);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberOfferedSlots));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(0));
	}
}
 

/**
 * Tests that a different slot can fulfill a pending slot request. If the
 * pending slot request has a pending task manager slot assigned, it should
 * be freed.
 */
@Test
public void testRegistrationOfDifferentSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		final ResourceProfile requestedSlotProfile = new ResourceProfile(1.0, 1);

		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId, requestedSlotProfile)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));

		final int numberOfferedSlots = 1;
		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final ResourceProfile offeredSlotProfile = new ResourceProfile(2.0, 2);
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberOfferedSlots, offeredSlotProfile);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberOfferedSlots));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(0));
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final WorkerRegistration<WorkerType> workerTypeWorkerRegistration = taskExecutors.get(taskManagerResourceId);

	if (workerTypeWorkerRegistration.getInstanceID().equals(taskManagerRegistrationId)) {
		slotManager.registerTaskManager(workerTypeWorkerRegistration, slotReport);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		return FutureUtils.completedExceptionally(new ResourceManagerException(String.format("Unknown TaskManager registration id %s.", taskManagerRegistrationId)));
	}
}
 

/**
 * Checks that un-registering a pending slot request will cancel it, removing it from all
 * assigned task manager slots and then remove it from the slot manager.
 */
@Test
public void testUnregisterPendingSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> new CompletableFuture<>())
		.createTestingTaskExecutorGateway();

	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final SlotRequest slotRequest = new SlotRequest(new JobID(), allocationId, resourceProfile, "foobar");

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		slotManager.registerSlotRequest(slotRequest);

		assertNotNull(slotManager.getSlotRequest(allocationId));

		assertTrue(slot.getState() == TaskManagerSlot.State.PENDING);

		slotManager.unregisterSlotRequest(allocationId);

		assertNull(slotManager.getSlotRequest(allocationId));

		slot = slotManager.getSlot(slotId);
		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
	}
}
 

@Override
public void heartbeatFromTaskManager(ResourceID heartbeatOrigin, SlotReport slotReport) {
	final BiConsumer<ResourceID, SlotReport> currentTaskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;

	if (currentTaskExecutorHeartbeatConsumer != null) {
		currentTaskExecutorHeartbeatConsumer.accept(heartbeatOrigin, slotReport);
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> currentSendSlotReportFunction = sendSlotReportFunction;

	if (currentSendSlotReportFunction != null) {
		return currentSendSlotReportFunction.apply(Tuple3.of(taskManagerResourceId, taskManagerRegistrationId, slotReport));
	} else {
		return CompletableFuture.completedFuture(Acknowledge.get());
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final WorkerRegistration<WorkerType> workerTypeWorkerRegistration = taskExecutors.get(taskManagerResourceId);

	if (workerTypeWorkerRegistration.getInstanceID().equals(taskManagerRegistrationId)) {
		slotManager.registerTaskManager(workerTypeWorkerRegistration, slotReport);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		return FutureUtils.completedExceptionally(new ResourceManagerException(String.format("Unknown TaskManager registration id %s.", taskManagerRegistrationId)));
	}
}
 

@Override
public void reportPayload(final ResourceID resourceID, final SlotReport payload) {
	validateRunsInMainThread();
	final WorkerRegistration<WorkerType> workerRegistration = taskExecutors.get(resourceID);

	if (workerRegistration == null) {
		log.debug("Received slot report from TaskManager {} which is no longer registered.", resourceID);
	} else {
		InstanceID instanceId = workerRegistration.getInstanceID();

		slotManager.reportSlotStatus(instanceId, payload);
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> currentSendSlotReportFunction = sendSlotReportFunction;

	if (currentSendSlotReportFunction != null) {
		return currentSendSlotReportFunction.apply(Tuple3.of(taskManagerResourceId, taskManagerRegistrationId, slotReport));
	} else {
		return CompletableFuture.completedFuture(Acknowledge.get());
	}
}