/*
 * 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.beam.sdk.io.kafka;

import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import org.apache.beam.sdk.coders.BigEndianLongCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.io.kafka.KafkaIO.WriteRecords;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.metrics.SinkMetrics;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.windowing.AfterPane;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.Repeatedly;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TimestampedValue.TimestampedValueCoder;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.CacheBuilder;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.CacheLoader;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.LoadingCache;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.RemovalCause;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterators;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Lists;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.joda.time.DateTimeUtils;
import org.joda.time.DateTimeZone;
import org.joda.time.format.DateTimeFormat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Exactly-once sink transform for Kafka. See {@link KafkaIO} for user visible documentation and
 * example usage.
 */
class KafkaExactlyOnceSink<K, V>
    extends PTransform<PCollection<ProducerRecord<K, V>>, PCollection<Void>> {

  // Dataflow ensures at-least once processing for side effects like sinks. In order to provide
  // exactly-once semantics, a sink needs to be idempotent or it should avoid writing records
  // that have already been written. This snk does the latter. All the the records are ordered
  // across a fixed number of shards and records in each shard are written in order. It drops
  // any records that are already written and buffers those arriving out of order.
  //
  // Exactly once sink involves two shuffles of the records:
  //    A : Assign a shard ---> B : Assign sequential ID ---> C : Write to Kafka in order
  //
  // Processing guarantees also require deterministic processing within user transforms.
  // Here, that requires order of the records committed to Kafka by C should not be affected by
  // restarts in C and its upstream stages.
  //
  // A : Assigns a random shard for message. Note that there are no ordering guarantees for
  //     writing user records to Kafka. User can still control partitioning among topic
  //     partitions as with regular sink (of course, there are no ordering guarantees in
  //     regular Kafka sink either).
  // B : Assigns an id sequentially for each messages within a shard.
  // C : Writes each shard to Kafka in sequential id order. In Dataflow, when C sees a record
  //     and id, it implies that record and the associated id are checkpointed to persistent
  //     storage and this record will always have same id, even in retries.
  //     Exactly-once semantics are achieved by writing records in the strict order of
  //     these check-pointed sequence ids.
  //
  // Parallelism for B and C is fixed to 'numShards', which defaults to number of partitions
  // for the topic. A few reasons for that:
  //  - B & C implement their functionality using per-key state. Shard id makes it independent
  //    of cardinality of user key.
  //  - We create one producer per shard, and its 'transactional id' is based on shard id. This
  //    requires that number of shards to be finite. This also helps with batching. and avoids
  //    initializing producers and transactions.
  //  - Most importantly, each of sharded writers stores 'next message id' in partition
  //    metadata, which is committed atomically with Kafka transactions. This is critical
  //    to handle retries of C correctly. Initial testing showed number of shards could be
  //    larger than number of partitions for the topic.
  //
  // Number of shards can change across multiple runs of a pipeline (job upgrade in Dataflow).
  //

  private static final Logger LOG = LoggerFactory.getLogger(KafkaExactlyOnceSink.class);
  private static final String METRIC_NAMESPACE = "KafkaExactlyOnceSink";

  private final WriteRecords<K, V> spec;

  static void ensureEOSSupport() {
    checkArgument(
        ProducerSpEL.supportsTransactions(),
        "%s %s",
        "This version of Kafka client does not support transactions required to support",
        "exactly-once semantics. Please use Kafka client version 0.11 or newer.");
  }

  KafkaExactlyOnceSink(WriteRecords<K, V> spec) {
    this.spec = spec;
  }

  @Override
  public PCollection<Void> expand(PCollection<ProducerRecord<K, V>> input) {

    int numShards = spec.getNumShards();
    if (numShards <= 0) {
      try (Consumer<?, ?> consumer = openConsumer(spec)) {
        numShards = consumer.partitionsFor(spec.getTopic()).size();
        LOG.info(
            "Using {} shards for exactly-once writer, matching number of partitions "
                + "for topic '{}'",
            numShards,
            spec.getTopic());
      }
    }
    checkState(numShards > 0, "Could not set number of shards");

    return input
        .apply(
            Window.<ProducerRecord<K, V>>into(new GlobalWindows()) // Everything into global window.
                .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
                .discardingFiredPanes())
        .apply(
            String.format("Shuffle across %d shards", numShards),
            ParDo.of(new Reshard<>(numShards)))
        .apply("Persist sharding", GroupByKey.create())
        .apply("Assign sequential ids", ParDo.of(new Sequencer<>()))
        .apply("Persist ids", GroupByKey.create())
        .apply(
            String.format("Write to Kafka topic '%s'", spec.getTopic()),
            ParDo.of(new ExactlyOnceWriter<>(spec, input.getCoder())));
  }

  /** Shuffle messages assigning each randomly to a shard. */
  private static class Reshard<K, V>
      extends DoFn<ProducerRecord<K, V>, KV<Integer, TimestampedValue<ProducerRecord<K, V>>>> {

    private final int numShards;
    private transient int shardId;

    Reshard(int numShards) {
      this.numShards = numShards;
    }

    @Setup
    public void setup() {
      shardId = ThreadLocalRandom.current().nextInt(numShards);
    }

    @ProcessElement
    public void processElement(ProcessContext ctx) {
      shardId = (shardId + 1) % numShards; // round-robin among shards.
      ctx.output(KV.of(shardId, TimestampedValue.of(ctx.element(), ctx.timestamp())));
    }
  }

  private static class Sequencer<K, V>
      extends DoFn<
          KV<Integer, Iterable<TimestampedValue<ProducerRecord<K, V>>>>,
          KV<Integer, KV<Long, TimestampedValue<ProducerRecord<K, V>>>>> {

    private static final String NEXT_ID = "nextId";

    @StateId(NEXT_ID)
    private final StateSpec<ValueState<Long>> nextIdSpec = StateSpecs.value();

    @ProcessElement
    public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState, ProcessContext ctx) {
      long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
      int shard = ctx.element().getKey();
      for (TimestampedValue<ProducerRecord<K, V>> value : ctx.element().getValue()) {
        ctx.output(KV.of(shard, KV.of(nextId, value)));
        nextId++;
      }
      nextIdState.write(nextId);
    }
  }

  private static class ExactlyOnceWriter<K, V>
      extends DoFn<KV<Integer, Iterable<KV<Long, TimestampedValue<ProducerRecord<K, V>>>>>, Void> {

    private static final String NEXT_ID = "nextId";
    private static final String MIN_BUFFERED_ID = "minBufferedId";
    private static final String OUT_OF_ORDER_BUFFER = "outOfOrderBuffer";
    private static final String WRITER_ID = "writerId";

    // Not sure of a good limit. This applies only for large bundles.
    private static final int MAX_RECORDS_PER_TXN = 1000;
    private static final ObjectMapper JSON_MAPPER = new ObjectMapper();

    @StateId(NEXT_ID)
    private final StateSpec<ValueState<Long>> sequenceIdSpec = StateSpecs.value();

    @StateId(MIN_BUFFERED_ID)
    private final StateSpec<ValueState<Long>> minBufferedIdSpec = StateSpecs.value();

    @StateId(OUT_OF_ORDER_BUFFER)
    private final StateSpec<BagState<KV<Long, TimestampedValue<ProducerRecord<K, V>>>>>
        outOfOrderBufferSpec;
    // A random id assigned to each shard. Helps with detecting when multiple jobs are mistakenly
    // started with same groupId used for storing state on Kafka side, including the case where
    // a job is restarted with same groupId, but the metadata from previous run was not cleared.
    // Better to be safe and error out with a clear message.
    @StateId(WRITER_ID)
    private final StateSpec<ValueState<String>> writerIdSpec = StateSpecs.value();

    private final WriteRecords<K, V> spec;

    // Metrics
    private final Counter elementsWritten = SinkMetrics.elementsWritten();
    // Elements buffered due to out of order arrivals.
    private final Counter elementsBuffered = Metrics.counter(METRIC_NAMESPACE, "elementsBuffered");
    private final Counter numTransactions = Metrics.counter(METRIC_NAMESPACE, "numTransactions");

    ExactlyOnceWriter(WriteRecords<K, V> spec, Coder<ProducerRecord<K, V>> elemCoder) {
      this.spec = spec;
      this.outOfOrderBufferSpec =
          StateSpecs.bag(KvCoder.of(BigEndianLongCoder.of(), TimestampedValueCoder.of(elemCoder)));
    }

    @Setup
    public void setup() {
      // This is on the worker. Ensure the runtime version is till compatible.
      KafkaExactlyOnceSink.ensureEOSSupport();
    }

    // Futures ignored as exceptions will be flushed out in the commitTxn
    @SuppressWarnings("FutureReturnValueIgnored")
    @RequiresStableInput
    @ProcessElement
    public void processElement(
        @StateId(NEXT_ID) ValueState<Long> nextIdState,
        @StateId(MIN_BUFFERED_ID) ValueState<Long> minBufferedIdState,
        @StateId(OUT_OF_ORDER_BUFFER)
            BagState<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> oooBufferState,
        @StateId(WRITER_ID) ValueState<String> writerIdState,
        ProcessContext ctx)
        throws IOException {

      int shard = ctx.element().getKey();

      minBufferedIdState.readLater();
      long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
      long minBufferedId = MoreObjects.firstNonNull(minBufferedIdState.read(), Long.MAX_VALUE);

      ShardWriterCache<K, V> cache =
          (ShardWriterCache<K, V>) CACHE_BY_GROUP_ID.getUnchecked(spec.getSinkGroupId());
      ShardWriter<K, V> writer = cache.removeIfPresent(shard);
      if (writer == null) {
        writer = initShardWriter(shard, writerIdState, nextId);
      }

      long committedId = writer.committedId;

      if (committedId >= nextId) {
        // This is a retry of an already committed batch.
        LOG.info(
            "{}: committed id {} is ahead of expected {}. {} records will be dropped "
                + "(these are already written).",
            shard,
            committedId,
            nextId - 1,
            committedId - nextId + 1);
        nextId = committedId + 1;
      }

      try {
        writer.beginTxn();
        int txnSize = 0;

        // Iterate in recordId order. The input iterator could be mostly sorted.
        // There might be out of order messages buffered in earlier iterations. These
        // will get merged if and when minBufferedId matches nextId.

        Iterator<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> iter =
            ctx.element().getValue().iterator();

        while (iter.hasNext()) {
          KV<Long, TimestampedValue<ProducerRecord<K, V>>> kv = iter.next();
          long recordId = kv.getKey();

          if (recordId < nextId) {
            LOG.info(
                "{}: dropping older record {}. Already committed till {}",
                shard,
                recordId,
                committedId);
            continue;
          }

          if (recordId > nextId) {
            // Out of order delivery. Should be pretty rare (what about in a batch pipeline?)
            LOG.info(
                "{}: Saving out of order record {}, next record id to be written is {}",
                shard,
                recordId,
                nextId);

            // checkState(recordId - nextId < 10000, "records are way out of order");

            oooBufferState.add(kv);
            minBufferedId = Math.min(minBufferedId, recordId);
            minBufferedIdState.write(minBufferedId);
            elementsBuffered.inc();
            continue;
          }

          // recordId and nextId match. Finally write the record.

          writer.sendRecord(kv.getValue(), elementsWritten);
          nextId++;

          if (++txnSize >= MAX_RECORDS_PER_TXN) {
            writer.commitTxn(recordId, numTransactions);
            txnSize = 0;
            writer.beginTxn();
          }

          if (minBufferedId == nextId) {
            // One or more of the buffered records can be committed now.
            // Read all of them in to memory and sort them. Reading into memory
            // might be problematic in extreme cases. Might need to improve it in future.

            List<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> buffered =
                Lists.newArrayList(oooBufferState.read());
            buffered.sort(new KV.OrderByKey<>());

            LOG.info(
                "{} : merging {} buffered records (min buffered id is {}).",
                shard,
                buffered.size(),
                minBufferedId);

            oooBufferState.clear();
            minBufferedIdState.clear();
            minBufferedId = Long.MAX_VALUE;

            iter =
                Iterators.mergeSorted(
                    ImmutableList.of(iter, buffered.iterator()), new KV.OrderByKey<>());
          }
        }

        writer.commitTxn(nextId - 1, numTransactions);
        nextIdState.write(nextId);

      } catch (ProducerSpEL.UnrecoverableProducerException e) {
        // Producer JavaDoc says these are not recoverable errors and producer should be closed.

        // Close the producer and a new producer will be initialized in retry.
        // It is possible that a rough worker keeps retrying and ends up fencing off
        // active producers. How likely this might be or how well such a scenario is handled
        // depends on the runner. For now we will leave it to upper layers, will need to revisit.

        LOG.warn(
            "{} : closing producer {} after unrecoverable error. The work might have migrated."
                + " Committed id {}, current id {}.",
            writer.shard,
            writer.producerName,
            writer.committedId,
            nextId - 1,
            e);

        writer.producer.close();
        writer = null; // No need to cache it.
        throw e;
      } finally {
        if (writer != null) {
          cache.insert(shard, writer);
        }
      }
    }

    private static class ShardMetadata {

      @JsonProperty("seq")
      public final long sequenceId;

      @JsonProperty("id")
      public final String writerId;

      private ShardMetadata() { // for json deserializer
        sequenceId = -1;
        writerId = null;
      }

      ShardMetadata(long sequenceId, String writerId) {
        this.sequenceId = sequenceId;
        this.writerId = writerId;
      }
    }

    /** A wrapper around Kafka producer. One for each of the shards. */
    private static class ShardWriter<K, V> {

      private final int shard;
      private final String writerId;
      private final Producer<K, V> producer;
      private final String producerName;
      private final WriteRecords<K, V> spec;
      private long committedId;

      ShardWriter(
          int shard,
          String writerId,
          Producer<K, V> producer,
          String producerName,
          WriteRecords<K, V> spec,
          long committedId) {
        this.shard = shard;
        this.writerId = writerId;
        this.producer = producer;
        this.producerName = producerName;
        this.spec = spec;
        this.committedId = committedId;
      }

      void beginTxn() {
        ProducerSpEL.beginTransaction(producer);
      }

      Future<RecordMetadata> sendRecord(
          TimestampedValue<ProducerRecord<K, V>> record, Counter sendCounter) {
        try {
          Long timestampMillis =
              spec.getPublishTimestampFunction() != null
                  ? spec.getPublishTimestampFunction()
                      .getTimestamp(record.getValue(), record.getTimestamp())
                      .getMillis()
                  : null;

          Future<RecordMetadata> result =
              producer.send(
                  new ProducerRecord<>(
                      spec.getTopic(),
                      null,
                      timestampMillis,
                      record.getValue().key(),
                      record.getValue().value()));
          sendCounter.inc();
          return result;
        } catch (KafkaException e) {
          ProducerSpEL.abortTransaction(producer);
          throw e;
        }
      }

      void commitTxn(long lastRecordId, Counter numTransactions) throws IOException {
        try {
          // Store id in consumer group metadata for the partition.
          // NOTE: Kafka keeps this metadata for 24 hours since the last update. This limits
          // how long the pipeline could be down before resuming it. It does not look like
          // this TTL can be adjusted (asked about it on Kafka users list).
          ProducerSpEL.sendOffsetsToTransaction(
              producer,
              ImmutableMap.of(
                  new TopicPartition(spec.getTopic(), shard),
                  new OffsetAndMetadata(
                      0L,
                      JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId, writerId)))),
              spec.getSinkGroupId());
          ProducerSpEL.commitTransaction(producer);

          numTransactions.inc();
          LOG.debug("{} : committed {} records", shard, lastRecordId - committedId);

          committedId = lastRecordId;
        } catch (KafkaException e) {
          ProducerSpEL.abortTransaction(producer);
          throw e;
        }
      }
    }

    private ShardWriter<K, V> initShardWriter(
        int shard, ValueState<String> writerIdState, long nextId) throws IOException {

      String producerName = String.format("producer_%d_for_%s", shard, spec.getSinkGroupId());
      Producer<K, V> producer = initializeExactlyOnceProducer(spec, producerName);

      // Fetch latest committed metadata for the partition (if any). Checks committed sequence ids.
      try {

        String writerId = writerIdState.read();

        OffsetAndMetadata committed;

        try (Consumer<?, ?> consumer = openConsumer(spec)) {
          committed = consumer.committed(new TopicPartition(spec.getTopic(), shard));
        }

        long committedSeqId = -1;

        if (committed == null || committed.metadata() == null || committed.metadata().isEmpty()) {
          checkState(
              nextId == 0 && writerId == null,
              "State exists for shard %s (nextId %s, writerId '%s'), but there is no state "
                  + "stored with Kafka topic '%s' group id '%s'",
              shard,
              nextId,
              writerId,
              spec.getTopic(),
              spec.getSinkGroupId());

          writerId =
              String.format(
                  "%X - %s",
                  new Random().nextInt(Integer.MAX_VALUE),
                  DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss")
                      .withZone(DateTimeZone.UTC)
                      .print(DateTimeUtils.currentTimeMillis()));
          writerIdState.write(writerId);
          LOG.info("Assigned writer id '{}' to shard {}", writerId, shard);

        } else {
          ShardMetadata metadata = JSON_MAPPER.readValue(committed.metadata(), ShardMetadata.class);

          checkNotNull(metadata.writerId);

          if (writerId == null) {
            // a) This might be a restart of the job from scratch, in which case metatdata
            // should be ignored and overwritten with new one.
            // b) This job might be started with an incorrect group id which is an error.
            // c) There is an extremely small chance that this is a retry of the first bundle
            // where metatdate was committed to Kafka but the bundle results were not committed
            // in Beam, in which case it should be treated as correct metadata.
            // How can we tell these three cases apart? Be safe and throw an exception.
            //
            // We could let users explicitly an option to override the existing metadata.
            //
            throw new IllegalStateException(
                String.format(
                    "Kafka metadata exists for shard %s, but there is no stored state for it. "
                        + "This mostly indicates groupId '%s' is used else where or in earlier runs. "
                        + "Try another group id. Metadata for this shard on Kafka : '%s'",
                    shard, spec.getSinkGroupId(), committed.metadata()));
          }

          checkState(
              writerId.equals(metadata.writerId),
              "Writer ids don't match. This is mostly a unintended misuse of groupId('%s')."
                  + "Beam '%s', Kafka '%s'",
              spec.getSinkGroupId(),
              writerId,
              metadata.writerId);

          committedSeqId = metadata.sequenceId;

          checkState(
              committedSeqId >= (nextId - 1),
              "Committed sequence id can not be lower than %s, partition metadata : %s",
              nextId - 1,
              committed.metadata());
        }

        LOG.info(
            "{} : initialized producer {} with committed sequence id {}",
            shard,
            producerName,
            committedSeqId);

        return new ShardWriter<>(shard, writerId, producer, producerName, spec, committedSeqId);

      } catch (IOException e) {
        producer.close();
        throw e;
      }
    }

    /**
     * A wrapper around guava cache to provide insert()/remove() semantics. A ShardWriter will be
     * closed if it is stays in cache for more than 1 minute, i.e. not used inside
     * KafkaExactlyOnceSink DoFn for a minute.
     */
    private static class ShardWriterCache<K, V> {

      static final ScheduledExecutorService SCHEDULED_CLEAN_UP_THREAD =
          Executors.newSingleThreadScheduledExecutor();

      static final int CLEAN_UP_CHECK_INTERVAL_MS = 10 * 1000;
      static final int IDLE_TIMEOUT_MS = 60 * 1000;

      private final Cache<Integer, ShardWriter<K, V>> cache;

      // Exceptions arising from the cache cleanup are ignored
      @SuppressWarnings("FutureReturnValueIgnored")
      ShardWriterCache() {
        this.cache =
            CacheBuilder.newBuilder()
                .expireAfterWrite(IDLE_TIMEOUT_MS, TimeUnit.MILLISECONDS)
                .<Integer, ShardWriter<K, V>>removalListener(
                    notification -> {
                      if (notification.getCause() != RemovalCause.EXPLICIT) {
                        ShardWriter writer = notification.getValue();
                        LOG.info(
                            "{} : Closing idle shard writer {} after 1 minute of idle time.",
                            writer.shard,
                            writer.producerName);
                        writer.producer.close();
                      }
                    })
                .build();

        // run cache.cleanUp() every 10 seconds.
        SCHEDULED_CLEAN_UP_THREAD.scheduleAtFixedRate(
            cache::cleanUp,
            CLEAN_UP_CHECK_INTERVAL_MS,
            CLEAN_UP_CHECK_INTERVAL_MS,
            TimeUnit.MILLISECONDS);
      }

      ShardWriter<K, V> removeIfPresent(int shard) {
        return cache.asMap().remove(shard);
      }

      void insert(int shard, ShardWriter<K, V> writer) {
        ShardWriter<K, V> existing = cache.asMap().putIfAbsent(shard, writer);
        checkState(
            existing == null, "Unexpected multiple instances of writers for shard %s", shard);
      }
    }

    // One cache for each sink (usually there is only one sink per pipeline)
    private static final LoadingCache<String, ShardWriterCache<?, ?>> CACHE_BY_GROUP_ID =
        CacheBuilder.newBuilder()
            .build(
                new CacheLoader<String, ShardWriterCache<?, ?>>() {
                  @Override
                  public ShardWriterCache<?, ?> load(String key) throws Exception {
                    return new ShardWriterCache<>();
                  }
                });
  }

  /**
   * Opens a generic consumer that is mainly meant for metadata operations like fetching number of
   * partitions for a topic rather than for fetching messages.
   */
  private static Consumer<?, ?> openConsumer(WriteRecords<?, ?> spec) {
    return spec.getConsumerFactoryFn()
        .apply(
            ImmutableMap.of(
                ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,
                spec.getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
                ConsumerConfig.GROUP_ID_CONFIG,
                spec.getSinkGroupId(),
                ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                ByteArrayDeserializer.class,
                ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                ByteArrayDeserializer.class));
  }

  private static <K, V> Producer<K, V> initializeExactlyOnceProducer(
      WriteRecords<K, V> spec, String producerName) {

    Map<String, Object> producerConfig = new HashMap<>(spec.getProducerConfig());
    producerConfig.putAll(
        ImmutableMap.of(
            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
            spec.getKeySerializer(),
            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,
            spec.getValueSerializer(),
            ProducerSpEL.ENABLE_IDEMPOTENCE_CONFIG,
            true,
            ProducerSpEL.TRANSACTIONAL_ID_CONFIG,
            producerName));

    Producer<K, V> producer =
        spec.getProducerFactoryFn() != null
            ? spec.getProducerFactoryFn().apply(producerConfig)
            : new KafkaProducer<>(producerConfig);

    ProducerSpEL.initTransactions(producer);
    return producer;
  }
}