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

private void testPartitionReleaseOrPromotionOnJobTermination(Function<TestSetup, CompletableFuture<ResourceID>> taskExecutorCallSelector, ExecutionState finalExecutionState) throws Exception {
	final CompletableFuture<TaskDeploymentDescriptor> taskDeploymentDescriptorFuture = new CompletableFuture<>();
	final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setSubmitTaskConsumer((tdd, ignored) -> {
			taskDeploymentDescriptorFuture.complete(tdd);
			return CompletableFuture.completedFuture(Acknowledge.get());
		})
		.createTestingTaskExecutorGateway();

	try (final TestSetup testSetup = new TestSetup(rpcService, testingFatalErrorHandler, testingTaskExecutorGateway)) {
		final JobMasterGateway jobMasterGateway = testSetup.getJobMasterGateway();

		// update the execution state of the only execution to target state
		// this should trigger the job to finish
		final TaskDeploymentDescriptor taskDeploymentDescriptor = taskDeploymentDescriptorFuture.get();
		jobMasterGateway.updateTaskExecutionState(
			new TaskExecutionState(
				taskDeploymentDescriptor.getJobId(),
				taskDeploymentDescriptor.getExecutionAttemptId(),
				finalExecutionState));

		assertThat(taskExecutorCallSelector.apply(testSetup).get(), equalTo(testSetup.getTaskExecutorResourceID()));
	}
}
 

/**
 * 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 that we can retrieve the correct {@link TaskManagerInfo} from the {@link ResourceManager}.
 */
@Test
public void testRequestTaskManagerInfo() throws Exception {
	final ResourceID taskManagerId = ResourceID.generate();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setAddress(UUID.randomUUID().toString()).createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	resourceManager = createAndStartResourceManager(heartbeatServices);
	final ResourceManagerGateway resourceManagerGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class);

	registerTaskExecutor(resourceManagerGateway, taskManagerId, taskExecutorGateway.getAddress());

	CompletableFuture<TaskManagerInfo> taskManagerInfoFuture = resourceManagerGateway.requestTaskManagerInfo(
		taskManagerId,
		TestingUtils.TIMEOUT());

	TaskManagerInfo taskManagerInfo = taskManagerInfoFuture.get();

	Assert.assertEquals(taskManagerId, taskManagerInfo.getResourceId());
	Assert.assertEquals(hardwareDescription, taskManagerInfo.getHardwareDescription());
	Assert.assertEquals(taskExecutorGateway.getAddress(), taskManagerInfo.getAddress());
	Assert.assertEquals(dataPort, taskManagerInfo.getDataPort());
	Assert.assertEquals(0, taskManagerInfo.getNumberSlots());
	Assert.assertEquals(0, taskManagerInfo.getNumberAvailableSlots());
}
 

@Test
public void testHeartbeatTimeoutWithTaskExecutor() throws Exception {
	final ResourceID taskExecutorId = ResourceID.generate();
	final CompletableFuture<ResourceID> heartbeatRequestFuture = new CompletableFuture<>();
	final CompletableFuture<Exception> disconnectFuture = new CompletableFuture<>();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setDisconnectResourceManagerConsumer(disconnectFuture::complete)
		.setHeartbeatResourceManagerConsumer(heartbeatRequestFuture::complete)
		.createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	runHeartbeatTimeoutTest(
		resourceManagerGateway -> {
			registerTaskExecutor(resourceManagerGateway, taskExecutorId, taskExecutorGateway.getAddress());
		},
		resourceManagerResourceId -> {
			// might have been completed or not depending whether the timeout was triggered first
			final ResourceID optionalHeartbeatRequestOrigin = heartbeatRequestFuture.getNow(null);
			assertThat(optionalHeartbeatRequestOrigin, anyOf(is(resourceManagerResourceId), is(nullValue())));
			assertThat(disconnectFuture.get(), instanceOf(TimeoutException.class));
		}
	);
}
 

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

