org.apache.flink.runtime.rpc.RpcUtils Java Examples

@Test
public void testSerializationFailureHandling() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	DispatcherGateway mockGateway = new TestingDispatcherGateway.Builder()
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance()), mockGateway);
		Assert.fail();
	} catch (RestHandlerException rhe) {
		Assert.assertEquals(HttpResponseStatus.BAD_REQUEST, rhe.getHttpResponseStatus());
	}
}
 

/**
 * Tests that the {@link AkkaRpcActor} only completes after the asynchronous
 * post stop action has completed.
 */
@Test
public void testActorTerminationWithAsynchronousOnStopAction() throws Exception {
	final CompletableFuture<Void> onStopFuture = new CompletableFuture<>();
	final AsynchronousOnStopEndpoint endpoint = new AsynchronousOnStopEndpoint(akkaRpcService, onStopFuture);

	try {
		endpoint.start();

		final CompletableFuture<Void> terminationFuture = endpoint.closeAsync();

		assertFalse(terminationFuture.isDone());

		onStopFuture.complete(null);

		// the onStopFuture completion should allow the endpoint to terminate
		terminationFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(endpoint, timeout);
	}
}
 

/**
 * Tests that the {@link AkkaRpcActor} only completes after the asynchronous
 * post stop action has completed.
 */
@Test
public void testActorTerminationWithAsynchronousOnStopAction() throws Exception {
	final CompletableFuture<Void> onStopFuture = new CompletableFuture<>();
	final AsynchronousOnStopEndpoint endpoint = new AsynchronousOnStopEndpoint(akkaRpcService, onStopFuture);

	try {
		endpoint.start();

		final CompletableFuture<Void> terminationFuture = endpoint.closeAsync();

		assertFalse(terminationFuture.isDone());

		onStopFuture.complete(null);

		// the onStopFuture completion should allow the endpoint to terminate
		terminationFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(endpoint, timeout);
	}
}
 

/**
 * Tests that the {@link MiniDispatcher} recovers the single job with which it
 * was started.
 */
@Test
public void testSingleJobRecovery() throws Exception {
	final MiniDispatcher miniDispatcher = createMiniDispatcher(ClusterEntrypoint.ExecutionMode.DETACHED);

	miniDispatcher.start();

	try {
		// wait until the Dispatcher is the leader
		dispatcherLeaderElectionService.isLeader(UUID.randomUUID()).get();

		final JobGraph actualJobGraph = jobGraphFuture.get();

		assertThat(actualJobGraph.getJobID(), is(jobGraph.getJobID()));
	} finally {
		RpcUtils.terminateRpcEndpoint(miniDispatcher, timeout);
	}
}
 

@Test
public void testMaximumRegistrationDuration() throws Exception {
	configuration.set(TaskManagerOptions.REGISTRATION_TIMEOUT, TimeUtils.parseDuration("10 ms"));

	final TaskExecutor taskExecutor = createTaskExecutor(new TaskManagerServicesBuilder().build());

	taskExecutor.start();

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

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

@After
public void tearDown() throws Exception {
	try {
		fatalErrorHandler.rethrowError();
	} finally {
		if (dispatcher != null) {
			RpcUtils.terminateRpcEndpoint(dispatcher, TIMEOUT);
		}
	}

	if (haServices != null) {
		haServices.closeAndCleanupAllData();
	}

	if (blobServer != null) {
		blobServer.close();
	}
}
 

@Test
public void createResourceManager_WithLessMemoryThanContainerizedHeapCutoffMin_ShouldSucceed() throws Exception {
	final StandaloneResourceManagerFactory resourceManagerFactory = StandaloneResourceManagerFactory.INSTANCE;

	final TestingRpcService rpcService = new TestingRpcService();
	try {
		final Configuration configuration = new Configuration();
		configuration.setString(TaskManagerOptions.TASK_MANAGER_HEAP_MEMORY, new MemorySize(128 * 1024 * 1024).toString());
		configuration.setInteger(ResourceManagerOptions.CONTAINERIZED_HEAP_CUTOFF_MIN, 600);

		final ResourceManager<ResourceID> ignored = resourceManagerFactory.createResourceManager(
			configuration,
			ResourceID.generate(),
			rpcService,
			new TestingHighAvailabilityServices(),
			new TestingHeartbeatServices(),
			NoOpMetricRegistry.INSTANCE,
			new TestingFatalErrorHandler(),
			new ClusterInformation("foobar", 1234),
			null,
			UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup());
	} finally {
		RpcUtils.terminateRpcService(rpcService, Time.seconds(10L));
	}
}
 

@Test
public void testSerializationFailureHandling() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	DispatcherGateway mockGateway = new TestingDispatcherGateway.Builder()
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance()), mockGateway);
		Assert.fail();
	} catch (RestHandlerException rhe) {
		Assert.assertEquals(HttpResponseStatus.BAD_REQUEST, rhe.getHttpResponseStatus());
	}
}
 

