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

import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.time.Deadline;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.configuration.CheckpointingOptions;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
import org.apache.flink.configuration.StateBackendOptions;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.client.JobStatusMessage;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
import org.apache.flink.runtime.testutils.TestingUtils;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.checkpoint.ListCheckpointed;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.test.util.MiniClusterWithClientResource;
import org.apache.flink.util.Collector;
import org.apache.flink.util.TestLogger;
import org.apache.flink.util.concurrent.FutureUtils;

import org.junit.AfterClass;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

import java.io.File;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static org.apache.flink.runtime.testutils.CommonTestUtils.waitForAllTaskRunning;
import static org.apache.flink.test.util.TestUtils.submitJobAndWaitForResult;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

/** Test savepoint rescaling. */
@RunWith(Parameterized.class)
public class RescalingITCase extends TestLogger {

    private static final int numTaskManagers = 2;
    private static final int slotsPerTaskManager = 2;
    private static final int numSlots = numTaskManagers * slotsPerTaskManager;

    @Parameterized.Parameters(name = "backend = {0}, buffersPerChannel = {1}")
    public static Collection<Object[]> data() {
        return Arrays.asList(
                new Object[][] {
                    {"filesystem", 2}, {"rocksdb", 0}, {"filesystem", 0}, {"rocksdb", 2}
                });
    }

    public RescalingITCase(String backend, int buffersPerChannel) {
        this.backend = backend;
        this.buffersPerChannel = buffersPerChannel;
    }

    private final String backend;

    private final int buffersPerChannel;

    private String currentBackend = null;

    enum OperatorCheckpointMethod {
        NON_PARTITIONED,
        CHECKPOINTED_FUNCTION,
        CHECKPOINTED_FUNCTION_BROADCAST,
        LIST_CHECKPOINTED
    }

    private static MiniClusterWithClientResource cluster;

    @ClassRule public static TemporaryFolder temporaryFolder = new TemporaryFolder();

    @Before
    public void setup() throws Exception {
        // detect parameter change
        if (currentBackend != backend) {
            shutDownExistingCluster();

            currentBackend = backend;

            Configuration config = new Configuration();

            final File checkpointDir = temporaryFolder.newFolder();
            final File savepointDir = temporaryFolder.newFolder();

            config.setString(StateBackendOptions.STATE_BACKEND, currentBackend);
            config.setString(
                    CheckpointingOptions.CHECKPOINTS_DIRECTORY, checkpointDir.toURI().toString());
            config.setString(
                    CheckpointingOptions.SAVEPOINT_DIRECTORY, savepointDir.toURI().toString());
            config.setInteger(
                    NettyShuffleEnvironmentOptions.NETWORK_BUFFERS_PER_CHANNEL, buffersPerChannel);

            cluster =
                    new MiniClusterWithClientResource(
                            new MiniClusterResourceConfiguration.Builder()
                                    .setConfiguration(config)
                                    .setNumberTaskManagers(numTaskManagers)
                                    .setNumberSlotsPerTaskManager(numSlots)
                                    .build());
            cluster.before();
        }
    }

    @AfterClass
    public static void shutDownExistingCluster() {
        if (cluster != null) {
            cluster.after();
            cluster = null;
        }
    }

    @Test
    public void testSavepointRescalingInKeyedState() throws Exception {
        testSavepointRescalingKeyedState(false, false);
    }

    @Test
    public void testSavepointRescalingOutKeyedState() throws Exception {
        testSavepointRescalingKeyedState(true, false);
    }

    @Test
    public void testSavepointRescalingInKeyedStateDerivedMaxParallelism() throws Exception {
        testSavepointRescalingKeyedState(false, true);
    }

    @Test
    public void testSavepointRescalingOutKeyedStateDerivedMaxParallelism() throws Exception {
        testSavepointRescalingKeyedState(true, true);
    }

