content: /*
   * 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.kafka.streams.integration;
  
  import org.apache.kafka.clients.consumer.ConsumerConfig;
  import org.apache.kafka.clients.consumer.ConsumerRecord;
  import org.apache.kafka.clients.producer.ProducerConfig;
  import org.apache.kafka.common.serialization.IntegerDeserializer;
  import org.apache.kafka.common.serialization.IntegerSerializer;
  import org.apache.kafka.common.serialization.Serdes;
  import org.apache.kafka.common.utils.Utils;
  import org.apache.kafka.streams.KafkaStreams;
  import org.apache.kafka.streams.KeyValue;
  import org.apache.kafka.streams.StoreQueryParameters;
  import org.apache.kafka.streams.StreamsBuilder;
  import org.apache.kafka.streams.StreamsConfig;
  import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
  import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
  import org.apache.kafka.streams.processor.api.Processor;
  import org.apache.kafka.streams.processor.api.ProcessorContext;
  import org.apache.kafka.streams.processor.api.Record;
  import org.apache.kafka.streams.state.KeyValueStore;
  import org.apache.kafka.streams.state.QueryableStoreTypes;
  import org.apache.kafka.streams.state.Stores;
  import org.apache.kafka.test.TestUtils;
  
  import org.hamcrest.Matchers;
  import org.junit.jupiter.api.AfterAll;
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeAll;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Tag;
  import org.junit.jupiter.api.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.time.Duration;
  import java.util.Collections;
  import java.util.List;
  import java.util.Properties;
  import java.util.UUID;
  import java.util.stream.Collectors;
  import java.util.stream.IntStream;
  
  import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
  import static org.apache.kafka.test.TestUtils.waitForCondition;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.is;
  
  @Tag("integration")
  public class StandbyTaskEOSMultiRebalanceIntegrationTest {
  
      private static final Logger LOG = LoggerFactory.getLogger(StandbyTaskEOSMultiRebalanceIntegrationTest.class);
  
      private static final long TWO_MINUTE_TIMEOUT = Duration.ofMinutes(2L).toMillis();
  
      private String appId;
      private String inputTopic;
      private String storeName;
      private String counterName;
  
      private String outputTopic;
  
      private KafkaStreams streamInstanceOne;
      private KafkaStreams streamInstanceTwo;
      private KafkaStreams streamInstanceThree;
  
      private static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(3);
  
      @BeforeAll
      public static void startCluster() throws IOException {
          CLUSTER.start();
      }
  
      @AfterAll
      public static void closeCluster() {
          CLUSTER.stop();
      }
  
      @BeforeEach
      public void createTopics() throws Exception {
          final String safeTestName = UUID.randomUUID().toString();
          appId = "app-" + safeTestName;
          inputTopic = "input-" + safeTestName;
          outputTopic = "output-" + safeTestName;
          storeName = "store-" + safeTestName;
          counterName = "counter-" + safeTestName;
  
          CLUSTER.deleteTopics(inputTopic, outputTopic);
          CLUSTER.createTopic(inputTopic, partitionCount, 3);
          CLUSTER.createTopic(outputTopic, partitionCount, 3);
      }
  
      private final int partitionCount = 12;
  
      @AfterEach
      public void cleanUp() {
          if (streamInstanceOne != null) {
              streamInstanceOne.close(Duration.ofSeconds(60));
          }
          if (streamInstanceTwo != null) {
              streamInstanceTwo.close(Duration.ofSeconds(60));
          }
          if (streamInstanceThree != null) {
              streamInstanceThree.close(Duration.ofSeconds(60));
          }
      }
  
      // The test produces a duplicate fee range of integers from 0 (inclusive) to initialBulk + secondBulk (exclusive) as input to the stream.
      // The stream is responsible for assigning a unique id to each of the integers.
      // The output topic must thus contain 63000 message:
      //      The Key is unique and from the range of input values
      //      The Values produced are unique.
      @Test
      public void shouldHonorEOSWhenUsingCachingAndStandbyReplicas() throws Exception {
          final Properties readCommitted = new Properties();
          readCommitted.setProperty("isolation.level", "read_committed");
          final long time = System.currentTimeMillis();
          final String base = TestUtils.tempDirectory(appId).getPath();
  
          final int initialBulk = 3000;
          final int secondBulk = 60000;
  
          IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp(
                  inputTopic,
                  IntStream.range(0, initialBulk).boxed().map(i -> new KeyValue<>(i, i)).collect(Collectors.toList()),
                  TestUtils.producerConfig(
                          CLUSTER.bootstrapServers(),
                          IntegerSerializer.class,
                          IntegerSerializer.class,
                          Utils.mkProperties(Collections.singletonMap(ProducerConfig.ACKS_CONFIG, "all"))
                  ),
                  10L + time
          );
  
          streamInstanceOne = buildWithUniqueIdAssignmentTopology(base + "-1");
          streamInstanceTwo = buildWithUniqueIdAssignmentTopology(base + "-2");
          streamInstanceThree = buildWithUniqueIdAssignmentTopology(base + "-3");
  
          LOG.info("start first instance and wait for completed processing");
          startApplicationAndWaitUntilRunning(Collections.singletonList(streamInstanceOne), Duration.ofSeconds(30));
          IntegrationTestUtils.waitUntilMinRecordsReceived(
                  TestUtils.consumerConfig(
                          CLUSTER.bootstrapServers(),
                          UUID.randomUUID().toString(),
                          IntegerDeserializer.class,
                          IntegerDeserializer.class,
                          readCommitted
                  ),
                  outputTopic,
                  initialBulk
          );
          LOG.info("Finished reading the initial bulk");
  
          LOG.info("start second instance and wait for standby replication");
          startApplicationAndWaitUntilRunning(Collections.singletonList(streamInstanceTwo), Duration.ofSeconds(30));
          waitForCondition(
                  () -> streamInstanceTwo.store(
                          StoreQueryParameters.fromNameAndType(
                                  storeName,
                                  QueryableStoreTypes.<Integer, Integer>keyValueStore()
                          ).enableStaleStores()
                  ).get(0) != null,
                  TWO_MINUTE_TIMEOUT,
                  "Could not get key from standby store"
          );
          LOG.info("Second stream have some data in the state store");
  
  
          LOG.info("Produce the second bulk");
          IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp(
                  inputTopic,
                  IntStream.range(initialBulk, initialBulk + secondBulk).boxed().map(i -> new KeyValue<>(i, i)).collect(Collectors.toList()),
                  TestUtils.producerConfig(
                          CLUSTER.bootstrapServers(),
                          IntegerSerializer.class,
                          IntegerSerializer.class,
                          new Properties()
                  ),
                  1000L + time
          );
  
          LOG.info("Start stream three which will introduce a re-balancing event and hopefully some redistribution of tasks.");
          startApplicationAndWaitUntilRunning(Collections.singletonList(streamInstanceThree), Duration.ofSeconds(90));
  
          LOG.info("Wait for the processing to be completed");
          final List<ConsumerRecord<Integer, Integer>> outputRecords = IntegrationTestUtils.waitUntilMinRecordsReceived(
                  TestUtils.consumerConfig(
                          CLUSTER.bootstrapServers(),
                          UUID.randomUUID().toString(),
                          IntegerDeserializer.class,
                          IntegerDeserializer.class,
                          readCommitted
                  ),
                  outputTopic,
                  initialBulk + secondBulk,
                  Duration.ofMinutes(10L).toMillis()
          );
          LOG.info("Processing completed");
  
          outputRecords.stream().collect(Collectors.groupingBy(ConsumerRecord::value)).forEach(this::logIfDuplicate);
  
          assertThat("Each output should correspond to one distinct value", outputRecords.stream().map(ConsumerRecord::value).distinct().count(), is(Matchers.equalTo((long) outputRecords.size())));
      }
  
      private void logIfDuplicate(final Integer id, final List<ConsumerRecord<Integer, Integer>> record) {
          assertThat("The id and the value in the records must match", record.stream().allMatch(r -> id.equals(r.value())));
          if (record.size() > 1) {
              LOG.warn("Id : " + id + " is assigned to the following " + record.stream().map(ConsumerRecord::key).collect(Collectors.toList()));
          }
      }
  
      private KafkaStreams buildWithUniqueIdAssignmentTopology(final String stateDirPath) {
          final StreamsBuilder builder = new StreamsBuilder();
  
          builder.addStateStore(Stores.keyValueStoreBuilder(
                  Stores.persistentKeyValueStore(storeName),
                  Serdes.Integer(),
                  Serdes.Integer())
          );
          builder.addStateStore(Stores.keyValueStoreBuilder(
                  Stores.persistentKeyValueStore(counterName),
                  Serdes.Integer(),
                  Serdes.Integer()).withCachingEnabled()
          );
          builder.<Integer, Integer>stream(inputTopic)
                  .process(
                          () -> new Processor<Integer, Integer, Integer, Integer>() {
                              private KeyValueStore<Integer, Integer> store;
                              private KeyValueStore<Integer, Integer> counter;
                              private ProcessorContext<Integer, Integer> context;
  
                              @Override
                              public void init(final ProcessorContext<Integer, Integer> context) {
                                  this.context = context;
                                  store = context.getStateStore(storeName);
                                  counter = context.getStateStore(counterName);
                              }
  
                              @Override
                              public void process(final Record<Integer, Integer> record) {
                                  final Integer key = record.key();
                                  final Integer unused = record.value();
                                  assertThat("Key and value mus be equal", key.equals(unused));
                                  Integer id = store.get(key);
                                  // Only assign a new id if the value have not been observed before
                                  if (id == null) {
                                      final int counterKey = 0;
                                      final Integer lastCounter = counter.get(counterKey);
                                      final int newCounter = lastCounter == null ? 0 : lastCounter + 1;
                                      counter.put(counterKey, newCounter);
                                      // Partitions assign ids from their own id space
                                      id = newCounter * partitionCount + context.recordMetadata().get().partition();
                                      store.put(key, id);
                                  }
                                  context.forward(record.withKey(id));
                              }
  
                              @Override
                              public void close() {
                              }
                          },
                          storeName, counterName
                  )
                  .to(outputTopic);
          return new KafkaStreams(builder.build(), props(stateDirPath));
      }
  
  
      private Properties props(final String stateDirPath) {
          final Properties streamsConfiguration = new Properties();
          streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, appId);
          streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
          streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, stateDirPath);
          streamsConfiguration.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
          streamsConfiguration.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, 4);
  
          streamsConfiguration.put(StreamsConfig.MAX_WARMUP_REPLICAS_CONFIG, 1);
          streamsConfiguration.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
          streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
          streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
          streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
          streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
  
          return streamsConfiguration;
      }
  }
  