private void runTest(TaskExecutorSetup taskExecutorBuilderSetup, TestAction testAction) throws Exception {
	final ResourceManagerGateway resourceManagerGateway = createAndStartResourceManager();

	TestingTaskExecutorGatewayBuilder testingTaskExecutorGateway1Builder = new TestingTaskExecutorGatewayBuilder();
	taskExecutorBuilderSetup.accept(testingTaskExecutorGateway1Builder);
	final TaskExecutorGateway taskExecutorGateway1 = testingTaskExecutorGateway1Builder
		.setAddress(UUID.randomUUID().toString())
		.createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway1.getAddress(), taskExecutorGateway1);

	final TaskExecutorGateway taskExecutorGateway2 = new TestingTaskExecutorGatewayBuilder()
		.setAddress(UUID.randomUUID().toString())
		.createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway2.getAddress(), taskExecutorGateway2);

	final ResourceID taskManagerId1 = ResourceID.generate();
	final ResourceID taskManagerId2 = ResourceID.generate();
	registerTaskExecutor(resourceManagerGateway, taskManagerId1, taskExecutorGateway1.getAddress());
	registerTaskExecutor(resourceManagerGateway, taskManagerId2, taskExecutorGateway2.getAddress());

	testAction.accept(resourceManagerGateway, taskManagerId1, taskManagerId2);
}
 

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

@Test
public void testHeartbeatTimeoutWithTaskExecutor() throws Exception {
	final ResourceID taskExecutorId = ResourceID.generate();
	final CompletableFuture<ResourceID> heartbeatRequestFuture = new CompletableFuture<>();
	final CompletableFuture<Exception> disconnectFuture = new CompletableFuture<>();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setDisconnectResourceManagerConsumer(disconnectFuture::complete)
		.setHeartbeatResourceManagerConsumer(heartbeatRequestFuture::complete)
		.createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	runHeartbeatTimeoutTest(
		resourceManagerGateway -> {
			registerTaskExecutor(resourceManagerGateway, taskExecutorId, taskExecutorGateway.getAddress());
		},
		resourceManagerResourceId -> {
			// might have been completed or not depending whether the timeout was triggered first
			final ResourceID optionalHeartbeatRequestOrigin = heartbeatRequestFuture.getNow(null);
			assertThat(optionalHeartbeatRequestOrigin, anyOf(is(resourceManagerResourceId), is(nullValue())));
			assertThat(disconnectFuture.get(), instanceOf(TimeoutException.class));
		}
	);
}
 

/**
 * Tests that we can retrieve the correct {@link TaskManagerInfo} from the {@link ResourceManager}.
 */
@Test
public void testRequestTaskManagerInfo() throws Exception {
	final ResourceID taskManagerId = ResourceID.generate();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setAddress(UUID.randomUUID().toString()).createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	resourceManager = createAndStartResourceManager(heartbeatServices);
	final ResourceManagerGateway resourceManagerGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class);

	registerTaskExecutor(resourceManagerGateway, taskManagerId, taskExecutorGateway.getAddress());

	CompletableFuture<TaskManagerInfo> taskManagerInfoFuture = resourceManagerGateway.requestTaskManagerInfo(
		taskManagerId,
		TestingUtils.TIMEOUT());

	TaskManagerInfo taskManagerInfo = taskManagerInfoFuture.get();

	assertEquals(taskManagerId, taskManagerInfo.getResourceId());
	assertEquals(hardwareDescription, taskManagerInfo.getHardwareDescription());
	assertEquals(taskExecutorGateway.getAddress(), taskManagerInfo.getAddress());
	assertEquals(dataPort, taskManagerInfo.getDataPort());
	assertEquals(0, taskManagerInfo.getNumberSlots());
	assertEquals(0, taskManagerInfo.getNumberAvailableSlots());
}
 

/**
 * Tests that we can retrieve the correct {@link TaskManagerInfo} from the {@link ResourceManager}.
 */
