org.apache.flink.runtime.jobmaster.JMTMRegistrationSuccess Java Examples

public void reconnect() {
	if (stopped) {
		LOG.debug("Cannot reconnect because the JobManagerLeaderListener has already been stopped.");
	} else {
		final RegisteredRpcConnection<JobMasterId, JobMasterGateway, JMTMRegistrationSuccess> currentRpcConnection = rpcConnection;

		if (currentRpcConnection != null) {
			if (currentRpcConnection.isConnected()) {

				if (currentRpcConnection.tryReconnect()) {
					// double check for concurrent stop operation
					if (stopped) {
						currentRpcConnection.close();
					}
				} else {
					LOG.debug("Could not reconnect to the JobMaster {}.", currentRpcConnection.getTargetAddress());
				}
			} else {
				LOG.debug("Ongoing registration to JobMaster {}.", currentRpcConnection.getTargetAddress());
			}
		} else {
			LOG.debug("Cannot reconnect to an unknown JobMaster.");
		}
	}
}
 

public void reconnect() {
	if (stopped) {
		LOG.debug("Cannot reconnect because the JobManagerLeaderListener has already been stopped.");
	} else {
		final RegisteredRpcConnection<JobMasterId, JobMasterGateway, JMTMRegistrationSuccess> currentRpcConnection = rpcConnection;

		if (currentRpcConnection != null) {
			if (currentRpcConnection.isConnected()) {

				if (currentRpcConnection.tryReconnect()) {
					// double check for concurrent stop operation
					if (stopped) {
						currentRpcConnection.close();
					}
				} else {
					LOG.debug("Could not reconnect to the JobMaster {}.", currentRpcConnection.getTargetAddress());
				}
			} else {
				LOG.debug("Ongoing registration to JobMaster {}.", currentRpcConnection.getTargetAddress());
			}
		} else {
			LOG.debug("Cannot reconnect to an unknown JobMaster.");
		}
	}
}
 

@Override
public void jobManagerGainedLeadership(
	final JobID jobId,
	final JobMasterGateway jobManagerGateway,
	final JMTMRegistrationSuccess registrationMessage) {
	runAsync(
		() ->
			establishJobManagerConnection(
				jobId,
				jobManagerGateway,
				registrationMessage));
}
 

@Override
protected RetryingRegistration<JobMasterId, JobMasterGateway, JMTMRegistrationSuccess> generateRegistration() {
	return new JobLeaderService.JobManagerRetryingRegistration(
			LOG,
			rpcService,
			"JobManager",
			JobMasterGateway.class,
			getTargetAddress(),
			getTargetLeaderId(),
			retryingRegistrationConfiguration,
			ownerAddress,
			ownLocation);
}
 

@Override
protected void onRegistrationSuccess(JMTMRegistrationSuccess success) {
	// filter out old registration attempts
	if (Objects.equals(getTargetLeaderId(), currentJobMasterId)) {
		log.info("Successful registration at job manager {} for job {}.", getTargetAddress(), jobId);

		jobLeaderListener.jobManagerGainedLeadership(jobId, getTargetGateway(), success);
	} else {
		log.debug("Encountered obsolete JobManager registration success from {} with leader session ID {}.", getTargetAddress(), getTargetLeaderId());
	}
}
 

@Override
protected void onRegistrationSuccess(JMTMRegistrationSuccess success) {
	// filter out old registration attempts
	if (Objects.equals(getTargetLeaderId(), getCurrentJobMasterId())) {
		log.info("Successful registration at job manager {} for job {}.", getTargetAddress(), jobId);

		jobLeaderListener.jobManagerGainedLeadership(jobId, getTargetGateway(), success);
	} else {
		log.debug("Encountered obsolete JobManager registration success from {} with leader session ID {}.", getTargetAddress(), getTargetLeaderId());
	}
}
 

@Override
public void jobManagerGainedLeadership(
	final JobID jobId,
	final JobMasterGateway jobManagerGateway,
	final JMTMRegistrationSuccess registrationMessage) {
	runAsync(
		() ->
			establishJobManagerConnection(
				jobId,
				jobManagerGateway,
				registrationMessage));
}
 

@Override
protected RetryingRegistration<JobMasterId, JobMasterGateway, JMTMRegistrationSuccess> generateRegistration() {
	return new DefaultJobLeaderService.JobManagerRetryingRegistration(
			LOG,
			rpcService,
			"JobManager",
			JobMasterGateway.class,
			getTargetAddress(),
			getTargetLeaderId(),
			retryingRegistrationConfiguration,
			ownerAddress,
			ownLocation);
}
 

@Override
protected RetryingRegistration<JobMasterId, JobMasterGateway, JMTMRegistrationSuccess> generateRegistration() {
	return new JobLeaderService.JobManagerRetryingRegistration(
			LOG,
			rpcService,
			"JobManager",
			JobMasterGateway.class,
			getTargetAddress(),
			getTargetLeaderId(),
			retryingRegistrationConfiguration,
			ownerAddress,
			ownLocation);
}
 

@Override
protected void onRegistrationSuccess(JMTMRegistrationSuccess success) {
	// filter out old registration attempts
	if (Objects.equals(getTargetLeaderId(), currentJobMasterId)) {
		log.info("Successful registration at job manager {} for job {}.", getTargetAddress(), jobId);

		jobLeaderListener.jobManagerGainedLeadership(jobId, getTargetGateway(), success);
	} else {
		log.debug("Encountered obsolete JobManager registration success from {} with leader session ID {}.", getTargetAddress(), getTargetLeaderId());
	}
}
 

