/*
 * 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.runtime.checkpoint;

import org.apache.flink.runtime.concurrent.FutureUtils;
import org.apache.flink.runtime.operators.coordination.OperatorCoordinator;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;

import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.Executor;

/**
 * All the logic related to taking checkpoints of the {@link OperatorCoordinator}s.
 *
 * <p>NOTE: This class has a simplified error handling logic. If one of the several coordinator checkpoints
 * fail, no cleanup is triggered for the other concurrent ones. That is okay, since they all produce just byte[]
 * as the result. We have to change that once we allow then to create external resources that actually need
 * to be cleaned up.
 */
final class OperatorCoordinatorCheckpoints {

	public static CompletableFuture<CoordinatorSnapshot> triggerCoordinatorCheckpoint(
			final OperatorCoordinatorCheckpointContext coordinatorInfo,
			final long checkpointId) throws Exception {

		final CompletableFuture<byte[]> checkpointFuture =
				coordinatorInfo.coordinator().checkpointCoordinator(checkpointId);

		return checkpointFuture.thenApply(
				(state) -> new CoordinatorSnapshot(
						coordinatorInfo, new ByteStreamStateHandle(coordinatorInfo.operatorId().toString(), state))
		);
	}

	public static CompletableFuture<AllCoordinatorSnapshots> triggerAllCoordinatorCheckpoints(
			final Collection<OperatorCoordinatorCheckpointContext> coordinators,
			final long checkpointId) throws Exception {

		final Collection<CompletableFuture<CoordinatorSnapshot>> individualSnapshots = new ArrayList<>(coordinators.size());

		for (final OperatorCoordinatorCheckpointContext coordinator : coordinators) {
			individualSnapshots.add(triggerCoordinatorCheckpoint(coordinator, checkpointId));
		}

		return FutureUtils.combineAll(individualSnapshots).thenApply(AllCoordinatorSnapshots::new);
	}

	public static CompletableFuture<Void> triggerAndAcknowledgeAllCoordinatorCheckpoints(
			final Collection<OperatorCoordinatorCheckpointContext> coordinators,
			final PendingCheckpoint checkpoint,
			final Executor acknowledgeExecutor) throws Exception {

		final CompletableFuture<AllCoordinatorSnapshots> snapshots =
				triggerAllCoordinatorCheckpoints(coordinators, checkpoint.getCheckpointId());

		return snapshots
				.thenAcceptAsync(
						(allSnapshots) -> {
							try {
								acknowledgeAllCoordinators(checkpoint, allSnapshots.snapshots);
							}
							catch (Exception e) {
								throw new CompletionException(e);
							}
						},
						acknowledgeExecutor);
	}

	public static CompletableFuture<Void> triggerAndAcknowledgeAllCoordinatorCheckpointsWithCompletion(
			final Collection<OperatorCoordinatorCheckpointContext> coordinators,
			final PendingCheckpoint checkpoint,
			final Executor acknowledgeExecutor) throws CompletionException {

		try {
			return triggerAndAcknowledgeAllCoordinatorCheckpoints(coordinators, checkpoint, acknowledgeExecutor);
		} catch (Exception e) {
			throw new CompletionException(e);
		}
	}

	// ------------------------------------------------------------------------

	private static void acknowledgeAllCoordinators(PendingCheckpoint checkpoint, Collection<CoordinatorSnapshot> snapshots) throws CheckpointException {
		for (final CoordinatorSnapshot snapshot : snapshots) {
			final PendingCheckpoint.TaskAcknowledgeResult result =
				checkpoint.acknowledgeCoordinatorState(snapshot.coordinator, snapshot.state);

			if (result != PendingCheckpoint.TaskAcknowledgeResult.SUCCESS) {
				throw new CheckpointException("Coordinator state not acknowledged successfully: " + result,
					CheckpointFailureReason.TRIGGER_CHECKPOINT_FAILURE);
			}
		}
	}

	// ------------------------------------------------------------------------

	static final class AllCoordinatorSnapshots {

		private final Collection<CoordinatorSnapshot> snapshots;

		AllCoordinatorSnapshots(Collection<CoordinatorSnapshot> snapshots) {
			this.snapshots = snapshots;
		}

		public Iterable<CoordinatorSnapshot> snapshots() {
			return snapshots;
		}
	}

	static final class CoordinatorSnapshot {

		final OperatorCoordinatorCheckpointContext coordinator;
		final ByteStreamStateHandle state;

		CoordinatorSnapshot(OperatorCoordinatorCheckpointContext coordinator, ByteStreamStateHandle state) {
			this.coordinator = coordinator;
			this.state = state;
		}
	}
}

@@ -4486,38 +4486,36 @@
 %09%09%09%09
-throw new CheckpointException(
+final String errorMessage = 
 %22Coo
@@ -4570,23 +4570,307 @@
 + result
-,%0A%09%09%09%09%09
+;%0A%09%09%09%09final Throwable error = checkpoint.isDiscarded() ? checkpoint.getFailureCause() : null;%0A%0A%09%09%09%09if (error != null) %7B%0A%09%09%09%09%09throw new CheckpointException(errorMessage, CheckpointFailureReason.TRIGGER_CHECKPOINT_FAILURE, error);%0A%09%09%09%09%7D else %7B%0A%09%09%09%09%09throw new CheckpointException(errorMessage, 
 Checkpoi
@@ -4914,16 +4914,22 @@
 ILURE);%0A
+%09%09%09%09%7D%0A
 %09%09%09%7D%0A%09%09%7D

Improve exception in case Coordinator State ack fails