org.apache.flink.runtime.registration.RegistrationResponse Java Examples

/**
 * Test receive normal registration from task executor and receive duplicate registration
 * from task executor.
 */
@Test
public void testRegisterTaskExecutor() throws Exception {
	// test response successful
	CompletableFuture<RegistrationResponse> successfulFuture = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress());

	RegistrationResponse response = successfulFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
	assertTrue(response instanceof TaskExecutorRegistrationSuccess);
	final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
		taskExecutorResourceID,
		TIMEOUT).get();
	assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));

	// test response successful with instanceID not equal to previous when receive duplicate registration from taskExecutor
	CompletableFuture<RegistrationResponse> duplicateFuture = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress());

	RegistrationResponse duplicateResponse = duplicateFuture.get();
	assertTrue(duplicateResponse instanceof TaskExecutorRegistrationSuccess);
	assertNotEquals(((TaskExecutorRegistrationSuccess) response).getRegistrationId(), ((TaskExecutorRegistrationSuccess) duplicateResponse).getRegistrationId());

	assertThat(rmGateway.requestResourceOverview(TIMEOUT).get().getNumberTaskManagers(), is(1));
}
 

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

/**
 * Test receive normal registration from task executor and receive duplicate registration
 * from task executor.
 */
@Test
public void testRegisterTaskExecutor() throws Exception {
	// test response successful
	CompletableFuture<RegistrationResponse> successfulFuture = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress());

	RegistrationResponse response = successfulFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
	assertTrue(response instanceof TaskExecutorRegistrationSuccess);
	final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
		taskExecutorResourceID,
		TIMEOUT).get();
	assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));

	// test response successful with instanceID not equal to previous when receive duplicate registration from taskExecutor
	CompletableFuture<RegistrationResponse> duplicateFuture = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress());

	RegistrationResponse duplicateResponse = duplicateFuture.get();
	assertTrue(duplicateResponse instanceof TaskExecutorRegistrationSuccess);
	assertNotEquals(((TaskExecutorRegistrationSuccess) response).getRegistrationId(), ((TaskExecutorRegistrationSuccess) duplicateResponse).getRegistrationId());

	assertThat(rmGateway.requestResourceOverview(TIMEOUT).get().getNumberTaskManagers(), is(1));
}
 

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

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

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

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

/**
 * Check and verify return RegistrationResponse. Decline when failed to start a
 * job master Leader retrieval listener.
 */
@Test
public void testRegisterJobMasterWithFailureLeaderListener() throws Exception {
	JobID unknownJobIDToHAServices = new JobID();

	// this should fail because we try to register a job leader listener for an unknown job id
	CompletableFuture<RegistrationResponse> registrationFuture = resourceManagerGateway.registerJobManager(
		jobMasterGateway.getFencingToken(),
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		unknownJobIDToHAServices,
		TIMEOUT);

	try {
		registrationFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
		fail("Expected to fail with a ResourceManagerException.");
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.stripExecutionException(e) instanceof ResourceManagerException);
	}

	// ignore the reported error
	testingFatalErrorHandler.clearError();
}
 

@Override
protected RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess> generateRegistration() {
	return new RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess>(
		log,
		getRpcService(),
		"ResourceManager",
		ResourceManagerGateway.class,
		getTargetAddress(),
		getTargetLeaderId(),
		jobMasterConfiguration.getRetryingRegistrationConfiguration()) {

		@Override
		protected CompletableFuture<RegistrationResponse> invokeRegistration(
				ResourceManagerGateway gateway, ResourceManagerId fencingToken, long timeoutMillis) {
			Time timeout = Time.milliseconds(timeoutMillis);

			return gateway.registerJobManager(
				jobMasterId,
				jobManagerResourceID,
				jobManagerRpcAddress,
				jobID,
				timeout);
		}
	};
}
 

/**
 * Test receive registration with unmatched leadershipId from job master.
 */
