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

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

/**
 * Tests that the {@link SlotReport} sent to the RM does not contain
 * out dated/stale information as slots are being requested from the
 * TM.
 *
 * <p>This is a probabilistic test case and needs to be executed
 * several times to produce a failure without the fix for FLINK-12865.
 */
@Test
public void testSlotReportDoesNotContainStaleInformation() throws Exception {
	final OneShotLatch receivedSlotRequest = new OneShotLatch();
	final CompletableFuture<Void> verifySlotReportFuture = new CompletableFuture<>();
	final OneShotLatch terminateSlotReportVerification = new OneShotLatch();
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	// Assertions for this test
	testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, slotReport) -> {
		try {
			final ArrayList<SlotStatus> slots = Lists.newArrayList(slotReport);
			assertThat(slots, hasSize(1));
			final SlotStatus slotStatus = slots.get(0);

			log.info("Received SlotStatus: {}", slotStatus);

			if (receivedSlotRequest.isTriggered()) {
				assertThat(slotStatus.getAllocationID(), is(notNullValue()));
			} else {
				assertThat(slotStatus.getAllocationID(), is(nullValue()));
			}
		} catch (AssertionError e) {
			verifySlotReportFuture.completeExceptionally(e);
		}

		if (terminateSlotReportVerification.isTriggered()) {
			verifySlotReportFuture.complete(null);
		}
	});
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();

	testingResourceManagerGateway.setSendSlotReportFunction(ignored -> {
		taskExecutorRegistrationFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

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

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(new AllocateSlotNotifyingTaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService, receivedSlotRequest))
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
	final ResourceID taskExecutorResourceId = taskManagerServices.getTaskManagerLocation().getResourceID();

	taskExecutor.start();

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

	final ScheduledExecutorService heartbeatExecutor = java.util.concurrent.Executors.newSingleThreadScheduledExecutor();

	try {
		taskExecutorRegistrationFuture.get();

		final OneShotLatch scheduleFirstHeartbeat = new OneShotLatch();
		final ResourceID resourceManagerResourceId = testingResourceManagerGateway.getOwnResourceId();
		final long heartbeatInterval = 5L;
		heartbeatExecutor.scheduleWithFixedDelay(
			() -> {
				scheduleFirstHeartbeat.trigger();
				taskExecutorGateway.heartbeatFromResourceManager(resourceManagerResourceId);
			},
			0L,
			heartbeatInterval,
			TimeUnit.MILLISECONDS);

		scheduleFirstHeartbeat.await();

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

		requestSlotFuture.get();

		terminateSlotReportVerification.trigger();

		verifySlotReportFuture.get();
	} finally {
		ExecutorUtils.gracefulShutdown(timeout.toMilliseconds(), TimeUnit.MILLISECONDS, heartbeatExecutor);
		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);
	}
}
 

/**
 * Tests that accepted slots go into state assigned and the others are returned to the resource
 * manager.
 */
