<commit_before>/*
 * 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.application;

import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.client.ClientUtils;
import org.apache.flink.client.cli.ClientOptions;
import org.apache.flink.client.deployment.application.executors.EmbeddedExecutor;
import org.apache.flink.client.deployment.application.executors.EmbeddedExecutorServiceLoader;
import org.apache.flink.client.program.PackagedProgram;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.PipelineOptionsInternal;
import org.apache.flink.core.execution.PipelineExecutorServiceLoader;
import org.apache.flink.runtime.clusterframework.ApplicationStatus;
import org.apache.flink.runtime.concurrent.FutureUtils;
import org.apache.flink.runtime.concurrent.ScheduledExecutor;
import org.apache.flink.runtime.dispatcher.DispatcherBootstrap;
import org.apache.flink.runtime.dispatcher.DispatcherGateway;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobmanager.HighAvailabilityMode;
import org.apache.flink.runtime.jobmaster.JobResult;
import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FlinkException;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

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

/**
 * A {@link DispatcherBootstrap} used for running the user's {@code main()} in "Application Mode" (see FLIP-85).
 *
 * <p>This dispatcher bootstrap submits the recovered {@link JobGraph job graphs} for re-execution
 * (in case of recovery from a failure), and then submits the remaining jobs of the application for execution.
 *
 * <p>To achieve this, it works in conjunction with the
 * {@link EmbeddedExecutor EmbeddedExecutor} which decides
 * if it should submit a job for execution (in case of a new job) or the job was already recovered and is running.
 */
@Internal
public class ApplicationDispatcherBootstrap implements DispatcherBootstrap {

	private static final Logger LOG = LoggerFactory.getLogger(ApplicationDispatcherBootstrap.class);

	public static final JobID ZERO_JOB_ID = new JobID(0, 0);

	private final PackagedProgram application;

	private final Collection<JobID> recoveredJobIds;

	private final Configuration configuration;

	private final FatalErrorHandler errorHandler;

	private final CompletableFuture<Void> applicationCompletionFuture;

	private final CompletableFuture<Acknowledge> clusterShutdownFuture;

	private ScheduledFuture<?> applicationExecutionTask;

	public ApplicationDispatcherBootstrap(
			final PackagedProgram application,
			final Collection<JobID> recoveredJobIds,
			final Configuration configuration,
			final DispatcherGateway dispatcherGateway,
			final ScheduledExecutor scheduledExecutor,
			final FatalErrorHandler errorHandler) {
		this.configuration = checkNotNull(configuration);
		this.recoveredJobIds = checkNotNull(recoveredJobIds);
		this.application = checkNotNull(application);
		this.errorHandler = checkNotNull(errorHandler);

		this.applicationCompletionFuture =
				fixJobIdAndRunApplicationAsync(dispatcherGateway, scheduledExecutor);

		this.clusterShutdownFuture =
				runApplicationAndShutdownClusterAsync(dispatcherGateway);
	}

	@Override
	public void stop() {
		if (applicationExecutionTask != null) {
			applicationExecutionTask.cancel(true);
		}

		if (applicationCompletionFuture != null) {
			applicationCompletionFuture.cancel(true);
		}
	}

	@VisibleForTesting
	ScheduledFuture<?> getApplicationExecutionFuture() {
		return applicationExecutionTask;
	}

	@VisibleForTesting
	CompletableFuture<Void> getApplicationCompletionFuture() {
		return applicationCompletionFuture;
	}

	@VisibleForTesting
	CompletableFuture<Acknowledge> getClusterShutdownFuture() {
		return clusterShutdownFuture;
	}

	/**
	 * Runs the user program entrypoint and shuts down the given dispatcherGateway when
	 * the application completes (either successfully or in case of failure).
	 */
	private CompletableFuture<Acknowledge> runApplicationAndShutdownClusterAsync(final DispatcherGateway dispatcherGateway) {
		return applicationCompletionFuture
				.handle((r, t) -> {
					if (t != null) {

						final Optional<ApplicationFailureException> exception =
								ExceptionUtils.findThrowable(t, ApplicationFailureException.class);

						if (exception.isPresent()) {
							final ApplicationStatus applicationStatus = exception.get().getStatus();

							if (applicationStatus == ApplicationStatus.CANCELED || applicationStatus == ApplicationStatus.FAILED) {
								LOG.info("Application {}: ", applicationStatus, t);
								return dispatcherGateway.shutDownCluster(applicationStatus);
							}
						}

						LOG.warn("Exiting with Application Status UNKNOWN: ", t);
						this.errorHandler.onFatalError(new FlinkException("Application failed unexpectedly.", t));

						return FutureUtils.<Acknowledge>completedExceptionally(t);
					}

					LOG.info("Application completed SUCCESSFULLY");
					return dispatcherGateway.shutDownCluster(ApplicationStatus.SUCCEEDED);
				})
				.thenCompose(Function.identity());
	}