@Test
public void testRegisterJobMasterWithUnmatchedLeaderSessionId1() throws Exception {
	final ResourceManagerGateway wronglyFencedGateway = rpcService.connect(resourceManager.getAddress(), ResourceManagerId.generate(), ResourceManagerGateway.class)
		.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);

	// test throw exception when receive a registration from job master which takes unmatched leaderSessionId
	CompletableFuture<RegistrationResponse> unMatchedLeaderFuture = wronglyFencedGateway.registerJobManager(
		jobMasterGateway.getFencingToken(),
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);

	try {
		unMatchedLeaderFuture.get(5L, TimeUnit.SECONDS);
		fail("Should fail because we are using the wrong fencing token.");
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.stripExecutionException(e) instanceof FencingTokenException);
	}
}
 

/**
 * Check and verify return RegistrationResponse. Decline when failed to start a
 * job master Leader retrieval listener.
 */
@Test
public void testRegisterJobMasterWithFailureLeaderListener() throws Exception {
	JobID unknownJobIDToHAServices = new JobID();

	// this should fail because we try to register a job leader listener for an unknown job id
	CompletableFuture<RegistrationResponse> registrationFuture = resourceManagerGateway.registerJobManager(
		jobMasterGateway.getFencingToken(),
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		unknownJobIDToHAServices,
		TIMEOUT);

	try {
		registrationFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
		fail("Expected to fail with a ResourceManagerException.");
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.stripExecutionException(e) instanceof ResourceManagerException);
	}

	// ignore the reported error
	testingFatalErrorHandler.clearError();
}
 

/**
 * Check and verify return RegistrationResponse. Decline when failed to start a
 * job master Leader retrieval listener.
 */
@Test
public void testRegisterJobMasterWithFailureLeaderListener() throws Exception {
	JobID unknownJobIDToHAServices = new JobID();

	// this should fail because we try to register a job leader listener for an unknown job id
	CompletableFuture<RegistrationResponse> registrationFuture = resourceManagerGateway.registerJobManager(
		jobMasterGateway.getFencingToken(),
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		unknownJobIDToHAServices,
		TIMEOUT);

	try {
		registrationFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
		fail("Expected to fail with a ResourceManagerException.");
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.stripExecutionException(e) instanceof ResourceManagerException);
	}

	// ignore the reported error
	testingFatalErrorHandler.clearError();
}
 

@Override
protected RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess> generateRegistration() {
	return new RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess>(
		log,
		getRpcService(),
		"ResourceManager",
		ResourceManagerGateway.class,
		getTargetAddress(),
		getTargetLeaderId(),
		jobMasterConfiguration.getRetryingRegistrationConfiguration()) {

		@Override
		protected CompletableFuture<RegistrationResponse> invokeRegistration(
				ResourceManagerGateway gateway, ResourceManagerId fencingToken, long timeoutMillis) {
			Time timeout = Time.milliseconds(timeoutMillis);

			return gateway.registerJobManager(
				jobMasterId,
				jobManagerResourceID,
				jobManagerRpcAddress,
				jobID,
				timeout);
		}
	};
}
 

@Override
protected RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess> generateRegistration() {
	return new RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess>(
		log,
		getRpcService(),
		"ResourceManager",
		ResourceManagerGateway.class,
		getTargetAddress(),
		getTargetLeaderId(),
		jobMasterConfiguration.getRetryingRegistrationConfiguration()) {

		@Override
		protected CompletableFuture<RegistrationResponse> invokeRegistration(
				ResourceManagerGateway gateway, ResourceManagerId fencingToken, long timeoutMillis) {
			Time timeout = Time.milliseconds(timeoutMillis);

			return gateway.registerJobManager(
				jobMasterId,
				jobManagerResourceID,
				jobManagerRpcAddress,
				jobID,
				timeout);
		}
	};
}
 

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

@Override
public CompletableFuture<RegistrationResponse> registerTaskExecutor(String taskExecutorAddress, ResourceID resourceId, int dataPort, HardwareDescription hardwareDescription, Time timeout) {
	final Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> currentFunction = registerTaskExecutorFunction;

	if (currentFunction != null) {
		return currentFunction.apply(Tuple4.of(taskExecutorAddress, resourceId, dataPort, hardwareDescription));
	} else {
		return CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				ownResourceId,
				new ClusterInformation("localhost", 1234)));
	}
}
 