@Test
public void testSlotAcceptance() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Arrays.asList(mock(ResourceProfile.class), mock(ResourceProfile.class)), timerService);
	final JobManagerTable jobManagerTable = new JobManagerTable();
	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final String resourceManagerAddress = "rm";
	final UUID resourceManagerLeaderId = UUID.randomUUID();

	final String jobManagerAddress = "jm";
	final UUID jobManagerLeaderId = UUID.randomUUID();

	resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, resourceManagerLeaderId);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobManagerLeaderId);

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	final ResourceID resourceManagerResourceId = resourceManagerGateway.getOwnResourceId();
	final InstanceID registrationId = new InstanceID();

	final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
	resourceManagerGateway.setRegisterTaskExecutorFunction(
		stringResourceIDIntegerHardwareDescriptionTuple4 -> {
			registrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, new ClusterInformation("localhost", 1234)));
		});

	final CompletableFuture<Tuple3<InstanceID, SlotID, AllocationID>> availableSlotFuture = new CompletableFuture<>();
	resourceManagerGateway.setNotifySlotAvailableConsumer(availableSlotFuture::complete);

	final ResourceID jmResourceId = new ResourceID(jobManagerAddress);

	final AllocationID allocationId1 = new AllocationID();
	final AllocationID allocationId2 = new AllocationID();

	final SlotOffer offer1 = new SlotOffer(allocationId1, 0, ResourceProfile.UNKNOWN);

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);

	when(jobMasterGateway.registerTaskManager(
			any(String.class),
			eq(taskManagerLocation),
			any(Time.class)
	)).thenReturn(CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jmResourceId)));
	when(jobMasterGateway.getHostname()).thenReturn(jobManagerAddress);

	when(jobMasterGateway.offerSlots(
			any(ResourceID.class), any(Collection.class), any(Time.class)))
		.thenReturn(CompletableFuture.completedFuture((Collection<SlotOffer>) Collections.singleton(offer1)));

	rpc.registerGateway(resourceManagerAddress, resourceManagerGateway);
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

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

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		assertThat(registrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));

		taskSlotTable.allocateSlot(0, jobId, allocationId1, Time.milliseconds(10000L));
		taskSlotTable.allocateSlot(1, jobId, allocationId2, Time.milliseconds(10000L));

		// we have to add the job after the TaskExecutor, because otherwise the service has not
		// been properly started.
		jobLeaderService.addJob(jobId, jobManagerAddress);

		final Tuple3<InstanceID, SlotID, AllocationID> instanceIDSlotIDAllocationIDTuple3 = availableSlotFuture.get();

		final Tuple3<InstanceID, SlotID, AllocationID> expectedResult = Tuple3.of(registrationId, new SlotID(taskManagerLocation.getResourceID(), 1), allocationId2);

		assertThat(instanceIDSlotIDAllocationIDTuple3, equalTo(expectedResult));

		assertTrue(taskSlotTable.tryMarkSlotActive(jobId, allocationId1));
		assertFalse(taskSlotTable.tryMarkSlotActive(jobId, allocationId2));
		assertTrue(taskSlotTable.isSlotFree(1));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that the {@link SlotReport} sent to the RM does not contain
 * out dated/stale information as slots are being requested from the
 * TM.
 *
 * <p>This is a probabilistic test case and needs to be executed
 * several times to produce a failure without the fix for FLINK-12865.
 */
@Test
public void testSlotReportDoesNotContainStaleInformation() throws Exception {
	final OneShotLatch receivedSlotRequest = new OneShotLatch();
	final CompletableFuture<Void> verifySlotReportFuture = new CompletableFuture<>();
	final OneShotLatch terminateSlotReportVerification = new OneShotLatch();
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	// Assertions for this test
	testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, slotReport) -> {
		try {
			final ArrayList<SlotStatus> slots = Lists.newArrayList(slotReport);
			assertThat(slots, hasSize(1));
			final SlotStatus slotStatus = slots.get(0);

			log.info("Received SlotStatus: {}", slotStatus);

			if (receivedSlotRequest.isTriggered()) {
				assertThat(slotStatus.getAllocationID(), is(notNullValue()));
			} else {
				assertThat(slotStatus.getAllocationID(), is(nullValue()));
			}
		} catch (AssertionError e) {
			verifySlotReportFuture.completeExceptionally(e);
		}

		if (terminateSlotReportVerification.isTriggered()) {
			verifySlotReportFuture.complete(null);
		}
	});
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();

	testingResourceManagerGateway.setSendSlotReportFunction(ignored -> {
		taskExecutorRegistrationFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

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

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(new AllocateSlotNotifyingTaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService, receivedSlotRequest))
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
	final ResourceID taskExecutorResourceId = taskManagerServices.getTaskManagerLocation().getResourceID();

	taskExecutor.start();

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

	final ScheduledExecutorService heartbeatExecutor = java.util.concurrent.Executors.newSingleThreadScheduledExecutor();

	try {
		taskExecutorRegistrationFuture.get();

		final OneShotLatch scheduleFirstHeartbeat = new OneShotLatch();
		final ResourceID resourceManagerResourceId = testingResourceManagerGateway.getOwnResourceId();
		final long heartbeatInterval = 5L;
		heartbeatExecutor.scheduleWithFixedDelay(
			() -> {
				scheduleFirstHeartbeat.trigger();
				taskExecutorGateway.heartbeatFromResourceManager(resourceManagerResourceId);
			},
			0L,
			heartbeatInterval,
			TimeUnit.MILLISECONDS);

		scheduleFirstHeartbeat.await();

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

		requestSlotFuture.get();

		terminateSlotReportVerification.trigger();

		verifySlotReportFuture.get();
	} finally {
		ExecutorUtils.gracefulShutdown(timeout.toMilliseconds(), TimeUnit.MILLISECONDS, heartbeatExecutor);
		RpcUtils.terminateRpcEndpoint(taskExecutor, 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);
	}
}
 