	private CompletableFuture<Void> fixJobIdAndRunApplicationAsync(
			final DispatcherGateway dispatcherGateway,
			final ScheduledExecutor scheduledExecutor) {

		final Optional<String> configuredJobId =
				configuration.getOptional(PipelineOptionsInternal.PIPELINE_FIXED_JOB_ID);
<commit_msg>[hotfix] minor refactoring in ApplicationDispatcherBootstrap<commit_after>/*
 * 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.application;

import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.client.ClientUtils;
import org.apache.flink.client.cli.ClientOptions;
import org.apache.flink.client.deployment.application.executors.EmbeddedExecutor;
import org.apache.flink.client.deployment.application.executors.EmbeddedExecutorServiceLoader;
import org.apache.flink.client.program.PackagedProgram;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.PipelineOptionsInternal;
import org.apache.flink.core.execution.PipelineExecutorServiceLoader;
import org.apache.flink.runtime.clusterframework.ApplicationStatus;
import org.apache.flink.runtime.concurrent.FutureUtils;
import org.apache.flink.runtime.concurrent.ScheduledExecutor;
import org.apache.flink.runtime.dispatcher.DispatcherBootstrap;
import org.apache.flink.runtime.dispatcher.DispatcherGateway;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobmanager.HighAvailabilityMode;
import org.apache.flink.runtime.jobmaster.JobResult;
import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FlinkException;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

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

/**
 * A {@link DispatcherBootstrap} used for running the user's {@code main()} in "Application Mode" (see FLIP-85).
 *
 * <p>This dispatcher bootstrap submits the recovered {@link JobGraph job graphs} for re-execution
 * (in case of recovery from a failure), and then submits the remaining jobs of the application for execution.
 *
 * <p>To achieve this, it works in conjunction with the
 * {@link EmbeddedExecutor EmbeddedExecutor} which decides
 * if it should submit a job for execution (in case of a new job) or the job was already recovered and is running.
 */
@Internal
public class ApplicationDispatcherBootstrap implements DispatcherBootstrap {

	private static final Logger LOG = LoggerFactory.getLogger(ApplicationDispatcherBootstrap.class);

	public static final JobID ZERO_JOB_ID = new JobID(0, 0);

	private final PackagedProgram application;

	private final Collection<JobID> recoveredJobIds;

	private final Configuration configuration;

	private final FatalErrorHandler errorHandler;

	private final CompletableFuture<Void> applicationCompletionFuture;

	private final CompletableFuture<Acknowledge> clusterShutdownFuture;

	private ScheduledFuture<?> applicationExecutionTask;

	public ApplicationDispatcherBootstrap(
			final PackagedProgram application,
			final Collection<JobID> recoveredJobIds,
			final Configuration configuration,
			final DispatcherGateway dispatcherGateway,
			final ScheduledExecutor scheduledExecutor,
			final FatalErrorHandler errorHandler) {
		this.configuration = checkNotNull(configuration);
		this.recoveredJobIds = checkNotNull(recoveredJobIds);
		this.application = checkNotNull(application);
		this.errorHandler = checkNotNull(errorHandler);

		this.applicationCompletionFuture =
				fixJobIdAndRunApplicationAsync(dispatcherGateway, scheduledExecutor);

		this.clusterShutdownFuture =
				runApplicationAndShutdownClusterAsync(dispatcherGateway);
	}

	@Override
	public void stop() {
		if (applicationExecutionTask != null) {
			applicationExecutionTask.cancel(true);
		}

		if (applicationCompletionFuture != null) {
			applicationCompletionFuture.cancel(true);
		}
	}

	@VisibleForTesting
	ScheduledFuture<?> getApplicationExecutionFuture() {
		return applicationExecutionTask;
	}

	@VisibleForTesting
	CompletableFuture<Void> getApplicationCompletionFuture() {
		return applicationCompletionFuture;
	}

	@VisibleForTesting
	CompletableFuture<Acknowledge> getClusterShutdownFuture() {
		return clusterShutdownFuture;
	}

	/**
	 * Runs the user program entrypoint and shuts down the given dispatcherGateway when
	 * the application completes (either successfully or in case of failure).
	 */
	private CompletableFuture<Acknowledge> runApplicationAndShutdownClusterAsync(final DispatcherGateway dispatcherGateway) {
		return applicationCompletionFuture
				.handle((r, t) -> {

					if (t == null) {
						LOG.info("Application completed SUCCESSFULLY");
						return dispatcherGateway.shutDownCluster(ApplicationStatus.SUCCEEDED);
					}

					final Optional<ApplicationFailureException> exception =
							ExceptionUtils.findThrowable(t, ApplicationFailureException.class);

					if (exception.isPresent()) {
						final ApplicationStatus applicationStatus = exception.get().getStatus();

						if (applicationStatus == ApplicationStatus.CANCELED || applicationStatus == ApplicationStatus.FAILED) {
							LOG.info("Application {}: ", applicationStatus, t);
							return dispatcherGateway.shutDownCluster(applicationStatus);
						}
					}

					LOG.warn("Application failed unexpectedly: ", t);
					this.errorHandler.onFatalError(new FlinkException("Application failed unexpectedly.", t));

					return FutureUtils.<Acknowledge>completedExceptionally(t);
				})
				.thenCompose(Function.identity());
	}

	private CompletableFuture<Void> fixJobIdAndRunApplicationAsync(
			final DispatcherGateway dispatcherGateway,
			final ScheduledExecutor scheduledExecutor) {

		final Optional<String> configuredJobId =
				configuration.getOptional(PipelineOptionsInternal.PIPELINE_FIXED_JOB_ID);

{'fraction_non_alphanumeric': 0.08122941822173436, 'fraction_numerical': 0.0014635931211123307, 'mean_word_length': 5.579682233991334, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '23765', 'n_tokens_mistral': 3902, 'n_tokens_neox': 3563, 'n_words': 934}