@Test
public void testRequestTaskManagerInfo() throws Exception {
	final ResourceID taskManagerId = ResourceID.generate();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setAddress(UUID.randomUUID().toString()).createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	resourceManager = createAndStartResourceManager(heartbeatServices);
	final ResourceManagerGateway resourceManagerGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class);

	registerTaskExecutor(resourceManagerGateway, taskManagerId, taskExecutorGateway.getAddress());

	CompletableFuture<TaskManagerInfo> taskManagerInfoFuture = resourceManagerGateway.requestTaskManagerInfo(
		taskManagerId,
		TestingUtils.TIMEOUT());

	TaskManagerInfo taskManagerInfo = taskManagerInfoFuture.get();

	assertEquals(taskManagerId, taskManagerInfo.getResourceId());
	assertEquals(hardwareDescription, taskManagerInfo.getHardwareDescription());
	assertEquals(taskExecutorGateway.getAddress(), taskManagerInfo.getAddress());
	assertEquals(dataPort, taskManagerInfo.getDataPort());
	assertEquals(0, taskManagerInfo.getNumberSlots());
	assertEquals(0, taskManagerInfo.getNumberAvailableSlots());
}
 

@Test
public void testHeartbeatTimeoutWithTaskExecutor() throws Exception {
	final ResourceID taskExecutorId = ResourceID.generate();
	final CompletableFuture<ResourceID> heartbeatRequestFuture = new CompletableFuture<>();
	final CompletableFuture<Exception> disconnectFuture = new CompletableFuture<>();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setDisconnectResourceManagerConsumer(disconnectFuture::complete)
		.setHeartbeatResourceManagerConsumer(heartbeatRequestFuture::complete)
		.createTestingTaskExecutorGateway();
	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	runHeartbeatTimeoutTest(
		resourceManagerGateway -> {
			registerTaskExecutor(resourceManagerGateway, taskExecutorId, taskExecutorGateway.getAddress());
		},
		resourceManagerResourceId -> {
			// might have been completed or not depending whether the timeout was triggered first
			final ResourceID optionalHeartbeatRequestOrigin = heartbeatRequestFuture.getNow(null);
			assertThat(optionalHeartbeatRequestOrigin, anyOf(is(resourceManagerResourceId), is(nullValue())));
			assertThat(disconnectFuture.get(), instanceOf(TimeoutException.class));
		}
	);
}
 

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

void registerTaskExecutor(ResourceID resourceID) throws Exception {
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.createTestingTaskExecutorGateway();
	((TestingRpcService) resourceManager.getRpcService()).registerGateway(resourceID.toString(), taskExecutorGateway);

	final ResourceManagerGateway rmGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class);

	final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(resourceID, 1), registerSlotProfile));

	final int numSlotsBeforeRegistering = CompletableFuture.supplyAsync(
		() -> slotManager.getNumberRegisteredSlots(),
		resourceManager.getMainThreadExecutorForTesting()).get();

	TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
		resourceID.toString(),
		resourceID,
		1234,
		new HardwareDescription(1, 2L, 3L, 4L),
		registerSlotProfile,
		registerSlotProfile);
	CompletableFuture<Integer> numberRegisteredSlotsFuture = rmGateway
		.registerTaskExecutor(
			taskExecutorRegistration,
			TIMEOUT)
		.thenCompose(
			(RegistrationResponse response) -> {
				assertThat(response, instanceOf(TaskExecutorRegistrationSuccess.class));
				final TaskExecutorRegistrationSuccess success = (TaskExecutorRegistrationSuccess) response;
				return rmGateway.sendSlotReport(
					resourceID,
					success.getRegistrationId(),
					slotReport,
					TIMEOUT);
			})
		.handleAsync(
			(Acknowledge ignored, Throwable throwable) -> slotManager.getNumberRegisteredSlots() - numSlotsBeforeRegistering,
			resourceManager.getMainThreadExecutorForTesting());
	Assert.assertEquals(1, numberRegisteredSlotsFuture.get().intValue());
}
 

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