@Test
public void testVersionMismatchBetweenRpcComponents() throws Exception {
	AkkaRpcActorTest.DummyRpcEndpoint rpcEndpoint = new AkkaRpcActorTest.DummyRpcEndpoint(akkaRpcService1);

	rpcEndpoint.start();

	try {
		try {
			wrongVersionAkkaRpcService.connect(rpcEndpoint.getAddress(), AkkaRpcActorTest.DummyRpcGateway.class).get();
			fail("Expected HandshakeException.");
		} catch (ExecutionException ee) {
			assertThat(ExceptionUtils.stripExecutionException(ee), instanceOf(HandshakeException.class));
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(rpcEndpoint, timeout);
	}
}
 

/**
 * Tests that the {@link TaskExecutor} sends the initial slot report after it
 * registered at the ResourceManager.
 */
@Test
public void testInitialSlotReport() throws Exception {
	final TaskExecutor taskExecutor = createTaskExecutor(1);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f0);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

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

		assertThat(initialSlotReportFuture.get(), equalTo(taskExecutor.getResourceID()));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Override
protected CompletableFuture<String> triggerOperation(
		final HandlerRequest<StopWithSavepointRequestBody, SavepointTriggerMessageParameters> request,
		final RestfulGateway gateway) throws RestHandlerException {

	final JobID jobId = request.getPathParameter(JobIDPathParameter.class);
	final String requestedTargetDirectory = request.getRequestBody().getTargetDirectory();

	if (requestedTargetDirectory == null && defaultSavepointDir == null) {
		throw new RestHandlerException(
				String.format("Config key [%s] is not set. Property [%s] must be provided.",
						CheckpointingOptions.SAVEPOINT_DIRECTORY.key(),
						StopWithSavepointRequestBody.FIELD_NAME_TARGET_DIRECTORY),
				HttpResponseStatus.BAD_REQUEST);
	}

	final boolean advanceToEndOfEventTime = request.getRequestBody().shouldDrain();
	final String targetDirectory = requestedTargetDirectory != null ? requestedTargetDirectory : defaultSavepointDir;
	return gateway.stopWithSavepoint(jobId, targetDirectory, advanceToEndOfEventTime, RpcUtils.INF_TIMEOUT);
}
 

@Override
protected CompletableFuture<Acknowledge> triggerOperation(HandlerRequest<EmptyRequestBody, RescalingTriggerMessageParameters> request, RestfulGateway gateway) throws RestHandlerException {
	final JobID jobId = request.getPathParameter(JobIDPathParameter.class);
	final List<Integer> queryParameter = request.getQueryParameter(RescalingParallelismQueryParameter.class);

	if (queryParameter.isEmpty()) {
		throw new RestHandlerException("No new parallelism was specified.", HttpResponseStatus.BAD_REQUEST);
	}

	final int newParallelism = queryParameter.get(0);

	final CompletableFuture<Acknowledge> rescalingFuture = gateway.rescaleJob(
		jobId,
		newParallelism,
		RescalingBehaviour.STRICT,
		RpcUtils.INF_TIMEOUT);

	return rescalingFuture;
}
 

@Override
protected CompletableFuture<String> triggerOperation(
		final HandlerRequest<StopWithSavepointRequestBody, SavepointTriggerMessageParameters> request,
		final RestfulGateway gateway) throws RestHandlerException {

	final JobID jobId = request.getPathParameter(JobIDPathParameter.class);
	final String requestedTargetDirectory = request.getRequestBody().getTargetDirectory();

	if (requestedTargetDirectory == null && defaultSavepointDir == null) {
		throw new RestHandlerException(
				String.format("Config key [%s] is not set. Property [%s] must be provided.",
						CheckpointingOptions.SAVEPOINT_DIRECTORY.key(),
						StopWithSavepointRequestBody.FIELD_NAME_TARGET_DIRECTORY),
				HttpResponseStatus.BAD_REQUEST);
	}

	final boolean advanceToEndOfEventTime = request.getRequestBody().shouldDrain();
	final String targetDirectory = requestedTargetDirectory != null ? requestedTargetDirectory : defaultSavepointDir;
	return gateway.stopWithSavepoint(jobId, targetDirectory, advanceToEndOfEventTime, RpcUtils.INF_TIMEOUT);
}
 

@After
public void teardown() throws Exception {
	if (rpc != null) {
		RpcUtils.terminateRpcService(rpc, timeout);
		rpc = null;
	}

	if (timerService != null) {
		timerService.stop();
		timerService = null;
	}

	if (dummyBlobCacheService != null) {
		dummyBlobCacheService.close();
		dummyBlobCacheService = null;
	}

	testingFatalErrorHandler.rethrowError();
}
 

/**
 * Tests that a JobMaster will restore the given JobGraph from its savepoint upon
 * initial submission.
 */
@Test
public void testRestoringFromSavepoint() throws Exception {

	// create savepoint data
	final long savepointId = 42L;
	final File savepointFile = createSavepoint(savepointId);

	// set savepoint settings
	final SavepointRestoreSettings savepointRestoreSettings = SavepointRestoreSettings.forPath(
		savepointFile.getAbsolutePath(),
		true);
	final JobGraph jobGraph = createJobGraphWithCheckpointing(savepointRestoreSettings);

	final StandaloneCompletedCheckpointStore completedCheckpointStore = new StandaloneCompletedCheckpointStore(1);
	final TestingCheckpointRecoveryFactory testingCheckpointRecoveryFactory = new TestingCheckpointRecoveryFactory(completedCheckpointStore, new StandaloneCheckpointIDCounter());
	haServices.setCheckpointRecoveryFactory(testingCheckpointRecoveryFactory);
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	try {
		// starting the JobMaster should have read the savepoint
		final CompletedCheckpoint savepointCheckpoint = completedCheckpointStore.getLatestCheckpoint(false);

		assertThat(savepointCheckpoint, Matchers.notNullValue());

		assertThat(savepointCheckpoint.getCheckpointID(), is(savepointId));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Tests that we can still run commands via the main thread executor when the onStop method
 * is called.
 */
@Test
public void testMainThreadExecutionOnStop() throws Exception {
	final MainThreadExecutorOnStopEndpoint endpoint = new MainThreadExecutorOnStopEndpoint(akkaRpcService);

	try {
		endpoint.start();

		CompletableFuture<Void> terminationFuture = endpoint.closeAsync();

		terminationFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(endpoint, timeout);
	}
}
 

/**
 * Tests that the {@link MiniDispatcher} only terminates in {@link ClusterEntrypoint.ExecutionMode#NORMAL}
 * after it has served the {@link org.apache.flink.runtime.jobmaster.JobResult} once.
 */
@Test
public void testJobResultRetrieval() throws Exception {
	final MiniDispatcher miniDispatcher = createMiniDispatcher(ClusterEntrypoint.ExecutionMode.NORMAL);

	miniDispatcher.start();

	try {
		// wait until we have submitted the job
		final TestingJobManagerRunner testingJobManagerRunner = testingJobManagerRunnerFactory.takeCreatedJobManagerRunner();

		testingJobManagerRunner.completeResultFuture(archivedExecutionGraph);

		assertFalse(miniDispatcher.getTerminationFuture().isDone());

		final DispatcherGateway dispatcherGateway = miniDispatcher.getSelfGateway(DispatcherGateway.class);

		final CompletableFuture<JobResult> jobResultFuture = dispatcherGateway.requestJobResult(jobGraph.getJobID(), timeout);

		final JobResult jobResult = jobResultFuture.get();

		assertThat(jobResult.getJobId(), is(jobGraph.getJobID()));
	}
	finally {
		RpcUtils.terminateRpcEndpoint(miniDispatcher, timeout);
	}
}
 

/**
 * Tests that we continue reconnecting to the latest known RM after a disconnection
 * message.
 */
@Test
public void testReconnectionAfterDisconnect() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
		final BlockingQueue<JobMasterId> registrationsQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationsQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		final ResourceManagerId resourceManagerId = testingResourceManagerGateway.getFencingToken();
		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		// wait for first registration attempt
		final JobMasterId firstRegistrationAttempt = registrationsQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		assertThat(registrationsQueue.isEmpty(), is(true));
		jobMasterGateway.disconnectResourceManager(resourceManagerId, new FlinkException("Test exception"));

		// wait for the second registration attempt after the disconnect call
		assertThat(registrationsQueue.take(), equalTo(jobMasterId));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@After
public void after() throws Exception {
	if (resourceManager != null) {
		RpcUtils.terminateRpcEndpoint(resourceManager, TIMEOUT);
	}

	if (highAvailabilityServices != null) {
		highAvailabilityServices.closeAndCleanupAllData();
	}

	if (testingFatalErrorHandler.hasExceptionOccurred()) {
		testingFatalErrorHandler.rethrowError();
	}
}
 

@Test
public void testFailedJobSubmission() throws Exception {
	final String errorMessage = "test";
	DispatcherGateway mockGateway = new TestingDispatcherGateway.Builder()
		.setSubmitFunction(jobgraph -> FutureUtils.completedExceptionally(new Exception(errorMessage)))
		.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();

	JobGraph jobGraph = new JobGraph("testjob");
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(jobGraph);
	}
	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(
				request,
				EmptyMessageParameters.getInstance(),
				Collections.emptyMap(),
				Collections.emptyMap(),
				Collections.singletonList(jobGraphFile.toFile())), mockGateway)
			.get();
	} catch (Exception e) {
		Throwable t = ExceptionUtils.stripExecutionException(e);
		Assert.assertEquals(errorMessage, t.getMessage());
	}
}
 

