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

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

@Override
public CompletableFuture<Collection<Tuple2<ResourceID, String>>> requestTaskManagerMetricQueryServicePaths(Time timeout) {
	final ArrayList<CompletableFuture<Optional<Tuple2<ResourceID, String>>>> metricQueryServicePathFutures = new ArrayList<>(taskExecutors.size());

	for (Map.Entry<ResourceID, WorkerRegistration<WorkerType>> workerRegistrationEntry : taskExecutors.entrySet()) {
		final ResourceID tmResourceId = workerRegistrationEntry.getKey();
		final WorkerRegistration<WorkerType> workerRegistration = workerRegistrationEntry.getValue();
		final TaskExecutorGateway taskExecutorGateway = workerRegistration.getTaskExecutorGateway();

		final CompletableFuture<Optional<Tuple2<ResourceID, String>>> metricQueryServicePathFuture = taskExecutorGateway
			.requestMetricQueryServiceAddress(timeout)
			.thenApply(optional -> optional.map(path -> Tuple2.of(tmResourceId, path)));

		metricQueryServicePathFutures.add(metricQueryServicePathFuture);
	}

	return FutureUtils.combineAll(metricQueryServicePathFutures).thenApply(
		collection -> collection
			.stream()
			.filter(Optional::isPresent)
			.map(Optional::get)
			.collect(Collectors.toList()));
}
 

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