    /**
     * Tests that a job with purely keyed state can be restarted from a savepoint with a different
     * parallelism.
     */
    public void testSavepointRescalingKeyedState(boolean scaleOut, boolean deriveMaxParallelism)
            throws Exception {
        final int numberKeys = 42;
        final int numberElements = 1000;
        final int numberElements2 = 500;
        final int parallelism = scaleOut ? numSlots / 2 : numSlots;
        final int parallelism2 = scaleOut ? numSlots : numSlots / 2;
        final int maxParallelism = 13;

        Duration timeout = Duration.ofMinutes(3);
        Deadline deadline = Deadline.now().plus(timeout);

        ClusterClient<?> client = cluster.getClusterClient();

        try {
            JobGraph jobGraph =
                    createJobGraphWithKeyedState(
                            parallelism, maxParallelism, numberKeys, numberElements, false, 100);

            final JobID jobID = jobGraph.getJobID();

            client.submitJob(jobGraph).get();

            // wait til the sources have emitted numberElements for each key and completed a
            // checkpoint
            assertTrue(
                    SubtaskIndexFlatMapper.workCompletedLatch.await(
                            deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS));

            // verify the current state

            Set<Tuple2<Integer, Integer>> actualResult = CollectionSink.getElementsSet();

            Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();

            for (int key = 0; key < numberKeys; key++) {
                int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);

                expectedResult.add(
                        Tuple2.of(
                                KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
                                        maxParallelism, parallelism, keyGroupIndex),
                                numberElements * key));
            }

            assertEquals(expectedResult, actualResult);

