/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * license agreements; and to You under the Apache License, version 2.0:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * This file is part of the Apache Pekko project, which was derived from Akka.
 */

/*
 * Copyright (C) 2014 - 2016 Softwaremill <https://softwaremill.com>
 * Copyright (C) 2016 - 2020 Lightbend Inc. <https://www.lightbend.com>
 */

package docs.javadsl;

import org.apache.pekko.Done;
import org.apache.pekko.NotUsed;
import org.apache.pekko.actor.AbstractLoggingActor;
import org.apache.pekko.actor.ActorRef;
import org.apache.pekko.actor.ActorSystem;
import org.apache.pekko.actor.Props;
import org.apache.pekko.actor.typed.Behavior;
import org.apache.pekko.actor.typed.javadsl.ActorContext;
import org.apache.pekko.actor.typed.javadsl.Behaviors;
// #withTypedRebalanceListenerActor
// #consumerActorTyped
// adds support for actors to a classic actor system and context
import org.apache.pekko.actor.typed.javadsl.Adapter;
// #consumerActorTyped
// #withTypedRebalanceListenerActor
import org.apache.pekko.japi.Pair;
import org.apache.pekko.kafka.*;
import org.apache.pekko.kafka.javadsl.Committer;
import org.apache.pekko.kafka.javadsl.Consumer;
import org.apache.pekko.kafka.javadsl.Producer;
import org.apache.pekko.kafka.javadsl.PartitionAssignmentHandler;
import org.apache.pekko.kafka.testkit.TestcontainersKafkaTest;
import org.apache.pekko.kafka.tests.javadsl.LogCapturingExtension;
import org.apache.pekko.stream.RestartSettings;
import org.apache.pekko.stream.javadsl.*;
import org.apache.pekko.testkit.javadsl.TestKit;
import com.typesafe.config.Config;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;
import org.junit.jupiter.api.extension.ExtendWith;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;

import static org.hamcrest.CoreMatchers.*;
import static org.hamcrest.MatcherAssert.*;
import static org.junit.Assert.assertEquals;

@TestInstance(TestInstance.Lifecycle.PER_CLASS)
@ExtendWith(LogCapturingExtension.class)
class ConsumerExampleTest extends TestcontainersKafkaTest {

  private static final ActorSystem system = ActorSystem.create("ConsumerExampleTest");
  private static final Executor executor = Executors.newSingleThreadExecutor();

  ConsumerExampleTest() {
    super(system);
  }

  @AfterAll
  void afterClass() {
    TestKit.shutdownActorSystem(system);
  }

  void assertDone(CompletionStage<Done> stage) throws Exception {
    assertEquals(Done.done(), resultOf(stage));
  }

  private <T> Flow<T, T, NotUsed> business() {
    return Flow.create();
  }

  // #settings
  final Config config = system.settings().config().getConfig("pekko.kafka.consumer");
  final ConsumerSettings<String, byte[]> consumerSettings =
      ConsumerSettings.create(config, new StringDeserializer(), new ByteArrayDeserializer())
          .withBootstrapServers("localhost:9092")
          .withGroupId("group1")
          .withProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
  // #settings

  final ConsumerSettings<String, byte[]> consumerSettingsWithAutoCommit =
      // #settings-autocommit
      consumerSettings
          .withProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
          .withProperty(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, "5000");
  // #settings-autocommit

  @Test
  void plainSourceWithExternalOffsetStorage() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    // #plainSource
    final OffsetStorage db = new OffsetStorage();