/**
 * Tests that the {@link SlotReport} sent to the RM does not contain
 * out dated/stale information as slots are being requested from the
 * TM.
 *
 * <p>This is a probabilistic test case and needs to be executed
 * several times to produce a failure without the fix for FLINK-12865.
 */
@Test
public void testSlotReportDoesNotContainStaleInformation() throws Exception {
	final OneShotLatch receivedSlotRequest = new OneShotLatch();
	final CompletableFuture<Void> verifySlotReportFuture = new CompletableFuture<>();
	final OneShotLatch terminateSlotReportVerification = new OneShotLatch();
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	// Assertions for this test
	testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, heartbeatPayload) -> {
		try {
			final ArrayList<SlotStatus> slots = Lists.newArrayList(heartbeatPayload.getSlotReport());
			assertThat(slots, hasSize(1));
			final SlotStatus slotStatus = slots.get(0);

			log.info("Received SlotStatus: {}", slotStatus);

			if (receivedSlotRequest.isTriggered()) {
				assertThat(slotStatus.getAllocationID(), is(notNullValue()));
			} else {
				assertThat(slotStatus.getAllocationID(), is(nullValue()));
			}
		} catch (AssertionError e) {
			verifySlotReportFuture.completeExceptionally(e);
		}

		if (terminateSlotReportVerification.isTriggered()) {
			verifySlotReportFuture.complete(null);
		}
	});
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();

	testingResourceManagerGateway.setSendSlotReportFunction(ignored -> {
		taskExecutorRegistrationFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

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

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(new AllocateSlotNotifyingTaskSlotTable(receivedSlotRequest))
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
	final ResourceID taskExecutorResourceId = taskManagerServices.getUnresolvedTaskManagerLocation().getResourceID();

	taskExecutor.start();

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

	final ScheduledExecutorService heartbeatExecutor = java.util.concurrent.Executors.newSingleThreadScheduledExecutor();

	try {
		taskExecutorRegistrationFuture.get();

		final OneShotLatch scheduleFirstHeartbeat = new OneShotLatch();
		final ResourceID resourceManagerResourceId = testingResourceManagerGateway.getOwnResourceId();
		final long heartbeatInterval = 5L;
		heartbeatExecutor.scheduleWithFixedDelay(
			() -> {
				scheduleFirstHeartbeat.trigger();
				taskExecutorGateway.heartbeatFromResourceManager(resourceManagerResourceId);
			},
			0L,
			heartbeatInterval,
			TimeUnit.MILLISECONDS);

		scheduleFirstHeartbeat.await();

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

		requestSlotFuture.get();

		terminateSlotReportVerification.trigger();

		verifySlotReportFuture.get();
	} finally {
		ExecutorUtils.gracefulShutdown(timeout.toMilliseconds(), TimeUnit.MILLISECONDS, heartbeatExecutor);
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}