org.apache.flink.runtime.resourcemanager.registration.TaskExecutorConnection Java Examples

/**
 * 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 = new TestingResourceActionsBuilder().build();

	final TaskExecutorConnection taskManagerConnection = createTaskExecutorConnection();
	final ResourceID resourceID = taskManagerConnection.getResourceID();

	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotID slotId = new SlotID(resourceID, 0);
	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 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 = ResourceProfile.fromResources(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 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 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 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().build();
	final ResourceProfile resourceProfile =
		SlotManagerImpl.generateDefaultSlotResourceProfile(WORKER_RESOURCE_SPEC, numberSlots);

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions, numberSlots)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId, resourceProfile)), 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, resourceProfile, SlotManagerImplTest::createEmptySlotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

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

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

/**
 * Tests that the unregister cause is being forwarded when failing allocations.
 */
@Test
public void unregisterTaskManager_withAllocatedSlot_failsAllocationsWithCause() throws Exception {
	CompletableFuture<Exception> allocationFailureCause = new CompletableFuture<>();
	TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setNotifyAllocationFailureConsumer(jobIDAllocationIDExceptionTuple3 -> allocationFailureCause.complete(jobIDAllocationIDExceptionTuple3.f2))
		.build();

	FlinkException failureCause = new FlinkException("unregisterTaskManager test exception.");

	try (SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		SlotReport slotReport = createSingleAllocatedSlotReport(taskExecutorConnection.getResourceID(), new JobID());
		slotManager.registerTaskManager(taskExecutorConnection, slotReport);
		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID(), failureCause);

		assertThat(allocationFailureCause.get(), FlinkMatchers.containsCause(failureCause));
	}
}
 

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

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions, numberSlots)) {
		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.ANY, 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 idle task managers time out after the configured timeout. A timed out task manager
 * will be removed from the slot manager and the resource manager will be notified about the
 * timeout, if it can be released.
 */
@Test
public void testTaskManagerTimeout() throws Exception {
	final long tmTimeout = 10L;

	final CompletableFuture<InstanceID> releaseFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releaseFuture.complete(instanceID))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

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

	final Executor mainThreadExecutor = TestingUtils.defaultExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(Time.milliseconds(tmTimeout))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		mainThreadExecutor.execute(() -> slotManager.registerTaskManager(taskManagerConnection, slotReport));

		assertThat(releaseFuture.get(), is(equalTo(taskManagerConnection.getInstanceID())));
	}
}
 

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

@Nonnull
private TaskManagerSlot createAndRegisterTaskManagerSlot(SlotID slotId, ResourceProfile resourceProfile, TaskExecutorConnection taskManagerConnection) {
	final TaskManagerSlot slot = new TaskManagerSlot(
		slotId,
		resourceProfile,
		taskManagerConnection);
	slots.put(slotId, slot);
	return slot;
}
 

/**
 * Test that the slot manager release resource when the number of slots exceed max limit when new TaskExecutor registered.
 */
@Test
public void testMaxSlotLimitRegisterResource() throws Exception {
	final int numberSlots = 1;
	final int maxSlotNum = 1;
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();

	final CompletableFuture<InstanceID> releasedResourceFuture = new CompletableFuture<>();
	ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releasedResourceFuture.complete(instanceID))
		.build();

	final TaskExecutorGateway taskExecutorGateway1 = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorGateway taskExecutorGateway2 = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final ResourceID resourceId1 = ResourceID.generate();
	final ResourceID resourceId2 = ResourceID.generate();
	final TaskExecutorConnection taskManagerConnection1 = new TaskExecutorConnection(resourceId1, taskExecutorGateway1);
	final TaskExecutorConnection taskManagerConnection2 = new TaskExecutorConnection(resourceId2, taskExecutorGateway2);

	final SlotID slotId1 = new SlotID(resourceId1, 0);
	final SlotID slotId2 = new SlotID(resourceId1, 0);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, ResourceProfile.UNKNOWN);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, ResourceProfile.UNKNOWN);
	final SlotReport slotReport1 = new SlotReport(Collections.singletonList(slotStatus1));
	final SlotReport slotReport2 = new SlotReport(Collections.singletonList(slotStatus2));

	try (SlotManagerImpl slotManager = createSlotManagerBuilder()
		.setNumSlotsPerWorker(numberSlots)
		.setMaxSlotNum(maxSlotNum)
		.buildAndStartWithDirectExec(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection1, slotReport1);
		slotManager.registerTaskManager(taskManagerConnection2, slotReport2);

		assertThat("The number registered slots does not equal the expected number.", slotManager.getNumberRegisteredSlots(), is(1));
		assertNotNull(slotManager.getSlot(slotId1));

		// The second registered task manager should be released.
		assertThat(releasedResourceFuture.get(), is(equalTo(taskManagerConnection2.getInstanceID())));
	}
}
 

/**
 * Registers a slot for the given task manager at the slot manager. The slot is identified by
 * the given slot id. The given resource profile defines the available resources for the slot.
 * The task manager connection can be used to communicate with the task manager.
 *
 * @param slotId identifying the slot on the task manager
 * @param allocationId which is currently deployed in the slot
 * @param resourceProfile of the slot
 * @param taskManagerConnection to communicate with the remote task manager
 */