    CompletionStage<Consumer.Control> controlCompletionStage =
        db.loadOffset()
            .thenApply(
                fromOffset ->
                    Consumer.plainSource(
                            consumerSettings,
                            Subscriptions.assignmentWithOffset(
                                new TopicPartition(topic, partition0), fromOffset))
                        .mapAsync(1, db::businessLogicAndStoreOffset)
                        .to(Sink.ignore())
                        .run(system));
    // #plainSource
    assertDone(produceString(topic, 10, partition0));
    while (db.offsetStore.get() < 9L) {
      sleepMillis(100, "until offsets have increased");
    }
    assertDone(resultOf(controlCompletionStage).shutdown());
  }

  // #plainSource

  class OffsetStorage {
    // #plainSource
    private final AtomicLong offsetStore = new AtomicLong(0L);

    // #plainSource
    public CompletionStage<Done> businessLogicAndStoreOffset(
        ConsumerRecord<String, String> record) { // ... }
      // #plainSource
      offsetStore.set(record.offset());
      return CompletableFuture.completedFuture(Done.getInstance());
    }

    // #plainSource
    public CompletionStage<Long> loadOffset() { // ... }
      // #plainSource

      return CompletableFuture.completedFuture(offsetStore.get());
    }

    public CompletionStage<Done> storeProcessedOffset(long offset) { // ... }
      offsetStore.set(offset);
      return CompletableFuture.completedFuture(Done.getInstance());
    }

    // #plainSource
  }
  // #plainSource

  @Test
  void configInheritance() throws Exception {
    // #config-inheritance
    Config config = system.settings().config().getConfig("our-kafka-consumer");
    ConsumerSettings<String, String> consumerSettings =
        ConsumerSettings.create(config, new StringDeserializer(), new StringDeserializer());
    // #config-inheritance
    assertEquals("kafka-host:9092", consumerSettings.getProperty("bootstrap.servers"));
  }

  @Test
  void atMostOnce() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    // #atMostOnce
    Consumer.Control control =
        Consumer.atMostOnceSource(consumerSettings, Subscriptions.topics(topic))
            .mapAsync(10, record -> business(record.key(), record.value()))
            .to(Sink.foreach(it -> System.out.println("Done with " + it)))
            .run(system);

    // #atMostOnce
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.shutdown());
  }

  // #atMostOnce #atLeastOnce
  CompletionStage<String> business(String key, String value) { // .... }
    // #atMostOnce #atLeastOnce
    return CompletableFuture.completedFuture("");
  }

  @Test
  void atLeastOnce() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    CommitterSettings committerSettings = committerDefaults();
    String topic = createTopic();
    // #atLeastOnce
    Consumer.DrainingControl<Done> control =
        Consumer.committableSource(consumerSettings, Subscriptions.topics(topic))
            .mapAsync(
                1,
                msg ->
                    business(msg.record().key(), msg.record().value())
                        .thenApply(done -> msg.committableOffset()))
            .toMat(
                Committer.sink(committerSettings.withMaxBatch(1)), Consumer::createDrainingControl)
            .run(system);

    // #atLeastOnce
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void atLeastOnceWithCommitterSink() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    Config config = system.settings().config().getConfig(CommitterSettings.configPath());
    // #committerSink
    CommitterSettings committerSettings = CommitterSettings.create(config);

    Consumer.DrainingControl<Done> control =
        Consumer.committableSource(consumerSettings, Subscriptions.topics(topic))
            .mapAsync(
                1,
                msg ->
                    business(msg.record().key(), msg.record().value())
                        .<ConsumerMessage.Committable>thenApply(done -> msg.committableOffset()))
            .toMat(Committer.sink(committerSettings), Consumer::createDrainingControl)
            .run(system);
    // #committerSink
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void commitWithMetadata() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    CommitterSettings committerSettings = committerDefaults();
    String topic = createTopic();
    // #commitWithMetadata
    Consumer.DrainingControl<Done> control =
        Consumer.commitWithMetadataSource(
                consumerSettings,
                Subscriptions.topics(topic),
                (record) -> Long.toString(record.timestamp()))
            .mapAsync(
                1,
                msg ->
                    business(msg.record().key(), msg.record().value())
                        .thenApply(done -> msg.committableOffset()))
            .toMat(Committer.sink(committerSettings), Consumer::createDrainingControl)
            .run(system);
    // #commitWithMetadata
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void consumerToProducer() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    ProducerSettings<String, String> producerSettings = producerDefaults();
    CommitterSettings committerSettings = committerDefaults();
    String topic1 = createTopic(1);
    String topic2 = createTopic(2);
    String targetTopic = createTopic(10);
    // #consumerToProducerSink
    Consumer.DrainingControl<Done> control =
        Consumer.committableSource(consumerSettings, Subscriptions.topics(topic1, topic2))
            .map(
                msg ->
                    ProducerMessage.<String, String, ConsumerMessage.Committable>single(
                        new ProducerRecord<>(targetTopic, msg.record().key(), msg.record().value()),
                        msg.committableOffset()))
            .toMat(
                Producer.committableSink(producerSettings, committerSettings),
                Consumer::createDrainingControl)
            .run(system);
    // #consumerToProducerSink
    assertDone(produceString(topic1, 10, partition0));
    assertDone(produceString(topic2, 10, partition0));
    Consumer.DrainingControl<List<ConsumerRecord<String, String>>> consumer =
        consumeString(targetTopic, 20);
    assertDone(consumer.isShutdown());
    assertEquals(20, resultOf(consumer.drainAndShutdown(executor)).size());
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void consumerToProducerWithContext() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    ProducerSettings<String, String> producerSettings = producerDefaults();
    CommitterSettings committerSettings = committerDefaults();
    String topic1 = createTopic(1);
    String topic2 = createTopic(2);
    String targetTopic = createTopic(10);
    // #consumerToProducerWithContext
    Consumer.DrainingControl<Done> control =
        Consumer.sourceWithOffsetContext(consumerSettings, Subscriptions.topics(topic1, topic2))
            .map(
                record ->
                    ProducerMessage.single(
                        new ProducerRecord<>(targetTopic, record.key(), record.value())))
            .toMat(
                Producer.committableSinkWithOffsetContext(producerSettings, committerSettings),
                Consumer::createDrainingControl)
            .run(system);
    // #consumerToProducerWithContext
    assertDone(produceString(topic1, 10, partition0));
    assertDone(produceString(topic2, 10, partition0));
    Consumer.DrainingControl<List<ConsumerRecord<String, String>>> consumer =
        consumeString(targetTopic, 20);
    assertDone(consumer.isShutdown());
    assertEquals(20, resultOf(consumer.drainAndShutdown(executor)).size());
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void committablePartitionedSource() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId()).withStopTimeout(Duration.ofMillis(10));
    CommitterSettings committerSettings = committerDefaults();
    int maxPartitions = 2;
    String topic = createTopic(1, maxPartitions);
    // #committablePartitionedSource
    Consumer.DrainingControl<Done> control =
        Consumer.committablePartitionedSource(consumerSettings, Subscriptions.topics(topic))
            .flatMapMerge(maxPartitions, Pair::second)
            .via(business())
            .map(msg -> msg.committableOffset())
            .toMat(Committer.sink(committerSettings), Consumer::createDrainingControl)
            .run(system);
    // #committablePartitionedSource
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void streamPerPartition() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId()).withStopTimeout(Duration.ofMillis(10));
    CommitterSettings committerSettings = committerDefaults();
    int maxPartitions = 2;
    String topic = createTopic(1, maxPartitions);
    // #committablePartitionedSource-stream-per-partition
    Consumer.DrainingControl<Done> control =
        Consumer.committablePartitionedSource(consumerSettings, Subscriptions.topics(topic))
            .mapAsyncUnordered(
                maxPartitions,
                pair -> {
                  Source<ConsumerMessage.CommittableMessage<String, String>, NotUsed> source =
                      pair.second();
                  return source
                      .via(business())
                      .map(message -> message.committableOffset())
                      .runWith(Committer.sink(committerSettings), system);
                })
            .toMat(Sink.ignore(), Consumer::createDrainingControl)
            .run(system);
    // #committablePartitionedSource-stream-per-partition
    assertDone(produceString(topic, 10, partition0));
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void consumerActor() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    ActorRef self = system.deadLetters();
    String topic = createTopic(1, 2);
    int partition0 = 0;
    int partition1 = 1;
    Behaviors.setup(
        ctx -> {
          // #consumerActorTyped

          // Consumer is represented by actor
          ActorRef consumer = Adapter.actorOf(ctx, KafkaConsumerActor.props(consumerSettings));
          // #consumerActorTyped
          return Behaviors.empty();
        });

    // #consumerActor
    // Consumer is represented by actor
    ActorRef consumer = system.actorOf((KafkaConsumerActor.props(consumerSettings)));

    // Manually assign topic partition to it
    Consumer.Control controlPartition1 =
        Consumer.plainExternalSource(
                consumer, Subscriptions.assignment(new TopicPartition(topic, partition0)))
            .via(business())
            .to(Sink.ignore())
            .run(system);

    // Manually assign another topic partition
    Consumer.Control controlPartition2 =
        Consumer.plainExternalSource(
                consumer, Subscriptions.assignment(new TopicPartition(topic, partition1)))
            .via(business())
            .to(Sink.ignore())
            .run(system);

    // #consumerActor
    CompletionStage<Done> producePartition0 = produceString(topic, 10, partition0);
    CompletionStage<Done> producePartition1 = produceString(topic, 10, partition1);
    assertDone(producePartition0);
    assertDone(producePartition1);
    assertDone(controlPartition1.shutdown());
    assertDone(controlPartition2.shutdown());
    // #consumerActor

    consumer.tell(KafkaConsumerActor.stop(), self);
    // #consumerActor
  }

  @Test
  void restartSource() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic(1, 2);
    // #restartSource
    AtomicReference<Consumer.Control> control = new AtomicReference<>(Consumer.createNoopControl());

    RestartSettings restartSettings =
        RestartSettings.create(Duration.ofSeconds(3), Duration.ofSeconds(30), 0.2);
    CompletionStage<Done> streamCompletion =
        RestartSource.onFailuresWithBackoff(
                restartSettings,
                () ->
                    Consumer.plainSource(consumerSettings, Subscriptions.topics(topic))
                        .mapMaterializedValue(
                            c -> {
                              // this is a hack to get access to the Consumer.Control
                              // instances of the latest Kafka Consumer source
                              control.set(c);
                              return c;
                            })
                        .via(business()))
            .runWith(Sink.ignore(), system);

    // #restartSource
    assertDone(produceString(topic, 10, partition0));
    // #restartSource
    control.get().drainAndShutdown(streamCompletion, system.getDispatcher());
    // #restartSource
  }

  // #withRebalanceListenerActor
  static class RebalanceListener extends AbstractLoggingActor {

    @Override
    public Receive createReceive() {
      return receiveBuilder()
          .match(
              TopicPartitionsAssigned.class,
              assigned -> {
                log().info("Assigned: {}", assigned);
              })
          .match(
              TopicPartitionsRevoked.class,
              revoked -> {
                log().info("Revoked: {}", revoked);
              })
          .build();
    }
  }

  // #withRebalanceListenerActor

  @Test
  void withRebalanceListener() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    int messageCount = 10;
    // #withRebalanceListenerActor
    ActorRef rebalanceListener = system.actorOf(Props.create(RebalanceListener.class));

    Subscription subscription =
        Subscriptions.topics(topic)
            // additionally, pass the actor reference:
            .withRebalanceListener(rebalanceListener);

    // use the subscription as usual:
    Consumer.DrainingControl<List<ConsumerRecord<String, String>>> control =
        Consumer.plainSource(consumerSettings, subscription)
            // #withRebalanceListenerActor
            .take(messageCount)
            // #withRebalanceListenerActor
            .toMat(Sink.seq(), Consumer::createDrainingControl)
            .run(system);
    // #withRebalanceListenerActor
    assertDone(produceString(topic, messageCount, partition0));
    assertDone(control.isShutdown());
    assertEquals(messageCount, resultOf(control.drainAndShutdown(executor)).size());
  }

  @Test
  void withTypedRebalanceListener() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    int messageCount = 10;

    // #withTypedRebalanceListenerActor

    Function<ActorContext<ConsumerRebalanceEvent>, Behavior<ConsumerRebalanceEvent>>
        rebalanceListener =
            (ActorContext<ConsumerRebalanceEvent> context) ->
                Behaviors.receive(ConsumerRebalanceEvent.class)
                    .onMessage(
                        TopicPartitionsAssigned.class,
                        assigned -> {
                          context.getLog().info("Assigned: {}", assigned);
                          return Behaviors.same();
                        })
                    .onMessage(
                        TopicPartitionsRevoked.class,
                        revoked -> {
                          context.getLog().info("Revoked: {}", revoked);
                          return Behaviors.same();
                        })
                    .build();
    // #withTypedRebalanceListenerActor

    Behavior<Object> guardian =
        Behaviors.setup(
            guardianCtx -> {
              // #withTypedRebalanceListenerActor

              Behavior<ConsumerRebalanceEvent> listener =
                  Behaviors.setup(ctx -> rebalanceListener.apply(ctx));

              org.apache.pekko.actor.typed.ActorRef<ConsumerRebalanceEvent> typedRef =
                  guardianCtx.spawn(listener, "rebalance-listener");

              org.apache.pekko.actor.ActorRef classicRef = Adapter.toClassic(typedRef);

              Subscription subscription =
                  Subscriptions.topics(topic)
                      // additionally, pass the actor reference:
                      .withRebalanceListener(classicRef);

              Consumer.DrainingControl<List<ConsumerRecord<String, String>>> control =
                  // use the subscription as usual:
                  Consumer.plainSource(consumerSettings, subscription)
                      // #withTypedRebalanceListenerActor
                      .take(messageCount)
                      // #withTypedRebalanceListenerActor
                      .toMat(Sink.seq(), Consumer::createDrainingControl)
                      .run(system);
              // #withTypedRebalanceListenerActor

              assertDone(produceString(topic, messageCount, partition0));
              assertDone(control.isShutdown());
              assertEquals(messageCount, resultOf(control.drainAndShutdown(executor)).size());

              return Behaviors.stopped();
            });

    org.apache.pekko.actor.typed.ActorSystem<Object> typed =
        org.apache.pekko.actor.typed.ActorSystem.create(
            guardian, "typed-rebalance-listener-example");
    assertDone(typed.getWhenTerminated());
  }

  @Test
  void withPartitionAssignmentHandler() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    TopicPartition tp = new TopicPartition(topic, 0);
    int messageCount = 10;
    AtomicReference<Set<TopicPartition>> revoked = new AtomicReference<>();
    AtomicReference<Set<TopicPartition>> assigned = new AtomicReference<>();
    AtomicReference<Set<TopicPartition>> stopped = new AtomicReference<>();

    // #partitionAssignmentHandler
    PartitionAssignmentHandler assignmentHandler =
        new PartitionAssignmentHandler() {
          public void onRevoke(Set<TopicPartition> revokedTps, RestrictedConsumer consumer) {
            // #partitionAssignmentHandler
            revoked.set(revokedTps);
            // #partitionAssignmentHandler
          }

          public void onAssign(Set<TopicPartition> assignedTps, RestrictedConsumer consumer) {
            // #partitionAssignmentHandler
            assigned.set(assignedTps);
            // #partitionAssignmentHandler
          }

          public void onLost(Set<TopicPartition> lostTps, RestrictedConsumer consumer) {}

          public void onStop(Set<TopicPartition> currentTps, RestrictedConsumer consumer) {
            // #partitionAssignmentHandler
            stopped.set(currentTps);
            // #partitionAssignmentHandler
          }
        };

    Subscription subscription =
        Subscriptions.topics(topic).withPartitionAssignmentHandler(assignmentHandler);
    // #partitionAssignmentHandler

    Consumer.DrainingControl<List<ConsumerRecord<String, String>>> control =
        Consumer.plainSource(consumerSettings, subscription)
            .take(messageCount)
            .toMat(Sink.seq(), Consumer::createDrainingControl)
            .run(system);
    assertDone(produceString(topic, messageCount, partition0));
    assertDone(control.isShutdown());
    assertEquals(messageCount, resultOf(control.drainAndShutdown(executor)).size());

    assertThat(assigned.get(), hasItem(tp));
    assertThat(stopped.get(), hasItem(tp));
    assertThat(revoked.get(), hasItem(tp)); // revoke of partitions occurs after stop
  }

  @Test
  void consumerMetrics() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    // #consumerMetrics
    // run the stream to obtain the materialized Control value
    Consumer.DrainingControl<Done> control =
        Consumer.plainSource(
                consumerSettings, Subscriptions.assignment(new TopicPartition(topic, 0)))
            .via(business())
            .toMat(Sink.ignore(), Consumer::createDrainingControl)
            .run(system);

    // #consumerMetrics
    sleepMillis(
        100,
        "to let the control establish itself (fails with `java.lang.IllegalStateException: not yet initialized: only setHandler is allowed in GraphStageLogic constructor)` otherwise");
    // #consumerMetrics
    CompletionStage<Map<MetricName, Metric>> metrics = control.getMetrics();
    metrics.thenAccept(map -> System.out.println("Metrics: " + map));
    // #consumerMetrics
    assertDone(control.drainAndShutdown(executor));
  }

  @Test
  void shutdownPlainSource() {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    // #shutdownPlainSource
    final OffsetStorage db = new OffsetStorage();

    CompletionStage<Consumer.DrainingControl<Done>> control =
        db.loadOffset()
            .thenApply(
                fromOffset ->
                    Consumer.plainSource(
                            consumerSettings,
                            Subscriptions.assignmentWithOffset(
                                new TopicPartition(topic, 0), fromOffset))
                        .mapAsync(
                            10,
                            record ->
                                business(record.key(), record.value())
                                    .thenApply(res -> db.storeProcessedOffset(record.offset())))
                        .toMat(Sink.ignore(), Consumer::createDrainingControl)
                        .run(system));

    // Shutdown the consumer when desired
    control.thenAccept(c -> c.drainAndShutdown(executor));
    // #shutdownPlainSource
  }

  @Test
  void shutdownCommittable() throws Exception {
    ConsumerSettings<String, String> consumerSettings =
        consumerDefaults().withGroupId(createGroupId());
    String topic = createTopic();
    CommitterSettings committerSettings = committerDefaults();
    // TODO there is a problem with combining `take` and committing.
    int messageCount = 1;
    // #shutdownCommittableSource
    final Executor ec = Executors.newCachedThreadPool();

    Consumer.DrainingControl<Done> control =
        Consumer.committableSource(
                consumerSettings.withStopTimeout(Duration.ZERO), Subscriptions.topics(topic))
            .mapAsync(
                1,
                msg ->
                    business(msg.record().key(), msg.record().value())
                        .thenApply(done -> msg.committableOffset()))
            // #shutdownCommittableSource
            .take(messageCount)
            // #shutdownCommittableSource
            .toMat(
                Committer.sink(committerSettings.withMaxBatch(1)), Consumer::createDrainingControl)
            .run(system);

    // #shutdownCommittableSource
    assertDone(produceString(topic, messageCount, partition0));
    assertDone(control.isShutdown());
    assertEquals(Done.done(), resultOf(control.drainAndShutdown(ec), Duration.ofSeconds(20)));
    // #shutdownCommittableSource
    control.drainAndShutdown(ec);
    // #shutdownCommittableSource
  }
}