@AfterClass
public static void teardownClass() throws IOException, InterruptedException, ExecutionException, TimeoutException {
	if (blobServer != null) {
		blobServer.close();
	}

	if (rpcService != null) {
		RpcUtils.terminateRpcService(rpcService, timeout);
	}
}
 

@After
public void teardown() throws Exception {
	if (resourceManager != null) {
		RpcUtils.terminateRpcEndpoint(resourceManager, TIMEOUT);
	}

	if (testingFatalErrorHandler != null && testingFatalErrorHandler.hasExceptionOccurred()) {
		testingFatalErrorHandler.rethrowError();
	}
}
 

/**
 * Tests that the {@link AkkaRpcActor} discards messages until the corresponding
 * {@link RpcEndpoint} has been started.
 */
@Test
public void testMessageDiscarding() throws Exception {
	int expectedValue = 1337;

	DummyRpcEndpoint rpcEndpoint = new DummyRpcEndpoint(akkaRpcService);

	DummyRpcGateway rpcGateway = rpcEndpoint.getSelfGateway(DummyRpcGateway.class);

	// this message should be discarded and completed with an AkkaRpcException
	CompletableFuture<Integer> result = rpcGateway.foobar();

	try {
		result.get(timeout.getSize(), timeout.getUnit());
		fail("Expected an AkkaRpcException.");
	} catch (ExecutionException ee) {
		// expected this exception, because the endpoint has not been started
		assertTrue(ee.getCause() instanceof AkkaRpcException);
	}

	// set a new value which we expect to be returned
	rpcEndpoint.setFoobar(expectedValue);

	// start the endpoint so that it can process messages
	rpcEndpoint.start();

	try {
		// send the rpc again
		result = rpcGateway.foobar();

		// now we should receive a result :-)
		Integer actualValue = result.get(timeout.getSize(), timeout.getUnit());

		assertThat("The new foobar value should have been returned.", actualValue, Is.is(expectedValue));
	} finally {
		RpcUtils.terminateRpcEndpoint(rpcEndpoint, timeout);
	}
}
 