@Override
public CompletableFuture<RegistrationResponse> registerJobManager(JobMasterId jobMasterId, ResourceID jobMasterResourceId, String jobMasterAddress, JobID jobId, Time timeout) {
	final Consumer<Tuple4<JobMasterId, ResourceID, String, JobID>> currentConsumer = registerJobManagerConsumer;

	if (currentConsumer != null) {
		currentConsumer.accept(Tuple4.of(jobMasterId, jobMasterResourceId, jobMasterAddress, jobId));
	}

	return CompletableFuture.completedFuture(
		new JobMasterRegistrationSuccess(
			resourceManagerId,
			ownResourceId));
}
 

private CompletableFuture<RegistrationResponse> createRegistrationResponse(TestingResourceManagerGateway rmGateway1) {
	return CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmGateway1.getOwnResourceId(),
			new ClusterInformation("localhost", 1234)));
}
 

/**
 * Test receive registration with invalid address from job master.
 */
@Test
public void testRegisterJobMasterFromInvalidAddress() throws Exception {
	// test throw exception when receive a registration from job master which takes invalid address
	String invalidAddress = "/jobMasterAddress2";
	CompletableFuture<RegistrationResponse> invalidAddressFuture = resourceManagerGateway.registerJobManager(
		new JobMasterId(HighAvailabilityServices.DEFAULT_LEADER_ID),
		jobMasterResourceId,
		invalidAddress,
		jobId,
		TIMEOUT);
	assertTrue(invalidAddressFuture.get(5, TimeUnit.SECONDS) instanceof RegistrationResponse.Decline);
}
 

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

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

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

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

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

/**
 * Test receive registration with invalid address from task executor.
 */
@Test
public void testRegisterTaskExecutorFromInvalidAddress() throws Exception {
	// test throw exception when receive a registration from taskExecutor which takes invalid address
	String invalidAddress = "/taskExecutor2";

	CompletableFuture<RegistrationResponse> invalidAddressFuture = registerTaskExecutor(rmGateway, invalidAddress);
	assertTrue(invalidAddressFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS) instanceof RegistrationResponse.Decline);
}
 

/**
 * Test worker registration after launch.
 */
@Test
public void testWorkerStarted() throws Exception {
	new Context() {{
		// set the initial state with a (recovered) launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

		final int dataPort = 1234;
		final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
		// send registration message
		CompletableFuture<RegistrationResponse> successfulFuture =
			resourceManager.registerTaskExecutor(task1Executor.address, task1Executor.resourceID, dataPort, hardwareDescription, timeout);
		RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);
		final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;

		final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);

		// check for errors
		initialSlotReportFuture.get();

		// verify the internal state
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
	}};
}
 

/**
 * Test worker registration after launch.
 */
@Test
public void testWorkerStarted() throws Exception {
	new Context() {{
		// set the initial state with a (recovered) launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

		final int dataPort = 1234;
		final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
		// send registration message
		TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
			task1Executor.address,
			task1Executor.resourceID,
			dataPort,
			hardwareDescription,
			ResourceProfile.ZERO,
			ResourceProfile.ZERO);
		CompletableFuture<RegistrationResponse> successfulFuture =
			resourceManager.registerTaskExecutor(taskExecutorRegistration, timeout);
		RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);
		final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;

		final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);

		// check for errors
		initialSlotReportFuture.get();

		// verify the internal state
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
	}};
}
 

@Override
protected CompletableFuture<RegistrationResponse> invokeRegistration(
		ResourceManagerGateway resourceManager, ResourceManagerId fencingToken, long timeoutMillis) throws Exception {

	Time timeout = Time.milliseconds(timeoutMillis);
	return resourceManager.registerTaskExecutor(
		taskExecutorRegistration,
		timeout);
}
 

/**
 * Test receive normal registration from job master and receive duplicate
 * registration from job master.
 */
@Test
public void testRegisterJobMaster() throws Exception {
	// test response successful
	CompletableFuture<RegistrationResponse> successfulFuture = resourceManagerGateway.registerJobManager(
		jobMasterGateway.getFencingToken(),
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);
	RegistrationResponse response = successfulFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
	assertTrue(response instanceof JobMasterRegistrationSuccess);
}
 

