Java Code Examples for org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway#setRegisterTaskExecutorFunction()

private TestingResourceManagerGateway createRmWithTmRegisterAndNotifySlotHooks(
		InstanceID registrationId,
		OneShotLatch taskExecutorIsRegistered,
		CompletableFuture<Tuple3<InstanceID, SlotID, AllocationID>> availableSlotFuture) {
	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	resourceManagerLeaderRetriever.notifyListener(
		resourceManagerGateway.getAddress(),
		resourceManagerGateway.getFencingToken().toUUID());

	resourceManagerGateway.setRegisterTaskExecutorFunction(
		taskExecutorRegistration -> CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(registrationId, resourceManagerGateway.getOwnResourceId(),
				new ClusterInformation("localhost", 1234))));

	resourceManagerGateway.setNotifySlotAvailableConsumer(availableSlotFuture::complete);

	resourceManagerGateway.setSendSlotReportFunction(ignored -> {
		taskExecutorIsRegistered.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());
	});
	return resourceManagerGateway;
}
 

/**
 * Tests that we ignore slot requests if the TaskExecutor is not
 * registered at a ResourceManager.
 */
@Test
public void testIgnoringSlotRequestsIfNotRegistered() throws Exception {
	final TaskExecutor taskExecutor = createTaskExecutor(1);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final CompletableFuture<RegistrationResponse> registrationFuture = new CompletableFuture<>();
		final CompletableFuture<ResourceID> taskExecutorResourceIdFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
			taskExecutorResourceIdFuture.complete(taskExecutorRegistration.getResourceId());
			return registrationFuture;
		});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		final ResourceID resourceId = taskExecutorResourceIdFuture.get();

		final SlotID slotId = new SlotID(resourceId, 0);
		final CompletableFuture<Acknowledge> slotRequestResponse = taskExecutorGateway.requestSlot(slotId, jobId, new AllocationID(), ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		try {
			slotRequestResponse.get();
			fail("We should not be able to request slots before the TaskExecutor is registered at the ResourceManager.");
		} catch (ExecutionException ee) {
			assertThat(ExceptionUtils.stripExecutionException(ee), instanceOf(TaskManagerException.class));
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Test
public void testImmediatelyRegistersIfLeaderIsKnown() throws Exception {
	final String resourceManagerAddress = "/resource/manager/address/one";

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CountDownLatch taskManagerRegisteredLatch = new CountDownLatch(1);
	testingResourceManagerGateway.setRegisterTaskExecutorFunction(FunctionUtils.uncheckedFunction(
		ignored -> {
			taskManagerRegisteredLatch.countDown();
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
				new InstanceID(), new ResourceID(resourceManagerAddress), new ClusterInformation("localhost", 1234)));
		}
	));

	rpc.registerGateway(resourceManagerAddress, testingResourceManagerGateway);

	final TaskSlotTable taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, UUID.randomUUID());

		assertTrue(taskManagerRegisteredLatch.await(timeout.toMilliseconds(), TimeUnit.MILLISECONDS));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

@Test
public void testImmediatelyRegistersIfLeaderIsKnown() throws Exception {
	final String resourceManagerAddress = "/resource/manager/address/one";

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CountDownLatch taskManagerRegisteredLatch = new CountDownLatch(1);
	testingResourceManagerGateway.setRegisterTaskExecutorFunction(FunctionUtils.uncheckedFunction(
		ignored -> {
			taskManagerRegisteredLatch.countDown();
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
				new InstanceID(), new ResourceID(resourceManagerAddress), new ClusterInformation("localhost", 1234)));
		}
	));

	rpc.registerGateway(resourceManagerAddress, testingResourceManagerGateway);

	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, UUID.randomUUID());

		assertTrue(taskManagerRegisteredLatch.await(timeout.toMilliseconds(), TimeUnit.MILLISECONDS));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

@Test
public void testRegisterWithDefaultSlotResourceProfile() throws Exception {
	final int numberOfSlots = 2;
	final TaskExecutor taskExecutor = createTaskExecutor(numberOfSlots);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final CompletableFuture<ResourceProfile> registeredDefaultSlotResourceProfileFuture = new CompletableFuture<>();

		final ResourceID ownResourceId = testingResourceManagerGateway.getOwnResourceId();
		testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
			registeredDefaultSlotResourceProfileFuture.complete(taskExecutorRegistration.getDefaultSlotResourceProfile());
			return CompletableFuture.completedFuture(
				new TaskExecutorRegistrationSuccess(
					new InstanceID(),
					ownResourceId,
					new ClusterInformation("localhost", 1234)));
		});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		assertThat(registeredDefaultSlotResourceProfileFuture.get(),
			equalTo(TaskExecutorResourceUtils.generateDefaultSlotResourceProfile(TM_RESOURCE_SPEC, numberOfSlots)));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report
 * fails.
 */
