@Override
    public CompletableFuture<RegistrationResponse> registerJobMaster(
            final JobMasterId jobMasterId,
            final ResourceID jobManagerResourceId,
            final String jobManagerAddress,
            final JobID jobId,
            final Time timeout) {

        checkNotNull(jobMasterId);
        checkNotNull(jobManagerResourceId);
        checkNotNull(jobManagerAddress);
        checkNotNull(jobId);

        if (!jobLeaderIdService.containsJob(jobId)) {
            try {
                jobLeaderIdService.addJob(jobId);
            } catch (Exception e) {
                ResourceManagerException exception =
                        new ResourceManagerException(
                                "Could not add the job " + jobId + " to the job id leader service.",
                                e);

                onFatalError(exception);

                log.error("Could not add job {} to job leader id service.", jobId, e);
                return FutureUtils.completedExceptionally(exception);
            }
        }

        log.info(
                "Registering job manager {}@{} for job {}.", jobMasterId, jobManagerAddress, jobId);

        CompletableFuture<JobMasterId> jobMasterIdFuture;

        try {
            jobMasterIdFuture = jobLeaderIdService.getLeaderId(jobId);
        } catch (Exception e) {
            // we cannot check the job leader id so let's fail
            // TODO: Maybe it's also ok to skip this check in case that we cannot check the leader
            // id
            ResourceManagerException exception =
                    new ResourceManagerException(
                            "Cannot obtain the "
                                    + "job leader id future to verify the correct job leader.",
                            e);

            onFatalError(exception);

            log.debug(
                    "Could not obtain the job leader id future to verify the correct job leader.");
            return FutureUtils.completedExceptionally(exception);
        }

        CompletableFuture<JobMasterGateway> jobMasterGatewayFuture =
                getRpcService().connect(jobManagerAddress, jobMasterId, JobMasterGateway.class);

        CompletableFuture<RegistrationResponse> registrationResponseFuture =
                jobMasterGatewayFuture.thenCombineAsync(
                        jobMasterIdFuture,
                        (JobMasterGateway jobMasterGateway, JobMasterId leadingJobMasterId) -> {
                            if (Objects.equals(leadingJobMasterId, jobMasterId)) {
                                return registerJobMasterInternal(
                                        jobMasterGateway,
                                        jobId,
                                        jobManagerAddress,
                                        jobManagerResourceId);
                            } else {
                                final String declineMessage =
                                        String.format(
                                                "The leading JobMaster id %s did not match the received JobMaster id %s. "
                                                        + "This indicates that a JobMaster leader change has happened.",
                                                leadingJobMasterId, jobMasterId);
                                log.debug(declineMessage);
                                return new RegistrationResponse.Failure(
                                        new FlinkException(declineMessage));
                            }
                        },
                        getMainThreadExecutor(jobId));

        // handle exceptions which might have occurred in one of the futures inputs of combine
        return registrationResponseFuture.handleAsync(
                (RegistrationResponse registrationResponse, Throwable throwable) -> {
                    if (throwable != null) {
                        if (log.isDebugEnabled()) {
                            log.debug(
                                    "Registration of job manager {}@{} failed.",
                                    jobMasterId,
                                    jobManagerAddress,
                                    throwable);
                        } else {
                            log.info(
                                    "Registration of job manager {}@{} failed.",
                                    jobMasterId,
                                    jobManagerAddress);
                        }

                        return new RegistrationResponse.Failure(throwable);
                    } else {
                        return registrationResponse;
                    }
                },
                ioExecutor);
    }
@Test
    void testJobMasterBecomesUnreachableTriggersDisconnect() throws Exception {
        final JobID jobId = new JobID();
        final ResourceID jobMasterResourceId = ResourceID.generate();
        final CompletableFuture<ResourceManagerId> disconnectFuture = new CompletableFuture<>();
        final TestingJobMasterGateway jobMasterGateway =
                new TestingJobMasterGatewayBuilder()
                        .setAddress(UUID.randomUUID().toString())
                        .setResourceManagerHeartbeatFunction(
                                resourceId ->
                                        FutureUtils.completedExceptionally(
                                                new RecipientUnreachableException(
                                                        "sender",
                                                        "recipient",
                                                        "task executor is unreachable")))
                        .setDisconnectResourceManagerConsumer(disconnectFuture::complete)
                        .build();
        rpcService.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);

        final LeaderRetrievalService jobMasterLeaderRetrievalService =
                new SettableLeaderRetrievalService(
                        jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());

        highAvailabilityServices.setJobMasterLeaderRetrieverFunction(
                requestedJobId -> {
                    assertThat(requestedJobId).isEqualTo(jobId);
                    return jobMasterLeaderRetrievalService;
                });

        runHeartbeatTargetBecomesUnreachableTest(
                (ignore) -> {},
                resourceManagerGateway -> {
                    final CompletableFuture<RegistrationResponse> registrationFuture =
                            resourceManagerGateway.registerJobMaster(
                                    jobMasterGateway.getFencingToken(),
                                    jobMasterResourceId,
                                    jobMasterGateway.getAddress(),
                                    jobId,
                                    TIMEOUT);

                    assertThatFuture(registrationFuture)
                            .eventuallySucceeds()
                            .isInstanceOf(RegistrationResponse.Success.class);
                },
                resourceManagerResourceId ->
                        assertThatFuture(disconnectFuture)
                                .eventuallySucceeds()
                                .isEqualTo(resourceManagerId));
    }