/**
 * Licensed 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.streaming.connectors.pulsar;

import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.scala.typeutils.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.runtime.client.JobCancellationException;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobStatus;
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.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.graph.StreamingJobGraphGenerator;
import org.apache.flink.streaming.api.operators.StreamSink;
import org.apache.flink.streaming.connectors.pulsar.internal.AvroDeser;
import org.apache.flink.streaming.connectors.pulsar.internal.CachedPulsarClient;
import org.apache.flink.streaming.connectors.pulsar.internal.JsonDeser;
import org.apache.flink.streaming.connectors.pulsar.internal.ReaderThread;
import org.apache.flink.streaming.connectors.pulsar.testutils.FailingIdentityMapper;
import org.apache.flink.streaming.connectors.pulsar.testutils.SingletonStreamSink;
import org.apache.flink.streaming.connectors.pulsar.testutils.ValidatingExactlyOnceSink;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.utils.LegacyTypeInfoDataTypeConverter;
import org.apache.flink.test.util.SuccessException;
import org.apache.flink.test.util.TestUtils;
import org.apache.flink.testutils.junit.RetryRule;
import org.apache.flink.types.Row;
import org.apache.flink.util.Collector;
import org.apache.flink.util.ExceptionUtils;

import org.apache.flink.shaded.guava18.com.google.common.collect.Iterables;
import org.apache.flink.shaded.guava18.com.google.common.collect.Sets;

import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.admin.LongRunningProcessStatus;
import org.apache.pulsar.client.admin.LongRunningProcessStatus.Status;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.PulsarAdminException;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Reader;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.SubscriptionType;
import org.apache.pulsar.client.impl.MessageIdImpl;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.policies.data.PersistentTopicInternalStats;
import org.apache.pulsar.common.schema.SchemaType;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.flink.streaming.connectors.pulsar.SchemaData.fooList;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.CLIENT_CACHE_SIZE_OPTION_KEY;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.FAIL_ON_WRITE_OPTION_KEY;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.FLUSH_ON_CHECKPOINT_OPTION_KEY;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.PARTITION_DISCOVERY_INTERVAL_MS_OPTION_KEY;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.PULSAR_READER_OPTION_KEY_PREFIX;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.TOPIC_ATTRIBUTE_NAME;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.TOPIC_MULTI_OPTION_KEY;
import static org.apache.flink.streaming.connectors.pulsar.internal.PulsarOptions.TOPIC_SINGLE_OPTION_KEY;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/**
 * Pulsar source sink integration tests.
 */
public class FlinkPulsarITest extends PulsarTestBaseWithFlink {

    private static final Logger log = LoggerFactory.getLogger(FlinkPulsarITest.class);

    @Rule
    public RetryRule retryRule = new RetryRule();

	@Before
	public void clearState() {
		SingletonStreamSink.clear();
		FailingIdentityMapper.failedBefore = false;
	}

    @Test
    public void testRunFailedOnWrongServiceUrl() {

        try {
            Properties props = MapUtils.toProperties(Collections.singletonMap(TOPIC_SINGLE_OPTION_KEY, "tp"));

            StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
            see.getConfig().disableSysoutLogging();
            see.setRestartStrategy(RestartStrategies.noRestart());
            see.setParallelism(1);

            FlinkPulsarSource<String> source =
                    new FlinkPulsarSource<String>("sev", "admin", new SimpleStringSchema(), props).setStartFromEarliest();

            DataStream<String> stream = see.addSource(source);
            stream.print();
            see.execute("wrong service url");
        } catch (Exception e) {
            final Optional<Throwable> optionalThrowable = ExceptionUtils.findThrowableWithMessage(e, "authority component is missing");
            assertTrue(optionalThrowable.isPresent());
            assertTrue(optionalThrowable.get() instanceof PulsarClientException);
        }
    }

    @Test
    public void testCaseSensitiveReaderConf() throws Exception {
        String tp = newTopic();
        List<Integer> messages =
                IntStream.range(0, 50).mapToObj(t -> Integer.valueOf(t)).collect(Collectors.toList());

        sendTypedMessages(tp, SchemaType.INT32, messages, Optional.empty());
        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();

        Properties props = MapUtils.toProperties(Collections.singletonMap(TOPIC_SINGLE_OPTION_KEY, tp));
        props.setProperty("pulsar.reader.receiverQueueSize", "1000000");

        FlinkPulsarSource<Integer> source =
                new FlinkPulsarSource<Integer>(serviceUrl, adminUrl, new IntegerDeserializer(), props)
                        .setStartFromEarliest();

        DataStream<Integer> stream = see.addSource(source);
        stream.addSink(new DiscardingSink<Integer>());

        final AtomicReference<Throwable> jobError = new AtomicReference<>();
        JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(see.getStreamGraph());
        JobID jobID = jobGraph.getJobID();

        final Thread runner = new Thread("runner") {
            @Override
            public void run() {
                try {
                    client.setDetached(false);
                    client.submitJob(jobGraph, this.getClass().getClassLoader());
                } catch (Throwable t) {
                    if (!(t instanceof JobCancellationException)) {
                        jobError.set(t);
                    }
                }
            }
        };
        runner.start();

        Thread.sleep(2000);
        Throwable failureCause = jobError.get();

        if (failureCause != null) {
            failureCause.printStackTrace();
            fail("Test failed prematurely with: " + failureCause.getMessage());
        }

        client.cancel(jobID);
        runner.join();

        assertEquals(client.getJobStatus(jobID).get(), JobStatus.CANCELED);
    }