private void registerSlot(
		SlotID slotId,
		AllocationID allocationId,
		JobID jobId,
		ResourceProfile resourceProfile,
		TaskExecutorConnection taskManagerConnection) {

	if (slots.containsKey(slotId)) {
		// remove the old slot first
		removeSlot(slotId);
	}

	final TaskManagerSlot slot = createAndRegisterTaskManagerSlot(slotId, resourceProfile, taskManagerConnection);

	final PendingTaskManagerSlot pendingTaskManagerSlot;

	if (allocationId == null) {
		pendingTaskManagerSlot = findExactlyMatchingPendingTaskManagerSlot(resourceProfile);
	} else {
		pendingTaskManagerSlot = null;
	}

	if (pendingTaskManagerSlot == null) {
		updateSlot(slotId, allocationId, jobId);
	} else {
		pendingSlots.remove(pendingTaskManagerSlot.getTaskManagerSlotId());
		final PendingSlotRequest assignedPendingSlotRequest = pendingTaskManagerSlot.getAssignedPendingSlotRequest();

		if (assignedPendingSlotRequest == null) {
			handleFreeSlot(slot);
		} else {
			assignedPendingSlotRequest.unassignPendingTaskManagerSlot();
			allocateSlot(slot, assignedPendingSlotRequest);
		}
	}
}
 

/**
 * 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(ignored -> allocateResourceCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = ResourceProfile.fromResources(1.0, 2);
	final ResourceProfile resourceProfile2 = ResourceProfile.fromResources(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));
}
 

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

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

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions, numberSlots)) {
		final JobID jobId = new JobID();
		final ResourceProfile requestedSlotProfile = ResourceProfile.fromResources(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 = ResourceProfile.fromResources(2.0, 2);
		final SlotReport slotReport = createSlotReport(
			taskExecutorConnection.getResourceID(),
			numberOfferedSlots,
			offeredSlotProfile,
			SlotManagerImplTest::createEmptySlotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

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

private static void registerFreeSlot(SlotManager slotManager, ResourceProfile slotProfile) {
	final ResourceID resourceID = ResourceID.generate();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotReport slotReport = new SlotReport(
		Collections.singleton(new SlotStatus(new SlotID(resourceID, 0), slotProfile)));

	slotManager.registerTaskManager(taskExecutorConnection, 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());
	}
}
 

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

public TaskManagerSlot(
		SlotID slotId,
		ResourceProfile resourceProfile,
		TaskExecutorConnection taskManagerConnection) {
	this.slotId = checkNotNull(slotId);
	this.resourceProfile = checkNotNull(resourceProfile);
	this.taskManagerConnection = checkNotNull(taskManagerConnection);

	this.state = State.FREE;
	this.allocationId = null;
	this.assignedSlotRequest = null;
}
 

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

public TaskManagerSlot(
		SlotID slotId,
		ResourceProfile resourceProfile,
		TaskExecutorConnection taskManagerConnection) {
	this.slotId = checkNotNull(slotId);
	this.resourceProfile = checkNotNull(resourceProfile);
	this.taskManagerConnection = checkNotNull(taskManagerConnection);

	this.state = State.FREE;
	this.allocationId = null;
	this.assignedSlotRequest = null;
}
 

private void registerTaskExecutorWithTwoSlots(SlotManagerImpl slotManager, CompletableFuture<JobID> firstRequestSlotFuture) {
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
			firstRequestSlotFuture.complete(slotIDJobIDAllocationIDStringResourceManagerIdTuple6.f1);
			return CompletableFuture.completedFuture(Acknowledge.get());
		})
		.createTestingTaskExecutorGateway();
	final TaskExecutorConnection firstTaskExecutorConnection = createTaskExecutorConnection(taskExecutorGateway);
	final SlotReport firstSlotReport = createSlotReport(firstTaskExecutorConnection.getResourceID(), 2);
	slotManager.registerTaskManager(firstTaskExecutorConnection, firstSlotReport);
}
 

/**
 * Registers a slot for the given task manager at the slot manager. The slot is identified by
 * the given slot id. The given resource profile defines the available resources for the slot.
 * The task manager connection can be used to communicate with the task manager.
 *
 * @param slotId identifying the slot on the task manager
 * @param allocationId which is currently deployed in the slot
 * @param resourceProfile of the slot
 * @param taskManagerConnection to communicate with the remote task manager
 */
private void registerSlot(
		SlotID slotId,
		AllocationID allocationId,
		JobID jobId,
		ResourceProfile resourceProfile,
		TaskExecutorConnection taskManagerConnection) {

	if (slots.containsKey(slotId)) {
		// remove the old slot first
		removeSlot(slotId);
	}

	final TaskManagerSlot slot = createAndRegisterTaskManagerSlot(slotId, resourceProfile, taskManagerConnection);

	final PendingTaskManagerSlot pendingTaskManagerSlot;

	if (allocationId == null) {
		pendingTaskManagerSlot = findExactlyMatchingPendingTaskManagerSlot(resourceProfile);
	} else {
		pendingTaskManagerSlot = null;
	}

	if (pendingTaskManagerSlot == null) {
		updateSlot(slotId, allocationId, jobId);
	} else {
		pendingSlots.remove(pendingTaskManagerSlot.getTaskManagerSlotId());
		final PendingSlotRequest assignedPendingSlotRequest = pendingTaskManagerSlot.getAssignedPendingSlotRequest();

		if (assignedPendingSlotRequest == null) {
			handleFreeSlot(slot);
		} else {
			assignedPendingSlotRequest.unassignPendingTaskManagerSlot();
			allocateSlot(slot, assignedPendingSlotRequest);
		}
	}
}
 

@Nonnull
private TaskManagerSlot createAndRegisterTaskManagerSlot(SlotID slotId, ResourceProfile resourceProfile, TaskExecutorConnection taskManagerConnection) {
	final TaskManagerSlot slot = new TaskManagerSlot(
		slotId,
		resourceProfile,
		taskManagerConnection);
	slots.put(slotId, slot);
	return slot;
}