@Test
public void testInitialSlotReportFailure() throws Exception {
	final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
	final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				try {
					return responseQueue.take();
				} catch (InterruptedException e) {
					return FutureUtils.completedExceptionally(e);
				}
			});

		final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				testingResourceManagerGateway.getOwnResourceId(),
				new ClusterInformation("foobar", 1234)));

		final CountDownLatch numberRegistrations = new CountDownLatch(2);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
				numberRegistrations.countDown();
				return registrationResponse;
		});

		responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
		responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		//wait for the second registration attempt
		numberRegistrations.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Test
public void testTriggerRegistrationOnLeaderChange() throws Exception {
	final UUID leaderId1 = UUID.randomUUID();
	final UUID leaderId2 = UUID.randomUUID();

	// register the mock resource manager gateways
	final CompletableFuture<TaskExecutorRegistration> rmGateway1TaskExecutorRegistration = new CompletableFuture<>();
	TestingResourceManagerGateway rmGateway1 = new TestingResourceManagerGateway();
	rmGateway1.setRegisterTaskExecutorFunction(
		taskExecutorRegistration -> {
			rmGateway1TaskExecutorRegistration.complete(taskExecutorRegistration);
			return createRegistrationResponse(rmGateway1);
		});

	final CompletableFuture<TaskExecutorRegistration> rmGateway2TaskExecutorRegistration = new CompletableFuture<>();
	TestingResourceManagerGateway rmGateway2 = new TestingResourceManagerGateway();
	rmGateway2.setRegisterTaskExecutorFunction(
		taskExecutorRegistration -> {
			rmGateway2TaskExecutorRegistration.complete(taskExecutorRegistration);
			return createRegistrationResponse(rmGateway2);
		});

	rpc.registerGateway(rmGateway1.getAddress(), rmGateway1);
	rpc.registerGateway(rmGateway2.getAddress(), rmGateway2);

	final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable
		.<Task>newBuilder()
		.createSlotReportSupplier(SlotReport::new)
		.closeAsyncReturns(CompletableFuture.completedFuture(null))
		.build();

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
		String taskManagerAddress = taskManager.getAddress();

		// no connection initially, since there is no leader
		assertNull(taskManager.getResourceManagerConnection());

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmGateway1.getAddress(), leaderId1);
		final TaskExecutorRegistration taskExecutorRegistration1 = rmGateway1TaskExecutorRegistration.join();
		assertThat(taskExecutorRegistration1.getTaskExecutorAddress(), is(taskManagerAddress));
		assertThat(taskExecutorRegistration1.getResourceId(), is(unresolvedTaskManagerLocation.getResourceID()));
		assertNotNull(taskManager.getResourceManagerConnection());

		// cancel the leader
		resourceManagerLeaderRetriever.notifyListener(null, null);

		// set a new leader, see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmGateway2.getAddress(), leaderId2);

		final TaskExecutorRegistration taskExecutorRegistration2 = rmGateway2TaskExecutorRegistration.join();
		assertThat(taskExecutorRegistration2.getTaskExecutorAddress(), is(taskManagerAddress));
		assertThat(taskExecutorRegistration2.getResourceId(), is(unresolvedTaskManagerLocation.getResourceID()));
		assertNotNull(taskManager.getResourceManagerConnection());
	}
	finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that the correct partition/slot report is sent as part of the heartbeat response.
 */