/**
 * 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 CompletableFuture<InstanceID> releaseResourceFuture = new CompletableFuture<>();
	final ResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceId, ignored) -> releaseResourceFuture.complete(instanceId))
		.build();
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();

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

	try (final SlotManager slotManager = createSlotManagerBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.buildAndStartWithDirectExec(resourceManagerId, resourceActions)) {

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		assertThat(releaseResourceFuture.get(), is(taskExecutorConnection.getInstanceID()));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID(), TEST_EXCEPTION);

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

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

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

/**
 * 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 = new TestingResourceActionsBuilder().build();
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

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

	final ResourceProfile resourceProfile = ResourceProfile.fromResources(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);
	}
}
 

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

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

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

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 pending request is removed if task executor reports a slot with its allocation id.
 */
@Test
public void testSlotRequestRemovedIfTMReportAllocation() throws Exception {
	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(),
			new TestingResourceActionsBuilder().build())) {

		final JobID jobID = new JobID();
		final SlotRequest slotRequest1 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest1);

		final BlockingQueue<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(1);

		final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
				.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
					requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple6);
					try {
						return responseQueue.take();
					} catch (InterruptedException ignored) {
						return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
					}
				})
				.createTestingTaskExecutorGateway();

		final ResourceID taskExecutorResourceId = ResourceID.generate();
		final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
		final SlotReport slotReport = new SlotReport(createEmptySlotStatus(new SlotID(taskExecutorResourceId, 0), ResourceProfile.ANY));

		final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(firstManualSlotRequestResponse);

		slotManager.registerTaskManager(taskExecutionConnection, slotReport);

		final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> firstRequest = requestSlotQueue.take();

		final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(secondManualSlotRequestResponse);

		final SlotRequest slotRequest2 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest2);

		// fail first request
		firstManualSlotRequestResponse.completeExceptionally(new TimeoutException("Test exception to fail first allocation"));

		final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> secondRequest = requestSlotQueue.take();

		// fail second request
		secondManualSlotRequestResponse.completeExceptionally(new SlotOccupiedException("Test exception", slotRequest1.getAllocationId(), jobID));

		assertThat(firstRequest.f2, equalTo(slotRequest1.getAllocationId()));
		assertThat(secondRequest.f2, equalTo(slotRequest2.getAllocationId()));
		assertThat(secondRequest.f0, equalTo(firstRequest.f0));

		secondManualSlotRequestResponse.complete(Acknowledge.get());

		final TaskManagerSlot slot = slotManager.getSlot(secondRequest.f0);
		assertThat(slot.getState(), equalTo(TaskManagerSlot.State.ALLOCATED));
		assertThat(slot.getAllocationId(), equalTo(firstRequest.f2));

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
	}
}
 

/**
 * Tests that free slots which are reported as allocated won't be considered for fulfilling
 * other pending slot requests.
 *
 * <p>See: FLINK-8505
 */
@Test
public void testReportAllocatedSlot() throws Exception {
	final ResourceID taskManagerId = ResourceID.generate();
	final ResourceActions resourceActions = new TestingResourceActionsBuilder().build();
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(taskManagerId, taskExecutorGateway);

	try (final SlotManagerImpl slotManager = SlotManagerBuilder.newBuilder().build()) {

		slotManager.start(ResourceManagerId.generate(), Executors.directExecutor(), resourceActions);

		// initially report a single slot as free
		final SlotID slotId = new SlotID(taskManagerId, 0);
		final SlotStatus initialSlotStatus = new SlotStatus(
			slotId,
			ResourceProfile.UNKNOWN);
		final SlotReport initialSlotReport = new SlotReport(initialSlotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(equalTo(1)));

		// Now report this slot as allocated
		final SlotStatus slotStatus = new SlotStatus(
			slotId,
			ResourceProfile.UNKNOWN,
			new JobID(),
			new AllocationID());
		final SlotReport slotReport = new SlotReport(
			slotStatus);

		slotManager.reportSlotStatus(
			taskExecutorConnection.getInstanceID(),
			slotReport);

		// this slot request should not be fulfilled
		final AllocationID allocationId = new AllocationID();
		final SlotRequest slotRequest = new SlotRequest(
			new JobID(),
			allocationId,
			ResourceProfile.UNKNOWN,
			"foobar");

		// This triggered an IllegalStateException before
		slotManager.registerSlotRequest(slotRequest);

		assertThat(slotManager.getSlotRequest(allocationId).isAssigned(), is(false));
	}
}
 