@Override
public void jobManagerGainedLeadership(
	final JobID jobId,
	final JobMasterGateway jobManagerGateway,
	final JMTMRegistrationSuccess registrationMessage) {
	runAsync(() ->
		jobTable
			.getJob(jobId)
			.ifPresent(job ->
				establishJobManagerConnection(
					job,
					jobManagerGateway,
					registrationMessage)));
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable<Task> taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			2,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

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

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

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

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Override
public void jobManagerGainedLeadership(JobID jobId, JobMasterGateway jobManagerGateway, JMTMRegistrationSuccess registrationMessage) {
	jobManagerGainedLeadership.accept(jobId);
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			Arrays.asList(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN),
			timerService,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

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

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

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

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * This tests makes sure that duplicate JobMaster gained leadership messages are filtered out
 * by the TaskExecutor. See FLINK-7526.
 */
@Test
public void testFilterOutDuplicateJobMasterRegistrations() throws Exception {
	final long verificationTimeout = 500L;
	final JobLeaderService jobLeaderService = mock(JobLeaderService.class);

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
	when(jobMasterGateway.getHostname()).thenReturn("localhost");
	final JMTMRegistrationSuccess registrationMessage = new JMTMRegistrationSuccess(ResourceID.generate());
	final JobManagerTable jobManagerTableMock = spy(new JobManagerTable());

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

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

	final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(taskManagerServices);

	try {
		taskExecutor.start();
		taskExecutor.waitUntilStarted();

		ArgumentCaptor<JobLeaderListener> jobLeaderListenerArgumentCaptor = ArgumentCaptor.forClass(JobLeaderListener.class);

		verify(jobLeaderService).start(anyString(), any(RpcService.class), any(HighAvailabilityServices.class), jobLeaderListenerArgumentCaptor.capture());

		JobLeaderListener taskExecutorListener = jobLeaderListenerArgumentCaptor.getValue();

		taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage);

		// duplicate job manager gained leadership message
		taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage);

		ArgumentCaptor<JobManagerConnection> jobManagerConnectionArgumentCaptor = ArgumentCaptor.forClass(JobManagerConnection.class);

		verify(jobManagerTableMock, Mockito.timeout(verificationTimeout).times(1)).put(eq(jobId), jobManagerConnectionArgumentCaptor.capture());

		JobManagerConnection jobManagerConnection = jobManagerConnectionArgumentCaptor.getValue();

		assertEquals(jobMasterGateway, jobManagerConnection.getJobManagerGateway());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, 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 TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			Arrays.asList(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN),
			timerService,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

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

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

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

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * This tests makes sure that duplicate JobMaster gained leadership messages are filtered out
 * by the TaskExecutor. See FLINK-7526.
 */
@Test
public void testFilterOutDuplicateJobMasterRegistrations() throws Exception {
	final long verificationTimeout = 500L;
	final JobLeaderService jobLeaderService = mock(JobLeaderService.class);
	final HeartbeatServices heartbeatServicesMock = mock(HeartbeatServices.class, Mockito.RETURNS_MOCKS);

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
	when(jobMasterGateway.getHostname()).thenReturn("localhost");
	final JMTMRegistrationSuccess registrationMessage = new JMTMRegistrationSuccess(ResourceID.generate());
	final JobManagerTable jobManagerTableMock = spy(new JobManagerTable());

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

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

	final TestingTaskExecutor taskExecutor = new TestingTaskExecutor(
		rpc,
		taskManagerConfiguration,
		haServices,
		taskManagerServices,
		heartbeatServicesMock,
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	try {
		taskExecutor.start();
		taskExecutor.waitUntilStarted();

		ArgumentCaptor<JobLeaderListener> jobLeaderListenerArgumentCaptor = ArgumentCaptor.forClass(JobLeaderListener.class);

		verify(jobLeaderService).start(anyString(), any(RpcService.class), any(HighAvailabilityServices.class), jobLeaderListenerArgumentCaptor.capture());

		JobLeaderListener taskExecutorListener = jobLeaderListenerArgumentCaptor.getValue();

		taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage);

		// duplicate job manager gained leadership message
		taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage);

		ArgumentCaptor<JobManagerConnection> jobManagerConnectionArgumentCaptor = ArgumentCaptor.forClass(JobManagerConnection.class);

		verify(jobManagerTableMock, Mockito.timeout(verificationTimeout).times(1)).put(eq(jobId), jobManagerConnectionArgumentCaptor.capture());

		JobManagerConnection jobManagerConnection = jobManagerConnectionArgumentCaptor.getValue();

		assertEquals(jobMasterGateway, jobManagerConnection.getJobManagerGateway());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Callback if a job manager has gained leadership for the job identified by the job id and a
 * connection could be established to this job manager.
 *
 * @param jobId identifying the job for which the job manager has gained leadership
 * @param jobManagerGateway to the job leader
 * @param registrationMessage containing further registration information
 */
void jobManagerGainedLeadership(JobID jobId, JobMasterGateway jobManagerGateway, JMTMRegistrationSuccess registrationMessage);
 

/**
 * Callback if a job manager has gained leadership for the job identified by the job id and a
 * connection could be established to this job manager.
 *
 * @param jobId identifying the job for which the job manager has gained leadership
 * @param jobManagerGateway to the job leader
 * @param registrationMessage containing further registration information
 */
void jobManagerGainedLeadership(JobID jobId, JobMasterGateway jobManagerGateway, JMTMRegistrationSuccess registrationMessage);
 

/**
 * Callback if a job manager has gained leadership for the job identified by the job id and a
 * connection could be established to this job manager.
 *
 * @param jobId identifying the job for which the job manager has gained leadership
 * @param jobManagerGateway to the job leader
 * @param registrationMessage containing further registration information
 */
void jobManagerGainedLeadership(JobID jobId, JobMasterGateway jobManagerGateway, JMTMRegistrationSuccess registrationMessage);