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

import org.apache.flink.api.common.JobID;
import org.apache.flink.client.program.MiniClusterClient;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.JobManagerOptions;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.configuration.StateRecoveryOptions;
import org.apache.flink.configuration.TaskManagerOptions;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph;
import org.apache.flink.runtime.executiongraph.ArchivedExecutionVertex;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobGraphBuilder;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.minicluster.MiniCluster;
import org.apache.flink.runtime.minicluster.MiniClusterConfiguration;
import org.apache.flink.runtime.scheduler.DefaultScheduler;
import org.apache.flink.runtime.testutils.CommonTestUtils;
import org.apache.flink.runtime.testutils.WaitingCancelableInvokable;
import org.apache.flink.streaming.util.RestartStrategyUtils;
import org.apache.flink.util.TestLoggerExtension;

import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;

import java.time.Duration;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

import static org.apache.flink.configuration.JobManagerOptions.EXECUTION_FAILOVER_STRATEGY;
import static org.assertj.core.api.Assertions.assertThat;

/** IT case to test local recovery using {@link DefaultScheduler}. */
@ExtendWith(TestLoggerExtension.class)
class DefaultSchedulerLocalRecoveryITCase {

    private static final long TIMEOUT = 10_000L;

    @Test
    // The AdaptiveScheduler doesn't update the ExecutionGraph but creates a new Execution during
    // local recovery. Recovering can also lead to a change in parallelism which makes the
    // executionHistory non-linear. The lack of a linear executionHistory prevents us from applying
    // the same test for the AdaptiveScheduler.
    @Tag("org.apache.flink.testutils.junit.FailsWithAdaptiveScheduler")
    void testLocalRecoveryFull() throws Exception {
        testLocalRecoveryInternal("full");
    }

    @Test
    // see comment in #testLocalRecoveryFull
    @Tag("org.apache.flink.testutils.junit.FailsWithAdaptiveScheduler")
    void testLocalRecoveryRegion() throws Exception {
        testLocalRecoveryInternal("region");
    }

    private void testLocalRecoveryInternal(String failoverStrategyValue) throws Exception {
        final Configuration configuration = new Configuration();
        configuration.set(StateRecoveryOptions.LOCAL_RECOVERY, true);
        configuration.setString(EXECUTION_FAILOVER_STRATEGY.key(), failoverStrategyValue);

        final int parallelism = 10;
        final ArchivedExecutionGraph graph = executeSchedulingTest(configuration, parallelism);
        assertNonLocalRecoveredTasksEquals(graph, 1);
    }

    private void assertNonLocalRecoveredTasksEquals(ArchivedExecutionGraph graph, int expected) {
        int nonLocalRecoveredTasks = 0;
        for (ArchivedExecutionVertex vertex : graph.getAllExecutionVertices()) {
            int currentAttemptNumber = vertex.getCurrentExecutionAttempt().getAttemptNumber();
            if (currentAttemptNumber == 0) {
                // the task had never restarted and do not need to recover
                continue;
            }
            AllocationID priorAllocation =
                    vertex.getExecutionHistory()
                            .getHistoricalExecution(currentAttemptNumber - 1)
                            .get()
                            .getAssignedAllocationID();
            AllocationID currentAllocation =
                    vertex.getCurrentExecutionAttempt().getAssignedAllocationID();

            assertThat(priorAllocation).isNotNull();
            assertThat(currentAllocation).isNotNull();
            if (!currentAllocation.equals(priorAllocation)) {
                nonLocalRecoveredTasks++;
            }
        }
        assertThat(nonLocalRecoveredTasks).isEqualTo(expected);
    }

    private ArchivedExecutionGraph executeSchedulingTest(
            Configuration configuration, int parallelism) throws Exception {
        final Duration slotIdleTimeout = Duration.ofMillis(TIMEOUT);
        configuration.set(JobManagerOptions.SLOT_IDLE_TIMEOUT, slotIdleTimeout);

        configuration.set(TaskManagerOptions.TOTAL_FLINK_MEMORY, MemorySize.parse("64mb"));
        configuration.set(TaskManagerOptions.FRAMEWORK_HEAP_MEMORY, MemorySize.parse("16mb"));
        configuration.set(TaskManagerOptions.FRAMEWORK_OFF_HEAP_MEMORY, MemorySize.parse("16mb"));

        final MiniClusterConfiguration miniClusterConfiguration =
                new MiniClusterConfiguration.Builder()
                        .withRandomPorts()
                        .setConfiguration(configuration)
                        .setNumTaskManagers(parallelism)
                        .setNumSlotsPerTaskManager(1)
                        .build();

        try (MiniCluster miniCluster = new MiniCluster(miniClusterConfiguration)) {
            miniCluster.start();

            MiniClusterClient miniClusterClient = new MiniClusterClient(configuration, miniCluster);

            JobGraph jobGraph = createJobGraph(parallelism);

            // wait for the submission to succeed
            JobID jobId = miniClusterClient.submitJob(jobGraph).get(TIMEOUT, TimeUnit.SECONDS);

            // wait until all tasks running before triggering task failures
            waitUntilAllVerticesRunning(jobId, miniCluster);

            // kill one TM to trigger task failure and remove one existing slot
            CompletableFuture<Void> terminationFuture = miniCluster.terminateTaskManager(0);
            terminationFuture.get();

            // restart a taskmanager as a replacement for the killed one
            miniCluster.startTaskManager();

            // wait until all tasks running again
            waitUntilAllVerticesRunning(jobId, miniCluster);

            ArchivedExecutionGraph graph =
                    miniCluster.getArchivedExecutionGraph(jobGraph.getJobID()).get();

            miniCluster.cancelJob(jobId).get();

            return graph;
        }
    }

    private void waitUntilAllVerticesRunning(JobID jobId, MiniCluster miniCluster)
            throws Exception {
        CommonTestUtils.waitForAllTaskRunning(
                () -> miniCluster.getExecutionGraph(jobId).get(TIMEOUT, TimeUnit.SECONDS), false);
    }

    private JobGraph createJobGraph(int parallelism) {
        final JobVertex source = new JobVertex("v1");
        source.setInvokableClass(WaitingCancelableInvokable.class);
        source.setParallelism(parallelism);

        JobGraph jobGraph =
                JobGraphBuilder.newStreamingJobGraphBuilder()
                        .addJobVertices(List.of(source))
                        .build();

        RestartStrategyUtils.configureFixedDelayRestartStrategy(jobGraph, 1, 10L);

        return jobGraph;
    }
}