    @Test
    public void testCaseSensitiveProducerConf() throws Exception {
        int numTopic = 5;
        int numElements = 20;

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(1);

        List<String> topics = new ArrayList<>();
        for (int i = 0; i < numTopic; i++) {
            topics.add(newTopic());
        }

        DataStream<Row> stream = see.addSource(new MultiTopicSource(topics, numElements));

        Properties sinkProp = sinkProperties();
        sinkProp.setProperty("pulsar.producer.blockIfQueueFull", "true");
        sinkProp.setProperty("pulsar.producer.maxPendingMessages", "100000");
        sinkProp.setProperty("pulsar.producer.maxPendingMessagesAcrossPartitions", "5000000");
        sinkProp.setProperty("pulsar.producer.sendTimeoutMs", "30000");
        produceIntoPulsar(stream, intRowWithTopicType(), sinkProp);
        see.execute("write with topics");
    }

    @Test
    public void testClientCacheParameterPassedToTasks() throws Exception {
        int numTopic = 5;
        int numElements = 20;

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(3);

        List<String> topics = new ArrayList<>();
        for (int i = 0; i < numTopic; i++) {
            topics.add(newTopic());
        }

        DataStream<Row> stream = see.addSource(new MultiTopicSource(topics, numElements));

        Properties sinkProp = sinkProperties();
        sinkProp.setProperty(FLUSH_ON_CHECKPOINT_OPTION_KEY, "true");
        sinkProp.setProperty(CLIENT_CACHE_SIZE_OPTION_KEY, "7");
        stream.addSink(new AssertSink(serviceUrl, adminUrl, 7, sinkProp, intRowWithTopicType()));
        see.execute("write with topics");
    }