/**
 * Tests notify the job manager of the allocations when the task manager is failed/killed.
 */
@Test
public void testNotifyFailedAllocationWhenTaskManagerTerminated() throws Exception {

	final Queue<Tuple2<JobID, AllocationID>> allocationFailures = new ArrayDeque<>(5);

	final TestingResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setNotifyAllocationFailureConsumer(
			(Tuple3<JobID, AllocationID, Exception> failureMessage) ->
				allocationFailures.offer(Tuple2.of(failureMessage.f0, failureMessage.f1)))
		.build();

	try (final SlotManager slotManager = createSlotManager(
		ResourceManagerId.generate(),
		resourceManagerActions)) {

		// register slot request for job1.
		JobID jobId1 = new JobID();
		final SlotRequest slotRequest11 = createSlotRequest(jobId1);
		final SlotRequest slotRequest12 = createSlotRequest(jobId1);
		slotManager.registerSlotRequest(slotRequest11);
		slotManager.registerSlotRequest(slotRequest12);

		// create task-manager-1 with 2 slots.
		final ResourceID taskExecutorResourceId1 = ResourceID.generate();
		final TestingTaskExecutorGateway testingTaskExecutorGateway1 = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
		final TaskExecutorConnection taskExecutionConnection1 = new TaskExecutorConnection(taskExecutorResourceId1, testingTaskExecutorGateway1);
		final SlotReport slotReport1 = createSlotReport(taskExecutorResourceId1, 2);

		// register the task-manager-1 to the slot manager, this will trigger the slot allocation for job1.
		slotManager.registerTaskManager(taskExecutionConnection1, slotReport1);

		// register slot request for job2.
		JobID jobId2 = new JobID();
		final SlotRequest slotRequest21 = createSlotRequest(jobId2);
		final SlotRequest slotRequest22 = createSlotRequest(jobId2);
		slotManager.registerSlotRequest(slotRequest21);
		slotManager.registerSlotRequest(slotRequest22);

		// register slot request for job3.
		JobID jobId3 = new JobID();
		final SlotRequest slotRequest31 = createSlotRequest(jobId3);
		slotManager.registerSlotRequest(slotRequest31);

		// create task-manager-2 with 3 slots.
		final ResourceID taskExecutorResourceId2 = ResourceID.generate();
		final TestingTaskExecutorGateway testingTaskExecutorGateway2 = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
		final TaskExecutorConnection taskExecutionConnection2 = new TaskExecutorConnection(taskExecutorResourceId2, testingTaskExecutorGateway2);
		final SlotReport slotReport2 = createSlotReport(taskExecutorResourceId2, 3);

		// register the task-manager-2 to the slot manager, this will trigger the slot allocation for job2 and job3.
		slotManager.registerTaskManager(taskExecutionConnection2, slotReport2);

		// validate for job1.
		slotManager.unregisterTaskManager(taskExecutionConnection1.getInstanceID(), TEST_EXCEPTION);

		assertThat(allocationFailures, hasSize(2));

		Tuple2<JobID, AllocationID> allocationFailure;
		final Set<AllocationID> failedAllocations = new HashSet<>(2);

		while ((allocationFailure = allocationFailures.poll()) != null) {
			assertThat(allocationFailure.f0, equalTo(jobId1));
			failedAllocations.add(allocationFailure.f1);
		}

		assertThat(failedAllocations, containsInAnyOrder(slotRequest11.getAllocationId(), slotRequest12.getAllocationId()));

		// validate the result for job2 and job3.
		slotManager.unregisterTaskManager(taskExecutionConnection2.getInstanceID(), TEST_EXCEPTION);

		assertThat(allocationFailures, hasSize(3));

		Map<JobID, List<Tuple2<JobID, AllocationID>>> job2AndJob3FailedAllocationInfo = allocationFailures.stream().collect(Collectors.groupingBy(tuple -> tuple.f0));

		assertThat(job2AndJob3FailedAllocationInfo.entrySet(), hasSize(2));

		final Set<AllocationID> job2FailedAllocations = extractFailedAllocationsForJob(jobId2, job2AndJob3FailedAllocationInfo);
		final Set<AllocationID> job3FailedAllocations = extractFailedAllocationsForJob(jobId3, job2AndJob3FailedAllocationInfo);

		assertThat(job2FailedAllocations, containsInAnyOrder(slotRequest21.getAllocationId(), slotRequest22.getAllocationId()));
		assertThat(job3FailedAllocations, containsInAnyOrder(slotRequest31.getAllocationId()));
	}
}
 