@Override
public CompletableFuture<RegistrationResponse> registerJobManager(JobMasterId jobMasterId, ResourceID jobMasterResourceId, String jobMasterAddress, JobID jobId, Time timeout) {
	final QuadFunction<JobMasterId, ResourceID, String, JobID, CompletableFuture<RegistrationResponse>> currentConsumer = registerJobManagerFunction;

	if (currentConsumer != null) {
		return currentConsumer.apply(jobMasterId, jobMasterResourceId, jobMasterAddress, jobId);
	}

	return CompletableFuture.completedFuture(getJobMasterRegistrationSuccess());
}
 

/**
 * Test receive registration with unmatched leadershipId from job master.
 */
@Test
public void testRegisterJobMasterWithUnmatchedLeaderSessionId2() throws Exception {
			// test throw exception when receive a registration from job master which takes unmatched leaderSessionId
	JobMasterId differentJobMasterId = JobMasterId.generate();
	CompletableFuture<RegistrationResponse> unMatchedLeaderFuture = resourceManagerGateway.registerJobManager(
		differentJobMasterId,
		jobMasterResourceId,
		jobMasterGateway.getAddress(),
		jobId,
		TIMEOUT);
	assertTrue(unMatchedLeaderFuture.get() instanceof RegistrationResponse.Decline);
}
 

@Override
protected CompletableFuture<RegistrationResponse> invokeRegistration(
		JobMasterGateway gateway,
		JobMasterId jobMasterId,
		long timeoutMillis) throws Exception {
	return gateway.registerTaskManager(taskManagerRpcAddress, taskManagerLocation, Time.milliseconds(timeoutMillis));
}
 

@Override
protected CompletableFuture<RegistrationResponse> invokeRegistration(
		ResourceManagerGateway resourceManager, ResourceManagerId fencingToken, long timeoutMillis) throws Exception {

	Time timeout = Time.milliseconds(timeoutMillis);
	return resourceManager.registerTaskExecutor(
		taskExecutorAddress,
		resourceID,
		dataPort,
		hardwareDescription,
		timeout);
}
 

/**
 * Test worker registration after launch.
 */
@Test
public void testWorkerStarted() throws Exception {
	new Context() {{
		// set the initial state with a (recovered) launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

		final int dataPort = 1234;
		final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
		// send registration message
		CompletableFuture<RegistrationResponse> successfulFuture =
			resourceManager.registerTaskExecutor(task1Executor.address, task1Executor.resourceID, dataPort, hardwareDescription, timeout);
		RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);
		final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;

		final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);

		// check for errors
		initialSlotReportFuture.get();

		// verify the internal state
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
	}};
}
 

/**
 * Test receive registration with invalid address from job master.
 */
@Test
public void testRegisterJobMasterFromInvalidAddress() throws Exception {
	// test throw exception when receive a registration from job master which takes invalid address
	String invalidAddress = "/jobMasterAddress2";
	CompletableFuture<RegistrationResponse> invalidAddressFuture = resourceManagerGateway.registerJobManager(
		new JobMasterId(HighAvailabilityServices.DEFAULT_LEADER_ID),
		jobMasterResourceId,
		invalidAddress,
		jobId,
		TIMEOUT);
	assertTrue(invalidAddressFuture.get(5, TimeUnit.SECONDS) instanceof RegistrationResponse.Decline);
}
 

private void registerTaskExecutor(ResourceManagerGateway resourceManagerGateway, ResourceID taskExecutorId, String taskExecutorAddress) throws Exception {
	TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
		taskExecutorAddress,
		taskExecutorId,
		dataPort,
		hardwareDescription,
		ResourceProfile.ZERO,
		ResourceProfile.ZERO);
	final CompletableFuture<RegistrationResponse> registrationFuture = resourceManagerGateway.registerTaskExecutor(
		taskExecutorRegistration,
		TestingUtils.TIMEOUT());

	assertThat(registrationFuture.get(), instanceOf(RegistrationResponse.Success.class));
}