@AfterClass
public static void teardown() throws InterruptedException, ExecutionException, TimeoutException {
	if (rpcService != null) {
		RpcUtils.terminateRpcService(rpcService, TIMEOUT);
		rpcService = null;
	}
}
 

@AfterClass
public static void teardownClass() throws IOException, InterruptedException, ExecutionException, TimeoutException {
	if (blobServer != null) {
		blobServer.close();
	}

	if (rpcService != null) {
		RpcUtils.terminateRpcService(rpcService, timeout);
	}
}
 

/**
 * Tests that the a JM connects to the leading RM after regaining leadership.
 */
@Test
public void testResourceManagerConnectionAfterRegainingLeadership() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();

		final BlockingQueue<JobMasterId> registrationQueue = new ArrayBlockingQueue<>(1);
		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		final JobMasterId firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		jobMaster.suspend(new FlinkException("Test exception.")).get();

		final JobMasterId jobMasterId2 = JobMasterId.generate();

		jobMaster.start(jobMasterId2).get();

		final JobMasterId secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(jobMasterId2));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@Test
public void testRequestKvStateWithIrrelevantRegistration() throws Exception {
	final JobGraph graph = createKvJobGraph();

	final JobMaster jobMaster = createJobMaster(
		configuration,
		graph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);
	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// register an irrelevant KvState
		try {
			jobMasterGateway.notifyKvStateRegistered(
				new JobID(),
				new JobVertexID(),
				new KeyGroupRange(0, 0),
				"any-name",
				new KvStateID(),
				new InetSocketAddress(InetAddress.getLocalHost(), 1233)).get();
			fail("Expected to fail with FlinkJobNotFoundException.");
		} catch (Exception e) {
			assertTrue(ExceptionUtils.findThrowable(e, FlinkJobNotFoundException.class).isPresent());
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@Test
public void testRequestKvStateOfWrongJob() throws Exception {
	final JobGraph graph = createKvJobGraph();

	final JobMaster jobMaster = createJobMaster(
		configuration,
		graph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);
	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// lookup location
		try {
			jobMasterGateway.requestKvStateLocation(new JobID(), "unknown").get();
			fail("Expected to fail with FlinkJobNotFoundException");
		} catch (Exception e) {
			assertTrue(ExceptionUtils.findThrowable(e, FlinkJobNotFoundException.class).isPresent());
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Start the rpc service and begin to run the job.
 *
 * @param newJobMasterId The necessary fencing token to run the job
 * @return Future acknowledge if the job could be started. Otherwise the future contains an exception
 */
public CompletableFuture<Acknowledge> start(final JobMasterId newJobMasterId) throws Exception {
	// make sure we receive RPC and async calls
	start();

	return callAsyncWithoutFencing(() -> startJobExecution(newJobMasterId), RpcUtils.INF_TIMEOUT);
}
 

/**
 * Tests that the a JM connects to the leading RM after regaining leadership.
 */
@Test
public void testResourceManagerConnectionAfterRegainingLeadership() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();

		final BlockingQueue<JobMasterId> registrationQueue = new ArrayBlockingQueue<>(1);
		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		final JobMasterId firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		jobMaster.suspend(new FlinkException("Test exception.")).get();

		final JobMasterId jobMasterId2 = JobMasterId.generate();

		jobMaster.start(jobMasterId2).get();

		final JobMasterId secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(jobMasterId2));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}