    @Test
    public void testProduceConsumeMultipleTopics() throws Exception {
        int numTopic = 5;
        int numElements = 20;

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(1);

        List<String> topics = new ArrayList<>();
        for (int i = 0; i < numTopic; i++) {
            topics.add(newTopic());
        }

        DataStream<Row> stream = see.addSource(new MultiTopicSource(topics, numElements));

        Properties sinkProp = sinkProperties();
        produceIntoPulsar(stream, intRowWithTopicType(), sinkProp);
        see.execute("write with topics");

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.getConfig().disableSysoutLogging();
        Properties sourceProp = sourceProperties();
        sourceProp.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(topics.toArray(), ','));
        DataStream<Row> stream1 = env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProp).setStartFromEarliest());

        stream1.flatMap(new CountMessageNumberFM(numElements)).setParallelism(1);
        TestUtils.tryExecute(env, "count elements from topics");
    }

    @Test
    public void testCommitOffsetsToPulsar() throws Exception {
        int numTopic = 3;
        List<Integer> messages = IntStream.range(0, 50).boxed().collect(Collectors.toList());
        List<String> topics = new ArrayList<>();
        Map<String, MessageId> expectedIds = new HashMap<>();

        for (int i = 0; i < numTopic; i++) {
            String topic = newTopic();
            topics.add(topic);
            List<MessageId> ids = sendTypedMessages(topic, SchemaType.INT32, messages, Optional.empty());
            expectedIds.put(topic, ids.get(ids.size() - 1));
        }

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(3);
        see.enableCheckpointing(200);

        String subName = UUID.randomUUID().toString();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(topics, ','));

        DataStream stream = see.addSource(
                new FlinkPulsarRowSourceSub(subName, serviceUrl, adminUrl, sourceProps).setStartFromEarliest());
        stream.addSink(new DiscardingSink());

        AtomicReference<Throwable> error = new AtomicReference<>();

        Thread runner = new Thread("runner") {
            @Override
            public void run() {
                try {
                    see.execute();
                } catch (Throwable e) {
                    if (!(e instanceof JobCancellationException)) {
                        error.set(e);
                    }
                }
            }
        };
        runner.start();

        Thread.sleep(3000);

        Long deadline = 30000000000L + System.nanoTime();

        boolean gotLast = false;

        do {
            Map<String, MessageId> ids = getCommittedOffsets(topics, subName);
            if (roughEquals(ids, expectedIds)) {
                gotLast = true;
            } else {
                Thread.sleep(100);
            }
        } while (System.nanoTime() < deadline && !gotLast);

        client.cancel(Iterables.getOnlyElement(getRunningJobs(client)));
        runner.join();

        Throwable t = error.get();
        if (t != null) {
            fail("Job failed with an exception " + ExceptionUtils.stringifyException(t));
        }

        assertTrue(gotLast);
    }

    @Test
    public void testAvro() throws Exception {
        String topic = newTopic();

        sendTypedMessages(topic, SchemaType.AVRO, fooList, Optional.empty(), SchemaData.Foo.class);

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.setParallelism(1);
        see.getConfig().disableSysoutLogging();
        see.setRestartStrategy(RestartStrategies.noRestart());

        String subName = UUID.randomUUID().toString();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_SINGLE_OPTION_KEY, topic);

        FlinkPulsarSource<SchemaData.Foo> source =
                new FlinkPulsarSource<>(serviceUrl, adminUrl, AvroDeser.of(SchemaData.Foo.class), sourceProps)
                .setStartFromEarliest();

        DataStream<Integer> ds = see.addSource(source)
                    .map(SchemaData.Foo::getI);

        ds.map(new FailingIdentityMapper<>(fooList.size()))
                .addSink(new SingletonStreamSink.StringSink<>()).setParallelism(1);

        try {
            see.execute("test read data of POJO using avro");
        } catch (Exception e) {

        }

        SingletonStreamSink.compareWithList(
                fooList.subList(0, fooList.size() - 1).stream().map(SchemaData.Foo::getI).map(Objects::toString).collect(Collectors.toList()));
    }

    @Test
    public void testJson() throws Exception {
        String topic = newTopic();

        sendTypedMessages(topic, SchemaType.JSON, fooList, Optional.empty(), SchemaData.Foo.class);

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.setParallelism(1);
        see.getConfig().disableSysoutLogging();
        see.setRestartStrategy(RestartStrategies.noRestart());

        String subName = UUID.randomUUID().toString();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_SINGLE_OPTION_KEY, topic);

        FlinkPulsarSource<SchemaData.Foo> source =
                new FlinkPulsarSource<>(serviceUrl, adminUrl, JsonDeser.of(SchemaData.Foo.class), sourceProps)
                        .setStartFromEarliest();

        DataStream<Integer> ds = see.addSource(source)
                .map(SchemaData.Foo::getI);

        ds.map(new FailingIdentityMapper<>(fooList.size()))
                .addSink(new SingletonStreamSink.StringSink<>()).setParallelism(1);

        try {
            see.execute("test read data of POJO using JSON");
        } catch (Exception e) {

        }

        SingletonStreamSink.compareWithList(
                fooList.subList(0, fooList.size() - 1).stream().map(SchemaData.Foo::getI).map(Objects::toString).collect(Collectors.toList()));
    }

    @Test
    public void testStartFromEarliest() throws Exception {
        int numTopic = 3;
        List<Integer> messages = IntStream.range(0, 50).boxed().collect(Collectors.toList());
        List<String> topics = new ArrayList<>();
        Map<String, Set<Integer>> expectedData = new HashMap<>();

        for (int i = 0; i < numTopic; i++) {
            String topic = newTopic();
            topics.add(topic);
            sendTypedMessages(topic, SchemaType.INT32, messages, Optional.empty());
            expectedData.put(topic, new HashSet<>(messages));
        }

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(3);

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(topics, ','));

        DataStream stream = see.addSource(
                new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromEarliest());

        stream.flatMap(new CheckAllMessageExist(expectedData, 150)).setParallelism(1);
        TestUtils.tryExecute(see, "start from earliest");
    }

    @Test
    public void testStartFromLatest() throws Exception {
        int numTopic = 3;
        List<Integer> messages = IntStream.range(0, 50).boxed().collect(Collectors.toList());
        List<Integer> newMessages = IntStream.range(50, 60).boxed().collect(Collectors.toList());
        Map<String, Set<Integer>> expectedData = new HashMap<>();

        List<String> topics = new ArrayList<>();

        for (int i = 0; i < numTopic; i++) {
            String topic = newTopic();
            topics.add(topic);
            sendTypedMessages(topic, SchemaType.INT32, messages, Optional.empty());
            expectedData.put(topic, new HashSet<>(newMessages));
        }

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(3);

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(topics, ','));

        DataStream stream = see.addSource(
                new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromLatest());
        stream.flatMap(new CheckAllMessageExist(expectedData, 30)).setParallelism(1);
        stream.addSink(new DiscardingSink());

        JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(see.getStreamGraph());
        JobID consumeJobId = jobGraph.getJobID();

        AtomicReference<Throwable> error = new AtomicReference<>();
        Thread consumerThread = new Thread() {
            @Override
            public void run() {
                try {
                    client.setDetached(false);
                    client.submitJob(jobGraph, FlinkPulsarITest.class.getClassLoader());
                } catch (Throwable e) {
                    if (!ExceptionUtils.findThrowable(e, JobCancellationException.class).isPresent()) {
                        error.set(e);
                    }
                }
            }
        };
        consumerThread.start();

        waitUntilJobIsRunning(client);

        Thread.sleep(3000);

        Thread extraProduceThread = new Thread() {
            @Override
            public void run() {
                try {
                    for (String tp : topics) {
                        sendTypedMessages(tp, SchemaType.INT32, newMessages, Optional.empty());
                    }
                } catch (PulsarClientException e) {
                    throw new RuntimeException(e);
                }
            }
        };
        extraProduceThread.start();

        consumerThread.join();

        Throwable consumerError = error.get();
        assertTrue(ExceptionUtils.findThrowable(consumerError, SuccessException.class).isPresent());
    }

    @Test
    public void testStartFromSpecific() throws Exception {
        String topic = newTopic();
        List<MessageId> mids = sendTypedMessages(topic, SchemaType.INT32, Arrays.asList(
                //  0,   1,   2, 3, 4, 5,  6,  7,  8
                -20, -21, -22, 1, 2, 3, 10, 11, 12), Optional.empty());

        Map<String, Set<Integer>> expectedData = new HashMap<>();
        expectedData.put(topic, new HashSet<>(Arrays.asList(2, 3, 10, 11, 12)));

        Map<String, MessageId> offset = new HashMap<>();
        offset.put(topic, mids.get(3));

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(1);

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_SINGLE_OPTION_KEY, topic);
        DataStream stream = see.addSource(
                new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromSpecificOffsets(offset));
        stream.flatMap(new CheckAllMessageExist(expectedData, 5)).setParallelism(1);

        TestUtils.tryExecute(see, "start from specific");
    }

    @Test
    public void testStartFromExternalSubscription() throws Exception {
        String topic = newTopic();
        List<MessageId> mids = sendTypedMessages(topic, SchemaType.INT32, Arrays.asList(
                //  0,   1,   2, 3, 4, 5,  6,  7,  8
                -20, -21, -22, 1, 2, 3, 10, 11, 12), Optional.empty());

        String subName = "sub-1";

        PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build();

        admin.topics().createSubscription(TopicName.get(topic).toString(), subName, mids.get(3));

        Map<String, Set<Integer>> expectedData = new HashMap<>();
        expectedData.put(topic, new HashSet<>(Arrays.asList(2, 3, 10, 11, 12)));

        StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        see.getConfig().disableSysoutLogging();
        see.setParallelism(1);

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_SINGLE_OPTION_KEY, topic);
        DataStream stream = see.addSource(
                new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromSubscription(subName));
        stream.flatMap(new CheckAllMessageExist(expectedData, 5)).setParallelism(1);

        TestUtils.tryExecute(see, "start from specific");

        assertTrue(Sets.newHashSet(admin.topics().getSubscriptions(topic)).contains(subName));

        admin.close();
    }

    @Test
    public void testOne2OneExactlyOnce() throws Exception {
        String topic = newTopic();
        int parallelism = 5;
        int numElementsPerPartition = 1000;
        int totalElements = parallelism * numElementsPerPartition;
        int failAfterElements = numElementsPerPartition / 3;

        List<String> allTopicNames = new ArrayList<>();
        for (int i = 0; i < parallelism; i++) {
            allTopicNames.add(topic + "-partition-" + i);
        }

        createTopic(topic, parallelism, adminUrl);

        generateRandomizedIntegerSequence(
                StreamExecutionEnvironment.getExecutionEnvironment(), topic, parallelism, numElementsPerPartition, true);

        // run the topology that fails and recovers

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(500);
        env.setParallelism(parallelism);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
        env.getConfig().disableSysoutLogging();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(allTopicNames, ','));

        env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromEarliest())
                .map(new PartitionValidationMapper(parallelism, 1))
                .map(new FailingIdentityMapper<Row>(failAfterElements))
                .addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

        FailingIdentityMapper.failedBefore = false;
        TestUtils.tryExecute(env, "one to one exactly once test");
    }

    @Test
    public void testOne2MultiSource() throws Exception {
        String topic = newTopic();
        int numPartitions = 5;
        int numElementsPerPartition = 1000;
        int totalElements = numPartitions * numElementsPerPartition;
        int failAfterElements = numElementsPerPartition / 3;

        List<String> allTopicNames = new ArrayList<>();
        for (int i = 0; i < numPartitions; i++) {
            allTopicNames.add(topic + "-partition-" + i);
        }

        createTopic(topic, numPartitions, adminUrl);

        generateRandomizedIntegerSequence(
                StreamExecutionEnvironment.getExecutionEnvironment(), topic, numPartitions, numElementsPerPartition, true);

        int parallelism = 2;

        // run the topology that fails and recovers

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(500);
        env.setParallelism(parallelism);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
        env.getConfig().disableSysoutLogging();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(allTopicNames, ','));

        env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromEarliest())
                .map(new PartitionValidationMapper(parallelism, 3))
                .map(new FailingIdentityMapper<Row>(failAfterElements))
                .addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

        FailingIdentityMapper.failedBefore = false;
        TestUtils.tryExecute(env, "One-source-multi-partitions exactly once test");
    }

    @Test
    public void testReadCompactedTopic() throws Exception {
	    String topic = newTopic();
	    String sub = "compacted-sub";
	    String outputTopic = topic + "-output";

	    int numKeys = 3;
	    int numMessagesPerKey = 10;

        PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build();
        admin.topics().createNonPartitionedTopic(topic);
        admin.topics().createSubscription(topic, sub, MessageId.earliest);

        PulsarClient client = PulsarClient.builder()
             .serviceUrl(serviceUrl)
             .build();

        try (Producer<String> producer = client.newProducer(Schema.STRING)
             .topic(topic)
             .create()) {
            for (int j = 0; j < numMessagesPerKey; j++) {
                for (int i = 0; i < numKeys; i++) {
                    producer.newMessage()
                        .key("key-" + i)
                        .value("value-" + i + "-" + j)
                        .send();
                }
            }
        }

	    admin.topics().triggerCompaction(topic);
	    LongRunningProcessStatus status = admin.topics().compactionStatus(topic);
	    while (status.status == Status.RUNNING) {
            TimeUnit.MILLISECONDS.sleep(500);
            status = admin.topics().compactionStatus(topic);
        }

	    assertEquals(Status.SUCCESS, status.status);

	    try (Reader<String> reader = client.newReader(Schema.STRING)
             .topic(topic)
             .startMessageId(MessageId.earliest)
             .readCompacted(true)
             .create()) {

	        for (int i = 0; i < numKeys; i++) {
	            Message<String> msg = reader.readNext();
	            log.info("Received message : {} - {}", msg.getKey(), msg.getValue());
            }
        }

	    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	    env.setParallelism(1);

	    Properties sourceProps = sourceProperties();
	    sourceProps.setProperty(
            PULSAR_READER_OPTION_KEY_PREFIX + "readCompacted", "true");
	    sourceProps.setProperty(
            TOPIC_SINGLE_OPTION_KEY, topic);

        Consumer<byte[]> outputConsumer = client.newConsumer()
            .topic(outputTopic)
            .subscriptionName(sub)
            .subscriptionType(SubscriptionType.Exclusive)
            .subscribe();

        env.addSource(new FlinkPulsarSource<>(serviceUrl, adminUrl, new SimpleStringSchema(), sourceProps)
                .setStartFromEarliest())
            .addSink(new FlinkPulsarSink<>(serviceUrl, adminUrl, Optional.of(outputTopic),
                new Properties(), null, String.class))
            .setParallelism(1);

        ExecutorService executorService = Executors.newSingleThreadExecutor();
        executorService.submit(() -> {
            try {
                TestUtils.tryExecute(env, "read from compacted topic");
            } catch (Exception e) {
                log.warn("Failed to execute a flink job", e);
            }
        });

        final Map<Integer, Integer> result = new HashMap<>();

        for (int i = 0; i < numKeys; i++) {
            Message<byte[]> msg = outputConsumer.receive();

            String value = new String(msg.getValue(), UTF_8);
            String[] parts = value.split("-");

            int keyIdx = Integer.parseInt(parts[1]);
            int valIdx = Integer.parseInt(parts[2]);

            assertEquals(numMessagesPerKey - 1, valIdx);
            assertNull(result.get(keyIdx));
            result.put(keyIdx, valIdx);
        }

        executorService.shutdown();
        outputConsumer.close();
        client.close();
        admin.close();
    }

    @Test
    public void testTaskNumberGreaterThanPartitionNumber() throws Exception {
        String topic = newTopic();
        int numPartitions = 5;
        int numElementsPerPartition = 1000;
        int totalElements = numPartitions * numElementsPerPartition;
        int failAfterElements = numElementsPerPartition / 3;

        List<String> allTopicNames = new ArrayList<>();
        for (int i = 0; i < numPartitions; i++) {
            allTopicNames.add(topic + "-partition-" + i);
        }

        createTopic(topic, numPartitions, adminUrl);

        generateRandomizedIntegerSequence(
                StreamExecutionEnvironment.getExecutionEnvironment(), topic, numPartitions, numElementsPerPartition, true);

        int parallelism = 8;

        // run the topology that fails and recovers

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(500);
        env.setParallelism(parallelism);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
        env.getConfig().disableSysoutLogging();

        Properties sourceProps = sourceProperties();
        sourceProps.setProperty(TOPIC_MULTI_OPTION_KEY, StringUtils.join(allTopicNames, ','));

        env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, sourceProps).setStartFromEarliest())
                .map(new PartitionValidationMapper(parallelism, 1))
                .map(new FailingIdentityMapper<Row>(failAfterElements))
                .addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

        FailingIdentityMapper.failedBefore = false;
        TestUtils.tryExecute(env, "source task number > partition number");
    }

    @Test
    public void testCancelingOnFullInput() throws Exception {
        String tp = newTopic();
        int parallelism = 3;
        createTopic(tp, parallelism, adminUrl);

        InfiniteStringGenerator generator = new InfiniteStringGenerator(tp);
        generator.start();

        // launch a consumer asynchronously

        AtomicReference<Throwable> jobError = new AtomicReference<>();

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        env.enableCheckpointing(100);
        env.getConfig().disableSysoutLogging();

        Properties prop = sourceProperties();
        prop.setProperty(TOPIC_SINGLE_OPTION_KEY, tp);
        env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, prop).setStartFromEarliest())
                .addSink(new DiscardingSink<>());

        JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
        JobID jobid = jobGraph.getJobID();

        Thread jobRunner = new Thread("program runner thread") {
            @Override
            public void run() {
                try {
                    client.setDetached(false);
                    client.submitJob(jobGraph, getClass().getClassLoader());
                } catch (Throwable e) {
                    jobError.set(e);
                }
            }
        };
        jobRunner.start();

        Thread.sleep(2000);
        Throwable failureCause = jobError.get();

        if (failureCause != null) {
            failureCause.printStackTrace();
            fail("Test failed prematurely with: " + failureCause.getMessage());
        }

        client.cancel(jobid);

        jobRunner.join();

        assertEquals(client.getJobStatus(jobid).get(), JobStatus.CANCELED);

        if (generator.isAlive()) {
            generator.shutdown();
            generator.join();
        } else {
            Throwable t = generator.getError();
            if (t != null) {
                t.printStackTrace();
                fail("Generator failed " + t.getMessage());
            } else {
                fail("Generator failed with no exception");
            }
        }
    }

    @Test
    public void testOnEmptyInput() throws Exception {
        String tp = newTopic();
        int parallelism = 3;
        createTopic(tp, parallelism, adminUrl);

        // launch a consumer asynchronously

        AtomicReference<Throwable> jobError = new AtomicReference<>();

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        env.enableCheckpointing(100);
        env.getConfig().disableSysoutLogging();

        Properties prop = sourceProperties();
        prop.setProperty(TOPIC_SINGLE_OPTION_KEY, tp);
        env.addSource(new FlinkPulsarRowSource(serviceUrl, adminUrl, prop).setStartFromEarliest())
                .addSink(new DiscardingSink<>());

        JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
        JobID jobid = jobGraph.getJobID();

        Thread jobRunner = new Thread("program runner thread") {
            @Override
            public void run() {
                try {
                    client.setDetached(false);
                    client.submitJob(jobGraph, getClass().getClassLoader());
                } catch (Throwable e) {
                    jobError.set(e);
                }
            }
        };
        jobRunner.start();

        Thread.sleep(2000);
        Throwable failureCause = jobError.get();

        if (failureCause != null) {
            failureCause.printStackTrace();
            fail("Test failed prematurely with: " + failureCause.getMessage());
        }

        client.cancel(jobid);

        jobRunner.join();

        assertEquals(client.getJobStatus(jobid).get(), JobStatus.CANCELED);

    }

    public static boolean isCause(
            Class<? extends Throwable> expected,
            Throwable exc) {
        return expected.isInstance(exc) || (
                exc != null && isCause(expected, exc.getCause())
        );
    }

    private static class MultiTopicSource extends RichParallelSourceFunction<Row> {
        private final List<String> topics;
        private final int numElements;
        private final int base;

        public MultiTopicSource(List<String> topics, int numElements, int base) {
            this.topics = topics;
            this.numElements = numElements;
            this.base = base;
        }

        public MultiTopicSource(List<String> topics, int numElements) {
            this(topics, numElements, 0);
        }

        @Override
        public void run(SourceContext<Row> ctx) throws Exception {
            for (String topic : topics) {
                for (int i = 0; i < numElements; i++) {
                    ctx.collect(intRowWithTopic(i + base, topic));
                }
            }
        }

        public Row intRowWithTopic(int i, String tp) {
            Row r = new Row(2);
            r.setField(0, tp);
            r.setField(1, i);
            return r;
        }

        @Override
        public void cancel() {

        }
    }

    private DataType intRowType() {
        return DataTypes.ROW(
                DataTypes.FIELD("v", DataTypes.INT()));
    }

    private TypeInformation<Row> intRowTypeInfo() {
        return (TypeInformation<Row>) LegacyTypeInfoDataTypeConverter
                .toLegacyTypeInfo(intRowType());
    }

    private DataType intRowWithTopicType() {
        return DataTypes.ROW(
                DataTypes.FIELD(TOPIC_ATTRIBUTE_NAME, DataTypes.STRING()),
                DataTypes.FIELD("v", DataTypes.INT())
        );
    }

    private TypeInformation<Row> intRowWithTopicTypeInfo() {
        return (TypeInformation<Row>) LegacyTypeInfoDataTypeConverter
                .toLegacyTypeInfo(intRowWithTopicType());
    }

    private Properties sinkProperties() {
        Properties props = new Properties();
        props.setProperty(FLUSH_ON_CHECKPOINT_OPTION_KEY, "true");
        props.setProperty(FAIL_ON_WRITE_OPTION_KEY, "true");
        return props;
    }

    private Properties sourceProperties() {
        Properties props = new Properties();
        props.setProperty(PARTITION_DISCOVERY_INTERVAL_MS_OPTION_KEY, "5000");
        return props;
    }

    private void produceIntoPulsar(DataStream<Row> stream, DataType dt, Properties props) {
        props.setProperty(FLUSH_ON_CHECKPOINT_OPTION_KEY, "true");
        stream.addSink(new FlinkPulsarRowSink(serviceUrl, adminUrl, Optional.empty(), props, dt));
    }

    private static class AssertSink extends FlinkPulsarRowSink {

        private final int cacheSize;

        public AssertSink(String serviceUrl, String adminUrl, int cacheSize, Properties properties, DataType dataType) {
            super(serviceUrl, adminUrl, Optional.empty(), properties, dataType);
            this.cacheSize = cacheSize;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            assertEquals(CachedPulsarClient.getCacheSize(), cacheSize);
        }
    }

    private static class CountMessageNumberFM implements FlatMapFunction<Row, Integer> {
        private final int numElements;

        private Map<String, Integer> map;

        private CountMessageNumberFM(int numElements) {
            this.numElements = numElements;
            this.map = new HashMap<>();
        }

        @Override
        public void flatMap(Row value, Collector<Integer> out) throws Exception {
            String topic = (String) value.getField(2);
            Integer old = map.getOrDefault(topic, 0);
            map.put(topic, old + 1);

            for (Map.Entry<String, Integer> entry : map.entrySet()) {
                if (entry.getValue() < numElements) {
                    return;
                } else if (entry.getValue() > numElements) {
                    throw new RuntimeException(entry.getKey() + " has " + entry.getValue() + "elements");
                }
            }
            throw new SuccessException();
        }
    }

    private static Map<String, MessageId> getCommittedOffsets(List<String> topics, String prefix) {
        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()) {
            Map<String, MessageId> results = new HashMap<>();
            String subName = "flink-pulsar-" + prefix;

            for (String topic : topics) {
                int index = TopicName.get(topic).getPartitionIndex();
                MessageId mid = null;

                try {
                    PersistentTopicInternalStats.CursorStats cursor = admin.topics().getInternalStats(topic).cursors.get(subName);
                    if (cursor != null) {
                        String[] le = cursor.readPosition.split(":");
                        long ledgerId = Long.parseLong(le[0]);
                        long entryId = Long.parseLong(le[1]);
                        mid = new MessageIdImpl(ledgerId, entryId, index);
                    }
                } catch (PulsarAdminException e) {
                    // do nothing
                }
                results.put(topic, mid);
            }
            return results;
        } catch (PulsarClientException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static boolean roughEquals(Map<String, MessageId> a, Map<String, MessageId> b) {
        for (Map.Entry<String, MessageId> aE : a.entrySet()) {
            MessageId bmid = b.getOrDefault(aE.getKey(), MessageId.latest);
            if (!ReaderThread.messageIdRoughEquals(bmid, aE.getValue())) {
                return false;
            }
        }
        return true;
    }

    public static void createTopic(String topic, int partition, String adminUrl) {
        assert partition > 1;
        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()) {
            admin.topics().createPartitionedTopic(topic, partition);
        } catch (PulsarClientException | PulsarAdminException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Check if all message exist.
     */
    public static class CheckAllMessageExist extends RichFlatMapFunction<Row, Row> {
        private final Map<String, Set<Integer>> expected;
        private final int total;

        private Map<String, List<Integer>> map = new HashMap<>();
        private int count = 0;

        public CheckAllMessageExist(Map<String, Set<Integer>> expected, int total) {
            this.expected = expected;
            this.total = total;
        }

        @Override
        public void flatMap(Row value, Collector<Row> out) throws Exception {
            String topic = (String) value.getField(2);
            int v = Integer.parseInt(value.getField(0).toString());
            List<Integer> current = map.getOrDefault(topic, new ArrayList<>());
            current.add(v);
            map.put(topic, current);

            count++;

            if (count == total) {
                for (Map.Entry<String, List<Integer>> e : map.entrySet()) {
                    Set<Integer> s = new HashSet<>(e.getValue());
                    if (s.size() != e.getValue().size()) {
                        throw new RuntimeException("duplicate elements in " + topic + " " + e.getValue().toString());
                    }
                    Set<Integer> expectedSet = expected.getOrDefault(e.getKey(), null);
                    if (expectedSet == null) {
                        throw new RuntimeException("Unknown topic seen " + e.getKey());
                    } else {
                        if (!expectedSet.equals(s)) {
                            throw new RuntimeException("" + expectedSet + "\n" + s);
                        }
                    }
                }
                throw new SuccessException();
            }
        }
    }

    private void generateRandomizedIntegerSequence(
            StreamExecutionEnvironment see,
            String tp,
            int numPartitions,
            int numElements,
            boolean randomizedOrder) throws Exception {
        see.setParallelism(numPartitions);
        see.getConfig().disableSysoutLogging();
        see.setRestartStrategy(RestartStrategies.noRestart());

        DataStream stream = see.addSource(new RandomizedIntegerSeq(tp, numPartitions, numElements, randomizedOrder));
        produceIntoPulsar(stream, intRowWithTopicType(), sinkProperties());
        see.execute("scrambles in sequence generator");
    }

    private static class RandomizedIntegerSeq extends RichParallelSourceFunction<Row> {
        private final String tp;
        private final int numPartitions;
        private final int numElements;
        private final boolean randomizedOrder;

        private volatile boolean running = true;

        private RandomizedIntegerSeq(String tp, int numPartitions, int numElements, boolean randomizedOrder) {
            this.tp = tp;
            this.numPartitions = numPartitions;
            this.numElements = numElements;
            this.randomizedOrder = randomizedOrder;
        }

        @Override
        public void run(SourceContext<Row> ctx) throws Exception {
            int subIndex = getRuntimeContext().getIndexOfThisSubtask();
            int all = getRuntimeContext().getNumberOfParallelSubtasks();

            List<Row> rows = new ArrayList<>();
            for (int i = 0; i < numElements; i++) {
                rows.add(intRowWithTopic(subIndex + i * all, topicName(tp, subIndex)));
            }

            if (randomizedOrder) {
                Random rand = new Random();
                for (int i = 0; i < numElements; i++) {
                    int otherPos = rand.nextInt(numElements);
                    Row tmp = rows.get(i);
                    rows.set(i, rows.get(otherPos));
                    rows.set(otherPos, tmp);
                }
            }

            rows.forEach(r -> ctx.collect(r));
        }

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

        String topicName(String tp, int index) {
            return tp + "-partition-" + index;
        }

        Row intRowWithTopic(int i, String tp) {
            Row r = new Row(2);
            r.setField(0, tp);
            r.setField(1, i);
            return r;
        }
    }

    private static class PartitionValidationMapper implements MapFunction<Row, Row> {
        private final int numPartitions;
        private final int maxPartitions;

        private HashSet<String> myTopics;

        private PartitionValidationMapper(int numPartitions, int maxPartitions) {
            this.numPartitions = numPartitions;
            this.maxPartitions = maxPartitions;
            myTopics = new HashSet<>();
        }

        @Override
        public Row map(Row value) throws Exception {
            String topic = (String) value.getField(2);
            myTopics.add(topic);
            if (myTopics.size() > maxPartitions) {
                throw new Exception(String.format(
                        "Error: Elements from too many different partitions: %s, Expected elements only from %d partitions",
                        myTopics.toString(), maxPartitions));
            }
            return value;
        }
    }

    private static class InfiniteStringGenerator extends Thread {
        private final String tp;

        private volatile boolean running = true;
        private volatile Throwable error = null;

        public InfiniteStringGenerator(String tp) {
            this.tp = tp;
        }

        @Override
        public void run() {
            try {
                Properties props = new Properties();
                props.setProperty(FLUSH_ON_CHECKPOINT_OPTION_KEY, "true");
                props.setProperty(FAIL_ON_WRITE_OPTION_KEY, "true");

                StreamSink<String> sink = new StreamSink<>(
                        new FlinkPulsarSinkBase<String>(serviceUrl, adminUrl, Optional.of(tp), props, new TopicKeyExtractor<String>() {
                            @Override
                            public byte[] serializeKey(String element) {
                                return new byte[0];
                            }

                            @Override
                            public String getTopic(String element) {
                                return null;
                            }
                        }) {
                            @Override
                            protected Schema<?> getPulsarSchema() {
                                return Schema.STRING;
                            }
                        });

                OneInputStreamOperatorTestHarness<String, Object> testHarness = new OneInputStreamOperatorTestHarness(sink);
                testHarness.open();

                StringBuilder bld = new StringBuilder();
                Random rnd = new Random();

                while (running) {
                    bld.setLength(0);
                    int len = rnd.nextInt(100) + 1;
                    for (int i = 0; i < len; i++) {
                        bld.append((char) (rnd.nextInt(20) + 'a'));
                    }

                    String next = bld.toString();
                    testHarness.processElement(new StreamRecord<>(next));
                }
            } catch (Exception e) {
                this.error = e;
            }
        }

        public void shutdown() {
            this.running = false;
            this.interrupt();
        }

        public Throwable getError() {
            return error;
        }
    }

    private static class IntegerDeserializer implements DeserializationSchema<Integer> {
        private final TypeInformation<Integer> ti;
        private final TypeSerializer<Integer> ser;

        public IntegerDeserializer() {
            this.ti = Types.INT();
            this.ser = ti.createSerializer(new ExecutionConfig());
        }

        @Override
        public Integer deserialize(byte[] message) throws IOException {

            DataInputView in = new DataInputViewStreamWrapper(new ByteArrayInputStream(message));
            Integer i = ser.deserialize(in);

            return i;
        }

        @Override
        public boolean isEndOfStream(Integer nextElement) {
            return false;
        }

        @Override
        public TypeInformation<Integer> getProducedType() {
            return ti;
        }
    }

    private static class FlinkPulsarRowSourceSub extends FlinkPulsarRowSource {

        private final String sub;

        public FlinkPulsarRowSourceSub(String sub, String serviceUrl, String adminUrl, Properties properties) {
            super(serviceUrl, adminUrl, properties);
            this.sub = sub;
        }

        @Override
        protected String getSubscriptionName() {
            return "flink-pulsar-" + sub;
        }
    }
}