@Test
public void testHeartbeatReporting() throws Exception {
	final String rmAddress = "rm";
	final UUID rmLeaderId = UUID.randomUUID();

	// register the mock resource manager gateway
	final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
	final ResourceID rmResourceId = rmGateway.getOwnResourceId();
	final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmResourceId,
			new ClusterInformation("localhost", 1234)));

	rmGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
		taskExecutorRegistrationFuture.complete(taskExecutorRegistration.getResourceId());
		return registrationResponse;
	});

	final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
	rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<TaskExecutorHeartbeatPayload> heartbeatPayloadCompletableFuture = new CompletableFuture<>();
	rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, heartbeatPayload) -> heartbeatPayloadCompletableFuture.complete(heartbeatPayload));

	rpc.registerGateway(rmAddress, rmGateway);

	final SlotID slotId = new SlotID(unresolvedTaskManagerLocation.getResourceID(), 0);
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotReport slotReport1 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile));
	final SlotReport slotReport2 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile,
			new JobID(),
			new AllocationID()));

	final Queue<SlotReport> reports = new ArrayDeque<>(Arrays.asList(slotReport1, slotReport2));
	final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable
		.<Task>newBuilder()
		.createSlotReportSupplier(reports::poll)
		.closeAsyncReturns(CompletableFuture.completedFuture(null))
		.build();

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutorPartitionTracker partitionTracker = createPartitionTrackerWithFixedPartitionReport(taskManagerServices.getShuffleEnvironment());

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices, HEARTBEAT_SERVICES, partitionTracker);

	try {
		taskManager.start();

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);

		// register resource manager success will trigger monitoring heartbeat target between tm and rm
		assertThat(taskExecutorRegistrationFuture.get(), equalTo(unresolvedTaskManagerLocation.getResourceID()));
		assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));

		TaskExecutorGateway taskExecutorGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// trigger the heartbeat asynchronously
		taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);

		// wait for heartbeat response
		SlotReport actualSlotReport = heartbeatPayloadCompletableFuture.get().getSlotReport();

		// the new slot report should be reported
		assertEquals(slotReport2, actualSlotReport);

		ClusterPartitionReport actualClusterPartitionReport = heartbeatPayloadCompletableFuture.get().getClusterPartitionReport();
		assertEquals(partitionTracker.createClusterPartitionReport(), actualClusterPartitionReport);
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

@Test
public void testHeartbeatTimeoutWithResourceManager() throws Exception {
	final String rmAddress = "rm";
	final ResourceID rmResourceId = new ResourceID(rmAddress);

	final long heartbeatInterval = 1L;
	final long heartbeatTimeout = 3L;

	final ResourceManagerId rmLeaderId = ResourceManagerId.generate();

	TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway(
		rmLeaderId,
		rmResourceId,
		rmAddress,
		rmAddress);

	final TaskExecutorRegistrationSuccess registrationResponse = new TaskExecutorRegistrationSuccess(
		new InstanceID(),
		rmResourceId,
		new ClusterInformation("localhost", 1234));

	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
	final CountDownLatch registrationAttempts = new CountDownLatch(2);
	rmGateway.setRegisterTaskExecutorFunction(
		registration -> {
			taskExecutorRegistrationFuture.complete(registration.getResourceId());
			registrationAttempts.countDown();
			return CompletableFuture.completedFuture(registrationResponse);
		});

	final CompletableFuture<ResourceID> taskExecutorDisconnectFuture = new CompletableFuture<>();
	rmGateway.setDisconnectTaskExecutorConsumer(
		disconnectInfo -> taskExecutorDisconnectFuture.complete(disconnectInfo.f0));

	rpc.registerGateway(rmAddress, rmGateway);

	final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable
		.<Task>newBuilder()
		.createSlotReportSupplier(SlotReport::new)
		.closeAsyncReturns(CompletableFuture.completedFuture(null))
		.build();

	HeartbeatServices heartbeatServices = new HeartbeatServices(heartbeatInterval, heartbeatTimeout);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices, heartbeatServices);

	try {
		taskManager.start();

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId.toUUID());

		// register resource manager success will trigger monitoring heartbeat target between tm and rm
		assertThat(taskExecutorRegistrationFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS), equalTo(unresolvedTaskManagerLocation.getResourceID()));

		// heartbeat timeout should trigger disconnect TaskManager from ResourceManager
		assertThat(taskExecutorDisconnectFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS), equalTo(unresolvedTaskManagerLocation.getResourceID()));

		assertTrue(
			"The TaskExecutor should try to reconnect to the RM",
			registrationAttempts.await(timeout.toMilliseconds(), TimeUnit.SECONDS));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report
 * fails.
 */