            // clear the CollectionSink set for the restarted job
            CollectionSink.clearElementsSet();

            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID());
            CompletableFuture<String> savepointPathFuture = client.triggerSavepoint(jobID, null);

            final String savepointPath =
                    savepointPathFuture.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);

            client.cancel(jobID).get();

            while (!getRunningJobs(client).isEmpty()) {
                Thread.sleep(50);
            }

            int restoreMaxParallelism =
                    deriveMaxParallelism ? JobVertex.MAX_PARALLELISM_DEFAULT : maxParallelism;

            JobGraph scaledJobGraph =
                    createJobGraphWithKeyedState(
                            parallelism2,
                            restoreMaxParallelism,
                            numberKeys,
                            numberElements2,
                            true,
                            100);

            scaledJobGraph.setSavepointRestoreSettings(
                    SavepointRestoreSettings.forPath(savepointPath));

            submitJobAndWaitForResult(client, scaledJobGraph, getClass().getClassLoader());

            Set<Tuple2<Integer, Integer>> actualResult2 = CollectionSink.getElementsSet();

            Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();

            for (int key = 0; key < numberKeys; key++) {
                int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
                expectedResult2.add(
                        Tuple2.of(
                                KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
                                        maxParallelism, parallelism2, keyGroupIndex),
                                key * (numberElements + numberElements2)));
            }

            assertEquals(expectedResult2, actualResult2);

        } finally {
            // clear the CollectionSink set for the restarted job
            CollectionSink.clearElementsSet();
        }
    }

    /**
     * Tests that a job cannot be restarted from a savepoint with a different parallelism if the
     * rescaled operator has non-partitioned state.
     *
     * @throws Exception
     */
    @Test
    public void testSavepointRescalingNonPartitionedStateCausesException() throws Exception {
        final int parallelism = numSlots / 2;
        final int parallelism2 = numSlots;
        final int maxParallelism = 13;

        Duration timeout = Duration.ofMinutes(3);
        Deadline deadline = Deadline.now().plus(timeout);

        ClusterClient<?> client = cluster.getClusterClient();

        try {
            JobGraph jobGraph =
                    createJobGraphWithOperatorState(
                            parallelism, maxParallelism, OperatorCheckpointMethod.NON_PARTITIONED);
            // make sure the job does not finish before we take the savepoint
            StateSourceBase.canFinishLatch = new CountDownLatch(1);

            final JobID jobID = jobGraph.getJobID();

            client.submitJob(jobGraph).get();

            // wait until the operator is started
            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID());
            // wait until the operator handles some data
            StateSourceBase.workStartedLatch.await();

            CompletableFuture<String> savepointPathFuture = client.triggerSavepoint(jobID, null);

            final String savepointPath =
                    savepointPathFuture.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);

            // we took a savepoint, the job can finish now
            StateSourceBase.canFinishLatch.countDown();
            client.cancel(jobID).get();
            while (!getRunningJobs(client).isEmpty()) {
                Thread.sleep(50);
            }

            // job successfully removed
            JobGraph scaledJobGraph =
                    createJobGraphWithOperatorState(
                            parallelism2, maxParallelism, OperatorCheckpointMethod.NON_PARTITIONED);

            scaledJobGraph.setSavepointRestoreSettings(
                    SavepointRestoreSettings.forPath(savepointPath));

            submitJobAndWaitForResult(client, scaledJobGraph, getClass().getClassLoader());
        } catch (JobExecutionException exception) {
            if (exception.getCause() instanceof IllegalStateException) {
                // we expect a IllegalStateException wrapped
                // in a JobExecutionException, because the job containing non-partitioned state
                // is being rescaled
            } else {
                throw exception;
            }
        }
    }

    /**
     * Tests that a job with non partitioned state can be restarted from a savepoint with a
     * different parallelism if the operator with non-partitioned state are not rescaled.
     *
     * @throws Exception
     */
    @Test
    public void testSavepointRescalingWithKeyedAndNonPartitionedState() throws Exception {
        int numberKeys = 42;
        int numberElements = 1000;
        int numberElements2 = 500;
        int parallelism = numSlots / 2;
        int parallelism2 = numSlots;
        int maxParallelism = 13;

        Duration timeout = Duration.ofMinutes(3);
        Deadline deadline = Deadline.now().plus(timeout);

        ClusterClient<?> client = cluster.getClusterClient();

        try {

            JobGraph jobGraph =
                    createJobGraphWithKeyedAndNonPartitionedOperatorState(
                            parallelism,
                            maxParallelism,
                            parallelism,
                            numberKeys,
                            numberElements,
                            false,
                            100);

            final JobID jobID = jobGraph.getJobID();

            // make sure the job does not finish before we take the savepoint
            StateSourceBase.canFinishLatch = new CountDownLatch(1);
            client.submitJob(jobGraph).get();

            // wait til the sources have emitted numberElements for each key and completed a
            // checkpoint
            assertTrue(
                    SubtaskIndexFlatMapper.workCompletedLatch.await(
                            deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS));

            // verify the current state

            Set<Tuple2<Integer, Integer>> actualResult = CollectionSink.getElementsSet();

            Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();

            for (int key = 0; key < numberKeys; key++) {
                int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);

                expectedResult.add(
                        Tuple2.of(
                                KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
                                        maxParallelism, parallelism, keyGroupIndex),
                                numberElements * key));
            }

            assertEquals(expectedResult, actualResult);

            // clear the CollectionSink set for the restarted job
            CollectionSink.clearElementsSet();

            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID());
            CompletableFuture<String> savepointPathFuture = client.triggerSavepoint(jobID, null);

            final String savepointPath =
                    savepointPathFuture.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);

            // we took a savepoint, the job can finish now
            StateSourceBase.canFinishLatch.countDown();
            client.cancel(jobID).get();

            while (!getRunningJobs(client).isEmpty()) {
                Thread.sleep(50);
            }

            JobGraph scaledJobGraph =
                    createJobGraphWithKeyedAndNonPartitionedOperatorState(
                            parallelism2,
                            maxParallelism,
                            parallelism,
                            numberKeys,
                            numberElements + numberElements2,
                            true,
                            100);

            scaledJobGraph.setSavepointRestoreSettings(
                    SavepointRestoreSettings.forPath(savepointPath));

            submitJobAndWaitForResult(client, scaledJobGraph, getClass().getClassLoader());

            Set<Tuple2<Integer, Integer>> actualResult2 = CollectionSink.getElementsSet();

            Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();

            for (int key = 0; key < numberKeys; key++) {
                int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
                expectedResult2.add(
                        Tuple2.of(
                                KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
                                        maxParallelism, parallelism2, keyGroupIndex),
                                key * (numberElements + numberElements2)));
            }

            assertEquals(expectedResult2, actualResult2);

        } finally {
            // clear the CollectionSink set for the restarted job
            CollectionSink.clearElementsSet();
        }
    }

    @Test
    public void testSavepointRescalingInPartitionedOperatorState() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                false, OperatorCheckpointMethod.CHECKPOINTED_FUNCTION);
    }

    @Test
    public void testSavepointRescalingOutPartitionedOperatorState() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                true, OperatorCheckpointMethod.CHECKPOINTED_FUNCTION);
    }

    @Test
    public void testSavepointRescalingInBroadcastOperatorState() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                false, OperatorCheckpointMethod.CHECKPOINTED_FUNCTION_BROADCAST);
    }

    @Test
    public void testSavepointRescalingOutBroadcastOperatorState() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                true, OperatorCheckpointMethod.CHECKPOINTED_FUNCTION_BROADCAST);
    }

    @Test
    public void testSavepointRescalingInPartitionedOperatorStateList() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                false, OperatorCheckpointMethod.LIST_CHECKPOINTED);
    }

    @Test
    public void testSavepointRescalingOutPartitionedOperatorStateList() throws Exception {
        testSavepointRescalingPartitionedOperatorState(
                true, OperatorCheckpointMethod.LIST_CHECKPOINTED);
    }

    /**
     * Tests rescaling of partitioned operator state. More specific, we test the mechanism with
     * {@link ListCheckpointed} as it subsumes {@link
     * org.apache.flink.streaming.api.checkpoint.CheckpointedFunction}.
     */
    public void testSavepointRescalingPartitionedOperatorState(
            boolean scaleOut, OperatorCheckpointMethod checkpointMethod) throws Exception {
        final int parallelism = scaleOut ? numSlots : numSlots / 2;
        final int parallelism2 = scaleOut ? numSlots / 2 : numSlots;
        final int maxParallelism = 13;

        Duration timeout = Duration.ofMinutes(3);
        Deadline deadline = Deadline.now().plus(timeout);

        ClusterClient<?> client = cluster.getClusterClient();

        int counterSize = Math.max(parallelism, parallelism2);

        if (checkpointMethod == OperatorCheckpointMethod.CHECKPOINTED_FUNCTION
                || checkpointMethod == OperatorCheckpointMethod.CHECKPOINTED_FUNCTION_BROADCAST) {
            PartitionedStateSource.checkCorrectSnapshot = new int[counterSize];
            PartitionedStateSource.checkCorrectRestore = new int[counterSize];
        } else {
            PartitionedStateSourceListCheckpointed.checkCorrectSnapshot = new int[counterSize];
            PartitionedStateSourceListCheckpointed.checkCorrectRestore = new int[counterSize];
        }

        try {
            JobGraph jobGraph =
                    createJobGraphWithOperatorState(parallelism, maxParallelism, checkpointMethod);
            // make sure the job does not finish before we take the savepoint
            StateSourceBase.canFinishLatch = new CountDownLatch(1);

            final JobID jobID = jobGraph.getJobID();

            client.submitJob(jobGraph).get();

            // wait until the operator is started
            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID());
            // wait until the operator handles some data
            StateSourceBase.workStartedLatch.await();

            CompletableFuture<String> savepointPathFuture =
                    FutureUtils.retryWithDelay(
                            () -> client.triggerSavepoint(jobID, null),
                            (int) deadline.timeLeft().getSeconds() / 10,
                            Time.seconds(10),
                            (throwable) -> true,
                            TestingUtils.defaultScheduledExecutor());

            final String savepointPath =
                    savepointPathFuture.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);

            // we took a savepoint, the job can finish now
            StateSourceBase.canFinishLatch.countDown();
            client.cancel(jobID).get();
            while (!getRunningJobs(client).isEmpty()) {
                Thread.sleep(50);
            }

            JobGraph scaledJobGraph =
                    createJobGraphWithOperatorState(parallelism2, maxParallelism, checkpointMethod);

            scaledJobGraph.setSavepointRestoreSettings(
                    SavepointRestoreSettings.forPath(savepointPath));

            submitJobAndWaitForResult(client, scaledJobGraph, getClass().getClassLoader());

            int sumExp = 0;
            int sumAct = 0;

            if (checkpointMethod == OperatorCheckpointMethod.CHECKPOINTED_FUNCTION) {
                for (int c : PartitionedStateSource.checkCorrectSnapshot) {
                    sumExp += c;
                }

                for (int c : PartitionedStateSource.checkCorrectRestore) {
                    sumAct += c;
                }
            } else if (checkpointMethod
                    == OperatorCheckpointMethod.CHECKPOINTED_FUNCTION_BROADCAST) {
                for (int c : PartitionedStateSource.checkCorrectSnapshot) {
                    sumExp += c;
                }

                for (int c : PartitionedStateSource.checkCorrectRestore) {
                    sumAct += c;
                }

                sumExp *= parallelism2;
            } else {
                for (int c : PartitionedStateSourceListCheckpointed.checkCorrectSnapshot) {
                    sumExp += c;
                }

                for (int c : PartitionedStateSourceListCheckpointed.checkCorrectRestore) {
                    sumAct += c;
                }
            }

            assertEquals(sumExp, sumAct);
        } finally {
        }
    }

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

    private static JobGraph createJobGraphWithOperatorState(
            int parallelism, int maxParallelism, OperatorCheckpointMethod checkpointMethod) {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        env.getConfig().setMaxParallelism(maxParallelism);
        env.enableCheckpointing(Long.MAX_VALUE);
        env.setRestartStrategy(RestartStrategies.noRestart());

        StateSourceBase.workStartedLatch = new CountDownLatch(parallelism);

        SourceFunction<Integer> src;

        switch (checkpointMethod) {
            case CHECKPOINTED_FUNCTION:
                src = new PartitionedStateSource(false);
                break;
            case CHECKPOINTED_FUNCTION_BROADCAST:
                src = new PartitionedStateSource(true);
                break;
            case LIST_CHECKPOINTED:
                src = new PartitionedStateSourceListCheckpointed();
                break;
            case NON_PARTITIONED:
                src = new NonPartitionedStateSource();
                break;
            default:
                throw new IllegalArgumentException();
        }

        DataStream<Integer> input = env.addSource(src);

        input.addSink(new DiscardingSink<Integer>());

        return env.getStreamGraph().getJobGraph();
    }

    private static JobGraph createJobGraphWithKeyedState(
            int parallelism,
            int maxParallelism,
            int numberKeys,
            int numberElements,
            boolean terminateAfterEmission,
            int checkpointingInterval) {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        if (0 < maxParallelism) {
            env.getConfig().setMaxParallelism(maxParallelism);
        }
        env.enableCheckpointing(checkpointingInterval);
        env.setRestartStrategy(RestartStrategies.noRestart());
        env.getConfig().setUseSnapshotCompression(true);

        DataStream<Integer> input =
                env.addSource(
                                new SubtaskIndexSource(
                                        numberKeys, numberElements, terminateAfterEmission))
                        .keyBy(
                                new KeySelector<Integer, Integer>() {
                                    private static final long serialVersionUID =
                                            -7952298871120320940L;

                                    @Override
                                    public Integer getKey(Integer value) throws Exception {
                                        return value;
                                    }
                                });

        SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);

        DataStream<Tuple2<Integer, Integer>> result =
                input.flatMap(new SubtaskIndexFlatMapper(numberElements));

        result.addSink(new CollectionSink<Tuple2<Integer, Integer>>());

        return env.getStreamGraph().getJobGraph();
    }

    private static JobGraph createJobGraphWithKeyedAndNonPartitionedOperatorState(
            int parallelism,
            int maxParallelism,
            int fixedParallelism,
            int numberKeys,
            int numberElements,
            boolean terminateAfterEmission,
            int checkpointingInterval) {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        env.getConfig().setMaxParallelism(maxParallelism);
        env.enableCheckpointing(checkpointingInterval);
        env.setRestartStrategy(RestartStrategies.noRestart());

        DataStream<Integer> input =
                env.addSource(
                                new SubtaskIndexNonPartitionedStateSource(
                                        numberKeys, numberElements, terminateAfterEmission))
                        .setParallelism(fixedParallelism)
                        .keyBy(
                                new KeySelector<Integer, Integer>() {
                                    private static final long serialVersionUID =
                                            -7952298871120320940L;

                                    @Override
                                    public Integer getKey(Integer value) throws Exception {
                                        return value;
                                    }
                                });

        SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);

        DataStream<Tuple2<Integer, Integer>> result =
                input.flatMap(new SubtaskIndexFlatMapper(numberElements));

        result.addSink(new CollectionSink<Tuple2<Integer, Integer>>());

        return env.getStreamGraph().getJobGraph();
    }

    private static class SubtaskIndexSource extends RichParallelSourceFunction<Integer> {

        private static final long serialVersionUID = -400066323594122516L;

        private final int numberKeys;
        private final int numberElements;
        private final boolean terminateAfterEmission;

        protected int counter = 0;

        private boolean running = true;

        SubtaskIndexSource(int numberKeys, int numberElements, boolean terminateAfterEmission) {

            this.numberKeys = numberKeys;
            this.numberElements = numberElements;
            this.terminateAfterEmission = terminateAfterEmission;
        }

        @Override
        public void run(SourceContext<Integer> ctx) throws Exception {
            final Object lock = ctx.getCheckpointLock();
            final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();

            while (running) {

                if (counter < numberElements) {
                    synchronized (lock) {
                        for (int value = subtaskIndex;
                                value < numberKeys;
                                value += getRuntimeContext().getNumberOfParallelSubtasks()) {

                            ctx.collect(value);
                        }

                        counter++;
                    }
                } else {
                    if (terminateAfterEmission) {
                        running = false;
                    } else {
                        Thread.sleep(100);
                    }
                }
            }
        }

        @Override
        public void cancel() {
            running = false;
        }
    }

    private static class SubtaskIndexNonPartitionedStateSource extends SubtaskIndexSource
            implements ListCheckpointed<Integer> {

        private static final long serialVersionUID = 8388073059042040203L;

        SubtaskIndexNonPartitionedStateSource(
                int numberKeys, int numberElements, boolean terminateAfterEmission) {
            super(numberKeys, numberElements, terminateAfterEmission);
        }

        @Override
        public List<Integer> snapshotState(long checkpointId, long timestamp) throws Exception {
            return Collections.singletonList(this.counter);
        }

        @Override
        public void restoreState(List<Integer> state) throws Exception {
            if (state.isEmpty() || state.size() > 1) {
                throw new RuntimeException(
                        "Test failed due to unexpected recovered state size " + state.size());
            }
            this.counter = state.get(0);
        }
    }

    private static class SubtaskIndexFlatMapper
            extends RichFlatMapFunction<Integer, Tuple2<Integer, Integer>>
            implements CheckpointedFunction {

        private static final long serialVersionUID = 5273172591283191348L;

        private static CountDownLatch workCompletedLatch = new CountDownLatch(1);

        private transient ValueState<Integer> counter;
        private transient ValueState<Integer> sum;

        private final int numberElements;

        SubtaskIndexFlatMapper(int numberElements) {
            this.numberElements = numberElements;
        }

        @Override
        public void flatMap(Integer value, Collector<Tuple2<Integer, Integer>> out)
                throws Exception {

            int count = counter.value() + 1;
            counter.update(count);

            int s = sum.value() + value;
            sum.update(s);

            if (count % numberElements == 0) {
                out.collect(Tuple2.of(getRuntimeContext().getIndexOfThisSubtask(), s));
                workCompletedLatch.countDown();
            }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // all managed, nothing to do.
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            counter =
                    context.getKeyedStateStore()
                            .getState(new ValueStateDescriptor<>("counter", Integer.class, 0));
            sum =
                    context.getKeyedStateStore()
                            .getState(new ValueStateDescriptor<>("sum", Integer.class, 0));
        }
    }

    private static class CollectionSink<IN> implements SinkFunction<IN> {

        private static Set<Object> elements =
                Collections.newSetFromMap(new ConcurrentHashMap<Object, Boolean>());

        private static final long serialVersionUID = -1652452958040267745L;

        public static <IN> Set<IN> getElementsSet() {
            return (Set<IN>) elements;
        }

        public static void clearElementsSet() {
            elements.clear();
        }

        @Override
        public void invoke(IN value) throws Exception {
            elements.add(value);
        }
    }

    private static class StateSourceBase extends RichParallelSourceFunction<Integer> {

        private static final long serialVersionUID = 7512206069681177940L;
        private static CountDownLatch workStartedLatch = new CountDownLatch(1);
        private static CountDownLatch canFinishLatch = new CountDownLatch(0);

        protected volatile int counter = 0;
        protected volatile boolean running = true;

        @Override
        public void run(SourceContext<Integer> ctx) throws Exception {
            final Object lock = ctx.getCheckpointLock();

            while (running) {
                synchronized (lock) {
                    ++counter;
                    ctx.collect(1);
                }

                Thread.sleep(2);

                if (counter == 10) {
                    workStartedLatch.countDown();
                }

                if (counter >= 500) {
                    break;
                }
            }

            canFinishLatch.await();
        }

        @Override
        public void cancel() {
            running = false;
        }
    }

    private static class NonPartitionedStateSource extends StateSourceBase
            implements ListCheckpointed<Integer> {

        private static final long serialVersionUID = -8108185918123186841L;

        @Override
        public List<Integer> snapshotState(long checkpointId, long timestamp) throws Exception {
            return Collections.singletonList(this.counter);
        }

        @Override
        public void restoreState(List<Integer> state) throws Exception {
            if (!state.isEmpty()) {
                this.counter = state.get(0);
            }
        }
    }

    private static class PartitionedStateSourceListCheckpointed extends StateSourceBase
            implements ListCheckpointed<Integer> {

        private static final long serialVersionUID = -4357864582992546L;
        private static final int NUM_PARTITIONS = 7;

        private static int[] checkCorrectSnapshot;
        private static int[] checkCorrectRestore;

        @Override
        public List<Integer> snapshotState(long checkpointId, long timestamp) throws Exception {

            checkCorrectSnapshot[getRuntimeContext().getIndexOfThisSubtask()] = counter;

            int div = counter / NUM_PARTITIONS;
            int mod = counter % NUM_PARTITIONS;

            List<Integer> split = new ArrayList<>();
            for (int i = 0; i < NUM_PARTITIONS; ++i) {
                int partitionValue = div;
                if (mod > 0) {
                    --mod;
                    ++partitionValue;
                }
                split.add(partitionValue);
            }
            return split;
        }

        @Override
        public void restoreState(List<Integer> state) throws Exception {
            for (Integer v : state) {
                counter += v;
            }
            checkCorrectRestore[getRuntimeContext().getIndexOfThisSubtask()] = counter;
        }
    }

    private static class PartitionedStateSource extends StateSourceBase
            implements CheckpointedFunction {

        private static final long serialVersionUID = -359715965103593462L;
        private static final int NUM_PARTITIONS = 7;

        private transient ListState<Integer> counterPartitions;
        private boolean broadcast;

        private static int[] checkCorrectSnapshot;
        private static int[] checkCorrectRestore;

        public PartitionedStateSource(boolean broadcast) {
            this.broadcast = broadcast;
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {

            counterPartitions.clear();

            checkCorrectSnapshot[getRuntimeContext().getIndexOfThisSubtask()] = counter;

            int div = counter / NUM_PARTITIONS;
            int mod = counter % NUM_PARTITIONS;

            for (int i = 0; i < NUM_PARTITIONS; ++i) {
                int partitionValue = div;
                if (mod > 0) {
                    --mod;
                    ++partitionValue;
                }
                counterPartitions.add(partitionValue);
            }
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {

            if (broadcast) {
                this.counterPartitions =
                        context.getOperatorStateStore()
                                .getUnionListState(
                                        new ListStateDescriptor<>(
                                                "counter_partitions", IntSerializer.INSTANCE));
            } else {
                this.counterPartitions =
                        context.getOperatorStateStore()
                                .getListState(
                                        new ListStateDescriptor<>(
                                                "counter_partitions", IntSerializer.INSTANCE));
            }

            if (context.isRestored()) {
                for (int v : counterPartitions.get()) {
                    counter += v;
                }
                checkCorrectRestore[getRuntimeContext().getIndexOfThisSubtask()] = counter;
            }
        }
    }

    private static List<JobID> getRunningJobs(ClusterClient<?> client) throws Exception {
        Collection<JobStatusMessage> statusMessages = client.listJobs().get();
        return statusMessages.stream()
                .filter(status -> !status.getJobState().isGloballyTerminalState())
                .map(JobStatusMessage::getJobId)
                .collect(Collectors.toList());
    }
}