private Collection<SlotOffer> registerSlotsAtJobMaster(
		int numberSlots,
		JobMasterGateway jobMasterGateway,
		TaskExecutorGateway taskExecutorGateway,
		TaskManagerLocation taskManagerLocation) throws ExecutionException, InterruptedException {
	final AllocationIdsResourceManagerGateway allocationIdsResourceManagerGateway = new AllocationIdsResourceManagerGateway();
	rpcService.registerGateway(allocationIdsResourceManagerGateway.getAddress(), allocationIdsResourceManagerGateway);
	notifyResourceManagerLeaderListeners(allocationIdsResourceManagerGateway);

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

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

	Collection<SlotOffer> slotOffers = IntStream
		.range(0, numberSlots)
		.mapToObj(
			index -> {
				final AllocationID allocationId = allocationIdsResourceManagerGateway.takeAllocationId();
				return new SlotOffer(allocationId, index, ResourceProfile.UNKNOWN);
			})
		.collect(Collectors.toList());

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

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

@Override
public CompletableFuture<Acknowledge> disconnectTaskManager(final ResourceID resourceID, final Exception cause) {
	log.debug("Disconnect TaskExecutor {} because: {}", resourceID, cause.getMessage());

	taskManagerHeartbeatManager.unmonitorTarget(resourceID);
	slotPool.releaseTaskManager(resourceID, cause);
	partitionTracker.stopTrackingPartitionsFor(resourceID);

	Tuple2<TaskManagerLocation, TaskExecutorGateway> taskManagerConnection = registeredTaskManagers.remove(resourceID);

	if (taskManagerConnection != null) {
		taskManagerConnection.f1.disconnectJobManager(jobGraph.getJobID(), cause);
	}

	return CompletableFuture.completedFuture(Acknowledge.get());
}
 

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

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

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

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

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

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

private Collection<SlotOffer> registerSlotsAtJobMaster(
		int numberSlots,
		JobMasterGateway jobMasterGateway,
		TaskExecutorGateway taskExecutorGateway,
		TaskManagerLocation taskManagerLocation) throws ExecutionException, InterruptedException {
	final AllocationIdsResourceManagerGateway allocationIdsResourceManagerGateway = new AllocationIdsResourceManagerGateway();
	rpcService.registerGateway(allocationIdsResourceManagerGateway.getAddress(), allocationIdsResourceManagerGateway);
	notifyResourceManagerLeaderListeners(allocationIdsResourceManagerGateway);

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

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

	Collection<SlotOffer> slotOffers = IntStream
		.range(0, numberSlots)
		.mapToObj(
			index -> {
				final AllocationID allocationId = allocationIdsResourceManagerGateway.takeAllocationId();
				return new SlotOffer(allocationId, index, ResourceProfile.UNKNOWN);
			})
		.collect(Collectors.toList());

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

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

@Override
public CompletableFuture<Acknowledge> disconnectTaskManager(final ResourceID resourceID, final Exception cause) {
	log.debug("Disconnect TaskExecutor {} because: {}", resourceID, cause.getMessage());

	taskManagerHeartbeatManager.unmonitorTarget(resourceID);
	slotPool.releaseTaskManager(resourceID, cause);
	partitionTracker.stopTrackingPartitionsFor(resourceID);

	Tuple2<TaskManagerLocation, TaskExecutorGateway> taskManagerConnection = registeredTaskManagers.remove(resourceID);

	if (taskManagerConnection != null) {
		taskManagerConnection.f1.disconnectJobManager(jobGraph.getJobID(), cause);
	}

	return CompletableFuture.completedFuture(Acknowledge.get());
}
 

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

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

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

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

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

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

@Override
public CompletableFuture<RegistrationResponse> registerTaskExecutor(
		final TaskExecutorRegistration taskExecutorRegistration,
		final Time timeout) {

	CompletableFuture<TaskExecutorGateway> taskExecutorGatewayFuture = getRpcService().connect(taskExecutorRegistration.getTaskExecutorAddress(), TaskExecutorGateway.class);
	taskExecutorGatewayFutures.put(taskExecutorRegistration.getResourceId(), taskExecutorGatewayFuture);

	return taskExecutorGatewayFuture.handleAsync(
		(TaskExecutorGateway taskExecutorGateway, Throwable throwable) -> {
			final ResourceID resourceId = taskExecutorRegistration.getResourceId();
			if (taskExecutorGatewayFuture == taskExecutorGatewayFutures.get(resourceId)) {
				taskExecutorGatewayFutures.remove(resourceId);
				if (throwable != null) {
					return new RegistrationResponse.Decline(throwable.getMessage());
				} else {
					return registerTaskExecutorInternal(taskExecutorGateway, taskExecutorRegistration);
				}
			} else {
				log.debug("Ignoring outdated TaskExecutorGateway connection for {}.", resourceId);
				return new RegistrationResponse.Decline("Decline outdated task executor registration.");
			}
		},
		getMainThreadExecutor());
}
 

@Override
public CompletableFuture<Collection<Tuple2<ResourceID, String>>> requestTaskManagerMetricQueryServiceAddresses(Time timeout) {
	final ArrayList<CompletableFuture<Optional<Tuple2<ResourceID, String>>>> metricQueryServiceAddressFutures = new ArrayList<>(taskExecutors.size());

	for (Map.Entry<ResourceID, WorkerRegistration<WorkerType>> workerRegistrationEntry : taskExecutors.entrySet()) {
		final ResourceID tmResourceId = workerRegistrationEntry.getKey();
		final WorkerRegistration<WorkerType> workerRegistration = workerRegistrationEntry.getValue();
		final TaskExecutorGateway taskExecutorGateway = workerRegistration.getTaskExecutorGateway();

		final CompletableFuture<Optional<Tuple2<ResourceID, String>>> metricQueryServiceAddressFuture = taskExecutorGateway
			.requestMetricQueryServiceAddress(timeout)
			.thenApply(o -> o.toOptional().map(address -> Tuple2.of(tmResourceId, address)));

		metricQueryServiceAddressFutures.add(metricQueryServiceAddressFuture);
	}

	return FutureUtils.combineAll(metricQueryServiceAddressFutures).thenApply(
		collection -> collection
			.stream()
			.filter(Optional::isPresent)
			.map(Optional::get)
			.collect(Collectors.toList()));
}
 

private Collection<SlotOffer> registerSlotsAtJobMaster(
		int numberSlots,
		JobMasterGateway jobMasterGateway,
		TaskExecutorGateway taskExecutorGateway,
		UnresolvedTaskManagerLocation unresolvedTaskManagerLocation) throws ExecutionException, InterruptedException {
	final AllocationIdsResourceManagerGateway allocationIdsResourceManagerGateway = new AllocationIdsResourceManagerGateway();
	rpcService.registerGateway(allocationIdsResourceManagerGateway.getAddress(), allocationIdsResourceManagerGateway);
	notifyResourceManagerLeaderListeners(allocationIdsResourceManagerGateway);

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

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

	Collection<SlotOffer> slotOffers = IntStream
		.range(0, numberSlots)
		.mapToObj(
			index -> {
				final AllocationID allocationId = allocationIdsResourceManagerGateway.takeAllocationId();
				return new SlotOffer(allocationId, index, ResourceProfile.ANY);
			})
		.collect(Collectors.toList());

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

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

@Override
public CompletableFuture<Collection<Tuple2<ResourceID, String>>> requestTaskManagerMetricQueryServiceAddresses(Time timeout) {
	final ArrayList<CompletableFuture<Optional<Tuple2<ResourceID, String>>>> metricQueryServiceAddressFutures = new ArrayList<>(taskExecutors.size());

	for (Map.Entry<ResourceID, WorkerRegistration<WorkerType>> workerRegistrationEntry : taskExecutors.entrySet()) {
		final ResourceID tmResourceId = workerRegistrationEntry.getKey();
		final WorkerRegistration<WorkerType> workerRegistration = workerRegistrationEntry.getValue();
		final TaskExecutorGateway taskExecutorGateway = workerRegistration.getTaskExecutorGateway();

		final CompletableFuture<Optional<Tuple2<ResourceID, String>>> metricQueryServiceAddressFuture = taskExecutorGateway
			.requestMetricQueryServiceAddress(timeout)
			.thenApply(o -> o.toOptional().map(address -> Tuple2.of(tmResourceId, address)));

		metricQueryServiceAddressFutures.add(metricQueryServiceAddressFuture);
	}

	return FutureUtils.combineAll(metricQueryServiceAddressFutures).thenApply(
		collection -> collection
			.stream()
			.filter(Optional::isPresent)
			.map(Optional::get)
			.collect(Collectors.toList()));
}
 

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 Collection<SlotOffer> registerSlotsAtJobMaster(
	int numberSlots,
	JobMasterGateway jobMasterGateway,
	TaskExecutorGateway taskExecutorGateway) throws ExecutionException, InterruptedException {
	return registerSlotsAtJobMaster(
		numberSlots,
		jobMasterGateway,
		taskExecutorGateway,
		new LocalTaskManagerLocation());
}
 

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

public WorkerRegistration(
		TaskExecutorGateway taskExecutorGateway,
		WorkerType worker,
		int dataPort,
		HardwareDescription hardwareDescription) {

	super(worker.getResourceID(), taskExecutorGateway);

	this.worker = Preconditions.checkNotNull(worker);
	this.dataPort = dataPort;
	this.hardwareDescription = Preconditions.checkNotNull(hardwareDescription);
}
 

public WorkerRegistration(
		TaskExecutorGateway taskExecutorGateway,
		WorkerType worker,
		int dataPort,
		HardwareDescription hardwareDescription) {

	super(worker.getResourceID(), taskExecutorGateway);

	this.worker = Preconditions.checkNotNull(worker);
	this.dataPort = dataPort;
	this.hardwareDescription = Preconditions.checkNotNull(hardwareDescription);
}
 

/**
 * 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 CompletableFuture<RegistrationResponse> registerTaskExecutor(
		final String taskExecutorAddress,
		final ResourceID taskExecutorResourceId,
		final int dataPort,
		final HardwareDescription hardwareDescription,
		final Time timeout) {

	CompletableFuture<TaskExecutorGateway> taskExecutorGatewayFuture = getRpcService().connect(taskExecutorAddress, TaskExecutorGateway.class);
	taskExecutorGatewayFutures.put(taskExecutorResourceId, taskExecutorGatewayFuture);

	return taskExecutorGatewayFuture.handleAsync(
		(TaskExecutorGateway taskExecutorGateway, Throwable throwable) -> {
			if (taskExecutorGatewayFuture == taskExecutorGatewayFutures.get(taskExecutorResourceId)) {
				taskExecutorGatewayFutures.remove(taskExecutorResourceId);
				if (throwable != null) {
					return new RegistrationResponse.Decline(throwable.getMessage());
				} else {
					return registerTaskExecutorInternal(
						taskExecutorGateway,
						taskExecutorAddress,
						taskExecutorResourceId,
						dataPort,
						hardwareDescription);
				}
			} else {
				log.info("Ignoring outdated TaskExecutorGateway connection.");
				return new RegistrationResponse.Decline("Decline outdated task executor registration.");
			}
		},
		getMainThreadExecutor());
}
 

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