@Test
public void testInitialSlotReportFailure() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				try {
					return responseQueue.take();
				} catch (InterruptedException e) {
					return FutureUtils.completedExceptionally(e);
				}
			});

		final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				testingResourceManagerGateway.getOwnResourceId(),
				new ClusterInformation("foobar", 1234)));

		final CountDownLatch numberRegistrations = new CountDownLatch(2);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(new Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>>() {
			@Override
			public CompletableFuture<RegistrationResponse> apply(Tuple4<String, ResourceID, Integer, HardwareDescription> stringResourceIDIntegerHardwareDescriptionTuple4) {
				numberRegistrations.countDown();
				return registrationResponse;
			}
		});

		responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
		responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		//wait for the second registration attempt
		numberRegistrations.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the TaskExecutor tries to reconnect to a ResourceManager from which it
 * was explicitly disconnected.
 */
@Test
public void testReconnectionAttemptIfExplicitlyDisconnected() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskExecutor taskExecutor = createTaskExecutor(new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build());

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final ClusterInformation clusterInformation = new ClusterInformation("foobar", 1234);
		final CompletableFuture<RegistrationResponse> registrationResponseFuture = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), ResourceID.generate(), clusterInformation));
		final BlockingQueue<ResourceID> registrationQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5 -> {
			registrationQueue.offer(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5.f1);
			return registrationResponseFuture;
		});
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		final ResourceID firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		assertThat(registrationQueue, is(empty()));

		taskExecutorGateway.disconnectResourceManager(new FlinkException("Test exception"));

		final ResourceID secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));

	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the TaskExecutor tries to reconnect to a ResourceManager from which it
 * was explicitly disconnected.
 */
@Test
public void testReconnectionAttemptIfExplicitlyDisconnected() throws Exception {
	final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
	final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	final TaskExecutor taskExecutor = createTaskExecutor(new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.build());

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final ClusterInformation clusterInformation = new ClusterInformation("foobar", 1234);
		final CompletableFuture<RegistrationResponse> registrationResponseFuture = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), ResourceID.generate(), clusterInformation));
		final BlockingQueue<ResourceID> registrationQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
			registrationQueue.offer(taskExecutorRegistration.getResourceId());
			return registrationResponseFuture;
		});
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		final ResourceID firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(unresolvedTaskManagerLocation.getResourceID()));

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		assertThat(registrationQueue, is(empty()));

		taskExecutorGateway.disconnectResourceManager(new FlinkException("Test exception"));

		final ResourceID secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(unresolvedTaskManagerLocation.getResourceID()));

	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Test
public void testMaximumRegistrationDurationAfterConnectionLoss() throws Exception {
	configuration.setString(TaskManagerOptions.REGISTRATION_TIMEOUT, "100 ms");
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices, new HeartbeatServices(10L, 10L));

	taskExecutor.start();

	final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
	final OneShotLatch secondRegistration = new OneShotLatch();
	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		testingResourceManagerGateway.setRegisterTaskExecutorFunction(
			tuple -> {
				if (registrationFuture.complete(tuple.f1)) {
					return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
						new InstanceID(),
						testingResourceManagerGateway.getOwnResourceId(),
						new ClusterInformation("localhost", 1234)));
				} else {
					secondRegistration.trigger();
					return CompletableFuture.completedFuture(new RegistrationResponse.Decline("Only the first registration should succeed."));
				}
			}
		);
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), UUID.randomUUID());

		final ResourceID registrationResourceId = registrationFuture.get();

		assertThat(registrationResourceId, equalTo(taskManagerServices.getTaskManagerLocation().getResourceID()));

		secondRegistration.await();

		final Throwable error = testingFatalErrorHandler.getErrorFuture().get();
		assertThat(error, is(notNullValue()));
		assertThat(ExceptionUtils.stripExecutionException(error), instanceOf(RegistrationTimeoutException.class));

		testingFatalErrorHandler.clearError();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the correct slot report is sent as part of the heartbeat response.
 */
