org.apache.flink.client.deployment.ClusterClientJobClientAdapter Maven / Gradle / Ivy

Go to download
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.client.deployment;

import org.apache.flink.api.common.JobExecutionResult;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.JobStatus;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.ClusterClientProvider;
import org.apache.flink.client.program.ProgramInvocationException;
import org.apache.flink.core.execution.JobClient;
import org.apache.flink.core.execution.SavepointFormatType;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.operators.coordination.CoordinationRequest;
import org.apache.flink.runtime.operators.coordination.CoordinationRequestGateway;
import org.apache.flink.runtime.operators.coordination.CoordinationResponse;
import org.apache.flink.util.concurrent.FutureUtils;

import org.apache.commons.io.IOUtils;

import javax.annotation.Nullable;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.function.Function;

import static org.apache.flink.util.Preconditions.checkNotNull;

/**
 * An implementation of the {@link JobClient} interface that uses a {@link ClusterClient}
 * underneath..
 */
public class ClusterClientJobClientAdapter
        implements JobClient, CoordinationRequestGateway {

    private final ClusterClientProvider clusterClientProvider;

    private final JobID jobID;

    private final ClassLoader classLoader;

    public ClusterClientJobClientAdapter(
            final ClusterClientProvider clusterClientProvider,
            final JobID jobID,
            final ClassLoader classLoader) {
        this.jobID = checkNotNull(jobID);
        this.clusterClientProvider = checkNotNull(clusterClientProvider);
        this.classLoader = classLoader;
    }

    @Override
    public JobID getJobID() {
        return jobID;
    }

    @Override
    public CompletableFuture getJobStatus() {
        return bridgeClientRequest(
                clusterClientProvider, (clusterClient -> clusterClient.getJobStatus(jobID)));
    }

    @Override
    public CompletableFuture cancel() {
        return bridgeClientRequest(
                clusterClientProvider,
                (clusterClient -> clusterClient.cancel(jobID).thenApply((ignored) -> null)));
    }

    @Override
    public CompletableFuture stopWithSavepoint(
            boolean advanceToEndOfEventTime,
            @Nullable String savepointDirectory,
            SavepointFormatType formatType) {
        return bridgeClientRequest(
                clusterClientProvider,
                (clusterClient ->
                        clusterClient.stopWithSavepoint(
                                jobID, advanceToEndOfEventTime, savepointDirectory, formatType)));
    }

    @Override
    public CompletableFuture triggerSavepoint(
            @Nullable String savepointDirectory, SavepointFormatType formatType) {
        return bridgeClientRequest(
                clusterClientProvider,
                (clusterClient ->
                        clusterClient.triggerSavepoint(jobID, savepointDirectory, formatType)));
    }

    @Override
    public CompletableFuture> getAccumulators() {
        checkNotNull(classLoader);

        return bridgeClientRequest(
                clusterClientProvider,
                (clusterClient -> clusterClient.getAccumulators(jobID, classLoader)));
    }

    @Override
    public CompletableFuture getJobExecutionResult() {
        checkNotNull(classLoader);

        return bridgeClientRequest(
                clusterClientProvider,
                (clusterClient ->
                        clusterClient
                                .requestJobResult(jobID)
                                .thenApply(
                                        (jobResult) -> {
                                            try {
                                                return jobResult.toJobExecutionResult(classLoader);
                                            } catch (Throwable t) {
                                                throw new CompletionException(
                                                        new ProgramInvocationException(
                                                                "Job failed", jobID, t));
                                            }
                                        })));
    }

    @Override
    public CompletableFuture sendCoordinationRequest(
            OperatorID operatorId, CoordinationRequest request) {
        return bridgeClientRequest(
                clusterClientProvider,
                clusterClient -> clusterClient.sendCoordinationRequest(jobID, operatorId, request));
    }

    @Override
    public void reportHeartbeat(long expiredTimestamp) {
        bridgeClientRequest(
                clusterClientProvider,
                (clusterClient -> clusterClient.reportHeartbeat(jobID, expiredTimestamp)));
    }

    private static  CompletableFuture bridgeClientRequest(
            ClusterClientProvider clusterClientProvider,
            Function, CompletableFuture> resultRetriever) {

        ClusterClient clusterClient = clusterClientProvider.getClusterClient();

        CompletableFuture resultFuture;
        try {
            resultFuture = resultRetriever.apply(clusterClient);
        } catch (Throwable throwable) {
            IOUtils.closeQuietly(clusterClient::close);
            return FutureUtils.completedExceptionally(throwable);
        }

        return resultFuture.whenCompleteAsync(
                (jobResult, throwable) -> IOUtils.closeQuietly(clusterClient::close));
    }
}