@Test
public void testHeartbeatTimeoutWithTaskManager() throws Exception {
	final CompletableFuture<ResourceID> heartbeatResourceIdFuture = new CompletableFuture<>();
	final CompletableFuture<JobID> disconnectedJobManagerFuture = new CompletableFuture<>();
	final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setHeartbeatJobManagerConsumer((taskManagerId, ignored) -> heartbeatResourceIdFuture.complete(taskManagerId))
		.setDisconnectJobManagerConsumer((jobId, throwable) -> disconnectedJobManagerFuture.complete(jobId))
		.createTestingTaskExecutorGateway();

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

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

	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		jobManagerSharedServices);

	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 JobID disconnectedJobManager = disconnectedJobManagerFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		assertThat(disconnectedJobManager, Matchers.equalTo(jobGraph.getJobID()));

		final ResourceID heartbeatResourceId = heartbeatResourceIdFuture.getNow(null);

		assertThat(heartbeatResourceId, anyOf(nullValue(), equalTo(jmResourceId)));
	} finally {
		jobManagerSharedServices.shutdown();
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

private TaskExecutorConnection createTaskExecutorConnection() {
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	return createTaskExecutorConnection(taskExecutorGateway);
}
 

/**
 * Tests that a batch job can be executed with fewer slots than its parallelism.
 * See FLINK-13187 for more information.
 */
@Test
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception {
	final int parallelism = 5;
	final Time batchSlotTimeout = Time.milliseconds(5L);
	final JobGraph jobGraph = createJobGraph(parallelism);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST);

	try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) {
		final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism);
		TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setSubmitTaskConsumer(
				(tdd, ignored) -> {
					submittedTasksQueue.offer(tdd.getExecutionAttemptId());
					return CompletableFuture.completedFuture(Acknowledge.get());
				})
			.createTestingTaskExecutorGateway();

		// register a single slot at the slot pool
		SlotPoolUtils.offerSlots(
			slotPool,
			mainThreadExecutor,
			Collections.singletonList(ResourceProfile.ANY),
			new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate()));

		final LegacyScheduler legacyScheduler = createLegacyScheduler(jobGraph, slotPool, mainThreadExecutor, batchSlotTimeout);

		final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener();
		legacyScheduler.registerJobStatusListener(jobStatusListener);
		startScheduling(legacyScheduler, mainThreadExecutor);

		// wait until the batch slot timeout has been reached
		Thread.sleep(batchSlotTimeout.toMilliseconds());

		final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture();

		for (int i = 0; i < parallelism; i++) {
			final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take));

			// wait until one of them is completed
			CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join();

			if (submittedTaskFuture.isDone()) {
				finishExecution(submittedTaskFuture.get(), legacyScheduler, mainThreadExecutor);
			} else {
				fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get()));
			}
		}

		assertThat(terminationFuture.get(), is(JobStatus.FINISHED));
	}
}