@Test
public void testHeartbeatSlotReporting() throws Exception {
	final String rmAddress = "rm";
	final UUID rmLeaderId = UUID.randomUUID();

	// register the mock resource manager gateway
	final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
	final ResourceID rmResourceId = rmGateway.getOwnResourceId();
	final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmResourceId,
			new ClusterInformation("localhost", 1234)));

	rmGateway.setRegisterTaskExecutorFunction(stringResourceIDIntegerHardwareDescriptionTuple4 -> {
		taskExecutorRegistrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
		return registrationResponse;
	});

	final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
	rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<SlotReport> heartbeatSlotReportFuture = new CompletableFuture<>();
	rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, slotReport) -> heartbeatSlotReportFuture.complete(slotReport));

	rpc.registerGateway(rmAddress, rmGateway);

	final SlotID slotId = new SlotID(taskManagerLocation.getResourceID(), 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotReport slotReport1 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile));
	final SlotReport slotReport2 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile,
			new JobID(),
			new AllocationID()));

	final TestingTaskSlotTable taskSlotTable = new TestingTaskSlotTable(new ArrayDeque<>(Arrays.asList(slotReport1, slotReport2)));

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);

		// register resource manager success will trigger monitoring heartbeat target between tm and rm
		assertThat(taskExecutorRegistrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));
		assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));

		TaskExecutorGateway taskExecutorGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// trigger the heartbeat asynchronously
		taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);

		// wait for heartbeat response
		SlotReport actualSlotReport = heartbeatSlotReportFuture.get();

		// the new slot report should be reported
		assertEquals(slotReport2, actualSlotReport);
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

@Test
public void testMaximumRegistrationDurationAfterConnectionLoss() throws Exception {
	configuration.set(TaskManagerOptions.REGISTRATION_TIMEOUT, TimeUtils.parseDuration("100 ms"));
	final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices, new HeartbeatServices(10L, 10L));

	taskExecutor.start();

	final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
	final OneShotLatch secondRegistration = new OneShotLatch();
	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		testingResourceManagerGateway.setRegisterTaskExecutorFunction(
			taskExecutorRegistration -> {
				if (registrationFuture.complete(taskExecutorRegistration.getResourceId())) {
					return createRegistrationResponse(testingResourceManagerGateway);
				} else {
					secondRegistration.trigger();
					return CompletableFuture.completedFuture(new Decline("Only the first registration should succeed."));
				}
			}
		);
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), UUID.randomUUID());

		final ResourceID registrationResourceId = registrationFuture.get();

		assertThat(registrationResourceId, equalTo(taskManagerServices.getUnresolvedTaskManagerLocation().getResourceID()));

		secondRegistration.await();

		final Throwable error = testingFatalErrorHandler.getErrorFuture().get();
		assertThat(error, is(notNullValue()));
		assertThat(ExceptionUtils.stripExecutionException(error), instanceOf(RegistrationTimeoutException.class));

		testingFatalErrorHandler.clearError();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report
 * fails.
 */
@Test
public void testInitialSlotReportFailure() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				try {
					return responseQueue.take();
				} catch (InterruptedException e) {
					return FutureUtils.completedExceptionally(e);
				}
			});

		final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				testingResourceManagerGateway.getOwnResourceId(),
				new ClusterInformation("foobar", 1234)));

		final CountDownLatch numberRegistrations = new CountDownLatch(2);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(new Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>>() {
			@Override
			public CompletableFuture<RegistrationResponse> apply(Tuple4<String, ResourceID, Integer, HardwareDescription> stringResourceIDIntegerHardwareDescriptionTuple4) {
				numberRegistrations.countDown();
				return registrationResponse;
			}
		});

		responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
		responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		//wait for the second registration attempt
		numberRegistrations.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the TaskExecutor tries to reconnect to a ResourceManager from which it
 * was explicitly disconnected.
 */
@Test
public void testReconnectionAttemptIfExplicitlyDisconnected() throws Exception {
	final long heartbeatInterval = 1000L;
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskExecutor taskExecutor = new TaskExecutor(
		rpc,
		TaskManagerConfiguration.fromConfiguration(configuration),
		haServices,
		new TaskManagerServicesBuilder()
			.setTaskSlotTable(taskSlotTable)
			.setTaskManagerLocation(taskManagerLocation)
			.build(),
		new HeartbeatServices(heartbeatInterval, 1000L),
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final ClusterInformation clusterInformation = new ClusterInformation("foobar", 1234);
		final CompletableFuture<RegistrationResponse> registrationResponseFuture = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), ResourceID.generate(), clusterInformation));
		final BlockingQueue<ResourceID> registrationQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5 -> {
			registrationQueue.offer(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5.f1);
			return registrationResponseFuture;
		});
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		final ResourceID firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		assertThat(registrationQueue, is(empty()));

		taskExecutorGateway.disconnectResourceManager(new FlinkException("Test exception"));

		final ResourceID secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));

	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that we ignore slot requests if the TaskExecutor is not
 * registered at a ResourceManager.
 */
@Test
public void testIgnoringSlotRequestsIfNotRegistered() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final CompletableFuture<RegistrationResponse> registrationFuture = new CompletableFuture<>();
		final CompletableFuture<ResourceID> taskExecutorResourceIdFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5 -> {
			taskExecutorResourceIdFuture.complete(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5.f1);
			return registrationFuture;
		});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		final ResourceID resourceId = taskExecutorResourceIdFuture.get();

		final SlotID slotId = new SlotID(resourceId, 0);
		final CompletableFuture<Acknowledge> slotRequestResponse = taskExecutorGateway.requestSlot(slotId, jobId, new AllocationID(), "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		try {
			slotRequestResponse.get();
			fail("We should not be able to request slots before the TaskExecutor is registered at the ResourceManager.");
		} catch (ExecutionException ee) {
			assertThat(ExceptionUtils.stripExecutionException(ee), instanceOf(TaskManagerException.class));
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Test
public void testMaximumRegistrationDurationAfterConnectionLoss() throws Exception {
	configuration.setString(TaskManagerOptions.REGISTRATION_TIMEOUT, "100 ms");
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);

	final long heartbeatInterval = 10L;
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
	final TaskExecutor taskExecutor = new TaskExecutor(
		rpc,
		TaskManagerConfiguration.fromConfiguration(configuration),
		haServices,
		taskManagerServices,
		new HeartbeatServices(heartbeatInterval, 10L),
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	taskExecutor.start();

	final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
	final OneShotLatch secondRegistration = new OneShotLatch();
	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		testingResourceManagerGateway.setRegisterTaskExecutorFunction(
			tuple -> {
				if (registrationFuture.complete(tuple.f1)) {
					return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
						new InstanceID(),
						testingResourceManagerGateway.getOwnResourceId(),
						new ClusterInformation("localhost", 1234)));
				} else {
					secondRegistration.trigger();
					return CompletableFuture.completedFuture(new RegistrationResponse.Decline("Only the first registration should succeed."));
				}
			}
		);
		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);

		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), UUID.randomUUID());

		final ResourceID registrationResourceId = registrationFuture.get();

		assertThat(registrationResourceId, equalTo(taskManagerServices.getTaskManagerLocation().getResourceID()));

		secondRegistration.await();

		final Throwable error = testingFatalErrorHandler.getErrorFuture().get();
		assertThat(error, is(notNullValue()));
		assertThat(ExceptionUtils.stripExecutionException(error), instanceOf(RegistrationTimeoutException.class));

		testingFatalErrorHandler.clearError();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the correct slot report is sent as part of the heartbeat response.
 */
@Test
public void testHeartbeatSlotReporting() throws Exception {
	final String rmAddress = "rm";
	final UUID rmLeaderId = UUID.randomUUID();

	// register the mock resource manager gateway
	final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
	final ResourceID rmResourceId = rmGateway.getOwnResourceId();
	final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmResourceId,
			new ClusterInformation("localhost", 1234)));

	rmGateway.setRegisterTaskExecutorFunction(stringResourceIDIntegerHardwareDescriptionTuple4 -> {
		taskExecutorRegistrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
		return registrationResponse;
	});

	final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
	rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<SlotReport> heartbeatSlotReportFuture = new CompletableFuture<>();
	rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, slotReport) -> heartbeatSlotReportFuture.complete(slotReport));

	rpc.registerGateway(rmAddress, rmGateway);

	final SlotID slotId = new SlotID(taskManagerLocation.getResourceID(), 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotReport slotReport1 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile));
	final SlotReport slotReport2 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile,
			new JobID(),
			new AllocationID()));

	final TestingTaskSlotTable taskSlotTable = new TestingTaskSlotTable(new ArrayDeque<>(Arrays.asList(slotReport1, slotReport2)));

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = new TaskExecutor(
		rpc,
		taskManagerConfiguration,
		haServices,
		taskManagerServices,
		HEARTBEAT_SERVICES,
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	try {
		taskManager.start();

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);

		// register resource manager success will trigger monitoring heartbeat target between tm and rm
		assertThat(taskExecutorRegistrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));
		assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));

		TaskExecutorGateway taskExecutorGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// trigger the heartbeat asynchronously
		taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);

		// wait for heartbeat response
		SlotReport actualSlotReport = heartbeatSlotReportFuture.get();

		// the new slot report should be reported
		assertEquals(slotReport2, actualSlotReport);
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}