/*
 * 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.admin.Admin;
import org.apache.kafka.clients.admin.AdminClientConfig;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.KafkaConsumer;
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.common.IsolationLevel;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.IntegerDeserializer;
import org.apache.kafka.common.serialization.IntegerSerializer;
import org.apache.kafka.common.serialization.LongDeserializer;
import org.apache.kafka.common.serialization.LongSerializer;
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.StreamsBuilder;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.StreamsConfig.InternalConfig;
import org.apache.kafka.streams.Topology;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.errors.StreamsUncaughtExceptionHandler;
import org.apache.kafka.streams.errors.TaskCorruptedException;
import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
import org.apache.kafka.streams.kstream.KStream;
import org.apache.kafka.streams.processor.StateRestoreListener;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.api.ContextualProcessor;
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.processor.internals.DefaultKafkaClientSupplier;
import org.apache.kafka.streams.processor.internals.StreamThread;
import org.apache.kafka.streams.query.QueryResult;
import org.apache.kafka.streams.query.RangeQuery;
import org.apache.kafka.streams.query.StateQueryRequest;
import org.apache.kafka.streams.query.StateQueryResult;
import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.KeyValueStore;
import org.apache.kafka.streams.state.StoreBuilder;
import org.apache.kafka.streams.state.Stores;
import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
import org.apache.kafka.test.StreamsTestUtils;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.DEFAULT_TIMEOUT;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.purgeLocalStreamsState;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitForEmptyConsumerGroup;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitUntilMinRecordsReceived;
import static org.apache.kafka.streams.query.StateQueryRequest.inStore;
import static org.apache.kafka.streams.utils.TestUtils.waitForApplicationState;
import static org.apache.kafka.test.TestUtils.consumerConfig;
import static org.apache.kafka.test.TestUtils.waitForCondition;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

@Tag("integration")
@Timeout(600)
public class EosIntegrationTest {
    private static final Logger LOG = LoggerFactory.getLogger(EosIntegrationTest.class);
    private static final int NUM_BROKERS = 3;
    private static final int MAX_POLL_INTERVAL_MS = 30_000;
    private static final int MAX_WAIT_TIME_MS = 120_000;

    public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(
        NUM_BROKERS,
        Utils.mkProperties(mkMap(
            mkEntry("auto.create.topics.enable", "true"),
            mkEntry("transaction.max.timeout.ms", "" + Integer.MAX_VALUE)
        ))
    );

    @BeforeAll
    public static void startCluster() throws IOException {
        CLUSTER.start();
    }

    @AfterAll
    public static void closeCluster() {
        CLUSTER.stop();
    }


    private String applicationId;
    private static final int NUM_TOPIC_PARTITIONS = 2;
    private static final String CONSUMER_GROUP_ID = "readCommitted";
    private static final String SINGLE_PARTITION_INPUT_TOPIC = "singlePartitionInputTopic";
    private static final String SINGLE_PARTITION_THROUGH_TOPIC = "singlePartitionThroughTopic";
    private static final String SINGLE_PARTITION_OUTPUT_TOPIC = "singlePartitionOutputTopic";
    private static final String MULTI_PARTITION_INPUT_TOPIC = "multiPartitionInputTopic";
    private static final String MULTI_PARTITION_THROUGH_TOPIC = "multiPartitionThroughTopic";
    private static final String MULTI_PARTITION_OUTPUT_TOPIC = "multiPartitionOutputTopic";
    private final String storeName = "store";

    private AtomicBoolean errorInjected;
    private AtomicBoolean stallInjected;
    private AtomicReference<String> stallingHost;
    private volatile boolean doStall = true;
    private AtomicInteger commitRequested;
    private Throwable uncaughtException;

    private static final AtomicInteger TEST_NUMBER = new AtomicInteger(0);

    private volatile boolean hasUnexpectedError = false;

    private String stateTmpDir;

    private static java.util.stream.Stream<Arguments> groupProtocolAndProcessingThreadsParameters() {
        return java.util.stream.Stream.of(
            Arguments.of("classic", true),
            Arguments.of("classic", false),
            Arguments.of("streams", true),
            Arguments.of("streams", false)
        );
    }

    @BeforeEach
    public void createTopics() throws Exception {
        applicationId = "appId-" + TEST_NUMBER.getAndIncrement();
        CLUSTER.deleteTopics(
            SINGLE_PARTITION_INPUT_TOPIC, MULTI_PARTITION_INPUT_TOPIC,
            SINGLE_PARTITION_THROUGH_TOPIC, MULTI_PARTITION_THROUGH_TOPIC,
            SINGLE_PARTITION_OUTPUT_TOPIC, MULTI_PARTITION_OUTPUT_TOPIC);

        CLUSTER.createTopics(SINGLE_PARTITION_INPUT_TOPIC, SINGLE_PARTITION_THROUGH_TOPIC, SINGLE_PARTITION_OUTPUT_TOPIC);
        CLUSTER.createTopic(MULTI_PARTITION_INPUT_TOPIC, NUM_TOPIC_PARTITIONS, 1);
        CLUSTER.createTopic(MULTI_PARTITION_THROUGH_TOPIC, NUM_TOPIC_PARTITIONS, 1);
        CLUSTER.createTopic(MULTI_PARTITION_OUTPUT_TOPIC, NUM_TOPIC_PARTITIONS, 1);
        CLUSTER.setGroupStandbyReplicas(applicationId, 1);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToRunWithEosEnabled(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, SINGLE_PARTITION_INPUT_TOPIC, null, SINGLE_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldCommitCorrectOffsetIfInputTopicIsTransactional(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, SINGLE_PARTITION_INPUT_TOPIC, null, SINGLE_PARTITION_OUTPUT_TOPIC, true, groupProtocol);

        try (final Admin adminClient = Admin.create(mkMap(mkEntry(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers())));
             final Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(mkMap(
                 mkEntry(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()),
                 mkEntry(ConsumerConfig.GROUP_ID_CONFIG, applicationId),
                 mkEntry(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class),
                 mkEntry(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class)))) {

            waitForEmptyConsumerGroup(adminClient, applicationId, 5 * MAX_POLL_INTERVAL_MS);

            final TopicPartition topicPartition = new TopicPartition(SINGLE_PARTITION_INPUT_TOPIC, 0);
            final Collection<TopicPartition> topicPartitions = Collections.singleton(topicPartition);

            final long committedOffset = adminClient.listConsumerGroupOffsets(applicationId).partitionsToOffsetAndMetadata().get().get(topicPartition).offset();

            consumer.assign(topicPartitions);
            final long consumerPosition = consumer.position(topicPartition);
            final long endOffset = consumer.endOffsets(topicPartitions).get(topicPartition);

            assertThat(committedOffset, equalTo(consumerPosition));
            assertThat(committedOffset, equalTo(endOffset));
        }
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToRestartAfterClose(final String groupProtocol) throws Exception {
        runSimpleCopyTest(2, SINGLE_PARTITION_INPUT_TOPIC, null, SINGLE_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToCommitToMultiplePartitions(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, SINGLE_PARTITION_INPUT_TOPIC, null, MULTI_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToCommitMultiplePartitionOffsets(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, MULTI_PARTITION_INPUT_TOPIC, null, SINGLE_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToRunWithTwoSubtopologies(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, SINGLE_PARTITION_INPUT_TOPIC, SINGLE_PARTITION_THROUGH_TOPIC, SINGLE_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToRunWithTwoSubtopologiesAndMultiplePartitions(final String groupProtocol) throws Exception {
        runSimpleCopyTest(1, MULTI_PARTITION_INPUT_TOPIC, MULTI_PARTITION_THROUGH_TOPIC, MULTI_PARTITION_OUTPUT_TOPIC, false, groupProtocol);
    }

    private void runSimpleCopyTest(final int numberOfRestarts,
                                   final String inputTopic,
                                   final String throughTopic,
                                   final String outputTopic,
                                   final boolean inputTopicTransactional,
                                   final String groupProtocol) throws Exception {
        final StreamsBuilder builder = new StreamsBuilder();
        final KStream<Long, Long> input = builder.stream(inputTopic);
        KStream<Long, Long> output = input;
        if (throughTopic != null) {
            input.to(throughTopic);
            output = builder.stream(throughTopic);
        }
        output.to(outputTopic);

        final Properties properties = new Properties();
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        properties.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_RECORDS_CONFIG), 1);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.METADATA_MAX_AGE_CONFIG), "1000");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG), "earliest");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), MAX_POLL_INTERVAL_MS - 1);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG), MAX_POLL_INTERVAL_MS);
        properties.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);

        for (int i = 0; i < numberOfRestarts; ++i) {
            final Properties config = StreamsTestUtils.getStreamsConfig(
                applicationId,
                CLUSTER.bootstrapServers(),
                Serdes.LongSerde.class.getName(),
                Serdes.LongSerde.class.getName(),
                properties);

            final List<KeyValue<Long, Long>> inputData = prepareData(i * 100, i * 100 + 10L, 0L, 1L);

            final Properties producerConfigs = new Properties();
            if (inputTopicTransactional) {
                producerConfigs.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, applicationId + "-input-producer");
            }

            IntegrationTestUtils.produceKeyValuesSynchronously(
                inputTopic,
                inputData,
                TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class, producerConfigs),
                CLUSTER.time,
                inputTopicTransactional
            );

            try (final KafkaStreams streams = new KafkaStreams(builder.build(), config)) {
                startApplicationAndWaitUntilRunning(streams);

                final List<KeyValue<Long, Long>> committedRecords = readResult(outputTopic, inputData.size(), CONSUMER_GROUP_ID);
                checkResultPerKey(committedRecords, inputData, "The committed records do not match what expected");
            }
        }
    }

    private void checkResultPerKey(final List<KeyValue<Long, Long>> result,
                                   final List<KeyValue<Long, Long>> expectedResult,
                                   final String reason) {
        final Set<Long> allKeys = new HashSet<>();
        addAllKeys(allKeys, result);
        addAllKeys(allKeys, expectedResult);

        for (final Long key : allKeys) {
            assertThat(reason, getAllRecordPerKey(key, result), equalTo(getAllRecordPerKey(key, expectedResult)));
        }
    }

    private void addAllKeys(final Set<Long> allKeys, final List<KeyValue<Long, Long>> records) {
        for (final KeyValue<Long, Long> record : records) {
            allKeys.add(record.key);
        }
    }

    private List<KeyValue<Long, Long>> getAllRecordPerKey(final Long key, final List<KeyValue<Long, Long>> records) {
        final List<KeyValue<Long, Long>> recordsPerKey = new ArrayList<>(records.size());

        for (final KeyValue<Long, Long> record : records) {
            if (record.key.equals(key)) {
                recordsPerKey.add(record);
            }
        }

        return recordsPerKey;
    }

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldBeAbleToPerformMultipleTransactions(final String groupProtocol) throws Exception {
        final StreamsBuilder builder = new StreamsBuilder();
        builder.stream(SINGLE_PARTITION_INPUT_TOPIC).to(SINGLE_PARTITION_OUTPUT_TOPIC);

        final Properties properties = new Properties();
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        properties.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
        properties.put(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "1000");
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        properties.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);

        final Properties config = StreamsTestUtils.getStreamsConfig(
            applicationId,
            CLUSTER.bootstrapServers(),
            Serdes.LongSerde.class.getName(),
            Serdes.LongSerde.class.getName(),
            properties);

        try (final KafkaStreams streams = new KafkaStreams(builder.build(), config)) {
            startApplicationAndWaitUntilRunning(streams);

            final List<KeyValue<Long, Long>> firstBurstOfData = prepareData(0L, 5L, 0L);
            final List<KeyValue<Long, Long>> secondBurstOfData = prepareData(5L, 8L, 0L);

            IntegrationTestUtils.produceKeyValuesSynchronously(
                SINGLE_PARTITION_INPUT_TOPIC,
                firstBurstOfData,
                TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class),
                CLUSTER.time
            );

            final List<KeyValue<Long, Long>> firstCommittedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, firstBurstOfData.size(), CONSUMER_GROUP_ID);
            assertThat(firstCommittedRecords, equalTo(firstBurstOfData));

            IntegrationTestUtils.produceKeyValuesSynchronously(
                SINGLE_PARTITION_INPUT_TOPIC,
                secondBurstOfData,
                TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class),
                CLUSTER.time
            );

            final List<KeyValue<Long, Long>> secondCommittedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, secondBurstOfData.size(), CONSUMER_GROUP_ID);
            assertThat(secondCommittedRecords, equalTo(secondBurstOfData));
        }
    }

    @ParameterizedTest
    @MethodSource("groupProtocolAndProcessingThreadsParameters")
    public void shouldNotViolateEosIfOneTaskFails(final String groupProtocol, final boolean processingThreadsEnabled) throws Exception {

        // this test writes 10 + 5 + 5 records per partition (running with 2 partitions)
        // the app is supposed to copy all 40 records into the output topic
        //
        // the app first commits after each 10 records per partition(total 20 records), and thus will have 2 * 5 uncommitted writes
        //
        // the failure gets inject after 20 committed and 30 uncommitted records got received
        // -> the failure only kills one thread
        // after fail over, we should read 40 committed records (even if 50 record got written)

        try (final KafkaStreams streams = getKafkaStreams("dummy", false, "appDir", 2, groupProtocol, processingThreadsEnabled)) {
            startApplicationAndWaitUntilRunning(streams);

            final List<KeyValue<Long, Long>> committedDataBeforeFailure = prepareData(0L, 10L, 0L, 1L);
            final List<KeyValue<Long, Long>> uncommittedDataBeforeFailure = prepareData(10L, 15L, 0L, 1L);

            final List<KeyValue<Long, Long>> dataBeforeFailure = new ArrayList<>(
                committedDataBeforeFailure.size() + uncommittedDataBeforeFailure.size());
            dataBeforeFailure.addAll(committedDataBeforeFailure);
            dataBeforeFailure.addAll(uncommittedDataBeforeFailure);

            final List<KeyValue<Long, Long>> dataAfterFailure = prepareData(15L, 20L, 0L, 1L);

            writeInputData(committedDataBeforeFailure);

            waitForCondition(
                () -> commitRequested.get() == 2, MAX_WAIT_TIME_MS,
                "StreamsTasks did not request commit.");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C
            // p-1: ---> 10 rec + C

            final List<KeyValue<Long, Long>> committedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, committedDataBeforeFailure.size(), CONSUMER_GROUP_ID);
            checkResultPerKey(
                committedRecords,
                committedDataBeforeFailure,
                "The committed records before failure do not match what expected");

            writeInputData(uncommittedDataBeforeFailure);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec (pending)
            // p-1: ---> 10 rec + C  + 5 rec (pending)

            final List<KeyValue<Long, Long>> uncommittedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, dataBeforeFailure.size(), null);
            checkResultPerKey(
                uncommittedRecords,
                dataBeforeFailure,
                "The uncommitted records before failure do not match what expected");
            errorInjected.set(true);
            writeInputData(dataAfterFailure);

            waitForCondition(
                () -> uncaughtException != null, MAX_WAIT_TIME_MS,
                "Should receive uncaught exception from one StreamThread.");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C

            final List<KeyValue<Long, Long>> allCommittedRecords = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                committedDataBeforeFailure.size() + uncommittedDataBeforeFailure.size() + dataAfterFailure.size(),
                CONSUMER_GROUP_ID + "_ALL");

            final List<KeyValue<Long, Long>> committedRecordsAfterFailure = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                uncommittedDataBeforeFailure.size() + dataAfterFailure.size(),
                CONSUMER_GROUP_ID);

            final int allCommittedRecordsAfterRecoverySize = committedDataBeforeFailure.size() +
                uncommittedDataBeforeFailure.size() + dataAfterFailure.size();
            final List<KeyValue<Long, Long>> allExpectedCommittedRecordsAfterRecovery = new ArrayList<>(allCommittedRecordsAfterRecoverySize);
            allExpectedCommittedRecordsAfterRecovery.addAll(committedDataBeforeFailure);
            allExpectedCommittedRecordsAfterRecovery.addAll(uncommittedDataBeforeFailure);
            allExpectedCommittedRecordsAfterRecovery.addAll(dataAfterFailure);

            final int committedRecordsAfterRecoverySize = uncommittedDataBeforeFailure.size() + dataAfterFailure.size();
            final List<KeyValue<Long, Long>> expectedCommittedRecordsAfterRecovery = new ArrayList<>(committedRecordsAfterRecoverySize);
            expectedCommittedRecordsAfterRecovery.addAll(uncommittedDataBeforeFailure);
            expectedCommittedRecordsAfterRecovery.addAll(dataAfterFailure);

            checkResultPerKey(
                allCommittedRecords,
                allExpectedCommittedRecordsAfterRecovery,
                "The all committed records after recovery do not match what expected");
            checkResultPerKey(
                committedRecordsAfterFailure,
                expectedCommittedRecordsAfterRecovery,
                "The committed records after recovery do not match what expected");

            assertThat("Should only get one uncaught exception from Streams.", hasUnexpectedError, is(false));
        }
    }

    @ParameterizedTest
    @MethodSource("groupProtocolAndProcessingThreadsParameters")
    public void shouldNotViolateEosIfOneTaskFailsWithState(final String groupProtocol, final boolean processingThreadsEnabled) throws Exception {

        // this test updates a store with 10 + 5 + 5 records per partition (running with 2 partitions)
        // the app is supposed to emit all 40 update records into the output topic
        //
        // the app first commits after each 10 records per partition (total 20 records), and thus will have 2 * 5 uncommitted writes
        // and store updates (ie, another 5 uncommitted writes to a changelog topic per partition)
        // in the uncommitted batch, sending some data for the new key to validate that upon resuming they will not be shown up in the store
        //
        // the failure gets inject after 20 committed and 30 uncommitted records got received
        // -> the failure only kills one thread
        // after fail over, we should read 40 committed records and the state stores should contain the correct sums
        // per key (even if some records got processed twice)

        // We need more processing time under "with state" situation, so increasing the max.poll.interval.ms
        // to avoid unexpected rebalance during test, which will cause unexpected fail over triggered
        try (final KafkaStreams streams = getKafkaStreams("dummy", true, "appDir", 2, groupProtocol, processingThreadsEnabled)) {
            startApplicationAndWaitUntilRunning(streams);

            final List<KeyValue<Long, Long>> committedDataBeforeFailure = prepareData(0L, 10L, 0L, 1L);
            final List<KeyValue<Long, Long>> uncommittedDataBeforeFailure = prepareData(10L, 15L, 0L, 1L, 2L, 3L);

            final List<KeyValue<Long, Long>> dataBeforeFailure = new ArrayList<>(
                committedDataBeforeFailure.size() + uncommittedDataBeforeFailure.size());
            dataBeforeFailure.addAll(committedDataBeforeFailure);
            dataBeforeFailure.addAll(uncommittedDataBeforeFailure);

            final List<KeyValue<Long, Long>> dataAfterFailure = prepareData(15L, 20L, 0L, 1L);

            writeInputData(committedDataBeforeFailure);

            waitForCondition(
                () -> commitRequested.get() == 2, MAX_WAIT_TIME_MS,
                "StreamsTasks did not request commit.");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C
            // p-1: ---> 10 rec + C

            final List<KeyValue<Long, Long>> committedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, committedDataBeforeFailure.size(), CONSUMER_GROUP_ID);
            checkResultPerKey(
                committedRecords,
                computeExpectedResult(committedDataBeforeFailure),
                "The committed records before failure do not match what expected");

            writeInputData(uncommittedDataBeforeFailure);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec (pending)
            // p-1: ---> 10 rec + C  + 5 rec (pending)

            final List<KeyValue<Long, Long>> uncommittedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, dataBeforeFailure.size(), null);
            final List<KeyValue<Long, Long>> expectedResultBeforeFailure = computeExpectedResult(dataBeforeFailure);


            checkResultPerKey(
                uncommittedRecords,
                expectedResultBeforeFailure,
                "The uncommitted records before failure do not match what expected");
            verifyStateStore(
                streams,
                getMaxPerKey(expectedResultBeforeFailure),
                "The state store content before failure do not match what expected");

            errorInjected.set(true);
            writeInputData(dataAfterFailure);

            waitForCondition(
                () -> uncaughtException != null, MAX_WAIT_TIME_MS,
                "Should receive uncaught exception from one StreamThread.");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C

            final List<KeyValue<Long, Long>> allCommittedRecords = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                committedDataBeforeFailure.size() + uncommittedDataBeforeFailure.size() + dataAfterFailure.size(),
                CONSUMER_GROUP_ID + "_ALL");

            final List<KeyValue<Long, Long>> committedRecordsAfterFailure = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                uncommittedDataBeforeFailure.size() + dataAfterFailure.size(),
                CONSUMER_GROUP_ID);

            final int allCommittedRecordsAfterRecoverySize = committedDataBeforeFailure.size() +
                uncommittedDataBeforeFailure.size() + dataAfterFailure.size();
            final List<KeyValue<Long, Long>> allExpectedCommittedRecordsAfterRecovery = new ArrayList<>(allCommittedRecordsAfterRecoverySize);
            allExpectedCommittedRecordsAfterRecovery.addAll(committedDataBeforeFailure);
            allExpectedCommittedRecordsAfterRecovery.addAll(uncommittedDataBeforeFailure);
            allExpectedCommittedRecordsAfterRecovery.addAll(dataAfterFailure);

            final List<KeyValue<Long, Long>> expectedResult = computeExpectedResult(allExpectedCommittedRecordsAfterRecovery);

            checkResultPerKey(
                allCommittedRecords,
                expectedResult,
                "The all committed records after recovery do not match what expected");

            checkResultPerKey(
                committedRecordsAfterFailure,
                expectedResult.subList(committedDataBeforeFailure.size(), expectedResult.size()),
                "The committed records after recovery do not match what expected");

            verifyStateStore(
                streams,
                getMaxPerKey(expectedResult),
                "The state store content after recovery do not match what expected");

            assertThat("Should only get one uncaught exception from Streams.", hasUnexpectedError, is(false));
        }
    }

    @ParameterizedTest
    @MethodSource("groupProtocolAndProcessingThreadsParameters")
    public void shouldNotViolateEosIfOneTaskGetsFencedUsingIsolatedAppInstances(final String groupProtocol, final boolean processingThreadsEnabled) throws Exception {
        // this test writes 10 + 5 + 5 + 10 records per partition (running with 2 partitions)
        // the app is supposed to copy all 60 records into the output topic
        //
        // the app first commits after each 10 records per partition, and thus will have 2 * 5 uncommitted writes
        //
        // Then, a stall gets injected after 20 committed and 30 uncommitted records got received
        // -> the stall only affects one thread and should trigger a rebalance
        // after rebalancing, we should read 40 committed records (even if 50 record got written)
        //
        // afterward, the "stalling" thread resumes, and another rebalance should get triggered
        // we write the remaining 20 records and verify to read 60 result records
        try (
            final KafkaStreams streams1 = getKafkaStreams("streams1", false, "appDir1", 1, groupProtocol, processingThreadsEnabled);
            final KafkaStreams streams2 = getKafkaStreams("streams2", false, "appDir2", 1, groupProtocol, processingThreadsEnabled)
        ) {
            startApplicationAndWaitUntilRunning(streams1);
            startApplicationAndWaitUntilRunning(streams2);

            final List<KeyValue<Long, Long>> committedDataBeforeStall = prepareData(0L, 10L, 0L, 1L);
            final List<KeyValue<Long, Long>> uncommittedDataBeforeStall = prepareData(10L, 15L, 0L, 1L);

            final List<KeyValue<Long, Long>> dataBeforeStall = new ArrayList<>(
                committedDataBeforeStall.size() + uncommittedDataBeforeStall.size());
            dataBeforeStall.addAll(committedDataBeforeStall);
            dataBeforeStall.addAll(uncommittedDataBeforeStall);

            final List<KeyValue<Long, Long>> dataToTriggerFirstRebalance = prepareData(15L, 20L, 0L, 1L);

            final List<KeyValue<Long, Long>> dataAfterSecondRebalance = prepareData(20L, 30L, 0L, 1L);

            writeInputData(committedDataBeforeStall);

            waitForCondition(
                () -> commitRequested.get() == 2, MAX_WAIT_TIME_MS,
                "StreamsTasks did not request commit.");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C
            // p-1: ---> 10 rec + C

            final List<KeyValue<Long, Long>> committedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, committedDataBeforeStall.size(), CONSUMER_GROUP_ID);
            checkResultPerKey(
                committedRecords,
                committedDataBeforeStall,
                "The committed records before stall do not match what expected");

            writeInputData(uncommittedDataBeforeStall);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec (pending)
            // p-1: ---> 10 rec + C  + 5 rec (pending)

            final List<KeyValue<Long, Long>> uncommittedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, dataBeforeStall.size(), null);
            checkResultPerKey(
                uncommittedRecords,
                dataBeforeStall,
                "The uncommitted records before stall do not match what expected");

            LOG.info("Injecting Stall");
            stallInjected.set(true);
            writeInputData(dataToTriggerFirstRebalance);
            LOG.info("Input Data Written");
            waitForCondition(
                () -> stallingHost.get() != null,
                MAX_WAIT_TIME_MS,
                "Expected a host to start stalling"
            );
            final String observedStallingHost = stallingHost.get();
            final KafkaStreams remainingInstance;
            if ("streams1".equals(observedStallingHost)) {
                remainingInstance = streams2;
            } else if ("streams2".equals(observedStallingHost)) {
                remainingInstance = streams1;
            } else {
                throw new IllegalArgumentException("unexpected host name: " + observedStallingHost);
            }

            // the stalling instance won't have an updated view, and it doesn't matter what it thinks
            // the assignment is. We only really care that the remaining instance only sees one host
            // that owns both partitions.
            waitForCondition(
                () -> remainingInstance.metadataForAllStreamsClients().size() == 1
                    && remainingInstance.metadataForAllStreamsClients().iterator().next().topicPartitions().size() == 2,
                MAX_WAIT_TIME_MS,
                () -> "Should have rebalanced.\n" +
                    "Streams1[" + streams1.metadataForAllStreamsClients() + "]\n" +
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C

            final List<KeyValue<Long, Long>> committedRecordsAfterRebalance = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size(),
                CONSUMER_GROUP_ID);

            final List<KeyValue<Long, Long>> expectedCommittedRecordsAfterRebalance = new ArrayList<>(
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size());
            expectedCommittedRecordsAfterRebalance.addAll(uncommittedDataBeforeStall);
            expectedCommittedRecordsAfterRebalance.addAll(dataToTriggerFirstRebalance);

            checkResultPerKey(
                committedRecordsAfterRebalance,
                expectedCommittedRecordsAfterRebalance,
                "The all committed records after rebalance do not match what expected");

            LOG.info("Releasing Stall");
            doStall = false;
            // Once the stalling host rejoins the group, we expect both instances to see both instances.
            // It doesn't really matter what the assignment is, but we might as well also assert that they
            // both see both partitions assigned exactly once
            waitForCondition(
                () -> streams1.metadataForAllStreamsClients().size() == 2
                    && streams2.metadataForAllStreamsClients().size() == 2
                    && streams1.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2
                    && streams2.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2,
                MAX_WAIT_TIME_MS,
                () -> "Should have rebalanced.\n" +
                    "Streams1[" + streams1.metadataForAllStreamsClients() + "]\n" +
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            writeInputData(dataAfterSecondRebalance);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C

            final List<KeyValue<Long, Long>> allCommittedRecords = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                committedDataBeforeStall.size() + uncommittedDataBeforeStall.size()
                + dataToTriggerFirstRebalance.size() + dataAfterSecondRebalance.size(),
                CONSUMER_GROUP_ID + "_ALL");

            final int allCommittedRecordsAfterRecoverySize = committedDataBeforeStall.size() +
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size() + dataAfterSecondRebalance.size();
            final List<KeyValue<Long, Long>> allExpectedCommittedRecordsAfterRecovery = new ArrayList<>(allCommittedRecordsAfterRecoverySize);
            allExpectedCommittedRecordsAfterRecovery.addAll(committedDataBeforeStall);
            allExpectedCommittedRecordsAfterRecovery.addAll(uncommittedDataBeforeStall);
            allExpectedCommittedRecordsAfterRecovery.addAll(dataToTriggerFirstRebalance);
            allExpectedCommittedRecordsAfterRecovery.addAll(dataAfterSecondRebalance);

            checkResultPerKey(
                allCommittedRecords,
                allExpectedCommittedRecordsAfterRecovery,
                "The all committed records after recovery do not match what expected");
        }
    }

    @ParameterizedTest
    @MethodSource("groupProtocolAndProcessingThreadsParameters")
    public void shouldWriteLatestOffsetsToCheckpointOnShutdown(final String groupProtocol, final boolean processingThreadsEnabled) throws Exception {
        final List<KeyValue<Long, Long>> writtenData = prepareData(0L, 10, 0L, 1L);
        final List<KeyValue<Long, Long>> expectedResult = computeExpectedResult(writtenData);

        try (final KafkaStreams streams = getKafkaStreams("streams", true, "appDir", 1, groupProtocol, processingThreadsEnabled)) {
            writeInputData(writtenData);

            startApplicationAndWaitUntilRunning(streams);

            waitForCondition(
                    () -> commitRequested.get() == 2, MAX_WAIT_TIME_MS,
                    "StreamsTasks did not request commit.");

            final List<KeyValue<Long, Long>> committedRecords = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, writtenData.size(), CONSUMER_GROUP_ID);

            checkResultPerKey(
                    committedRecords,
                    expectedResult,
                    "The committed records do not match what expected");

            verifyStateStore(
                    streams,
                    getMaxPerKey(expectedResult),
                    "The state store content do not match what expected");
        }

        verifyOffsetsAreInCheckpoint(0);
        verifyOffsetsAreInCheckpoint(1);
    }

    @ParameterizedTest
    @MethodSource("groupProtocolAndProcessingThreadsParameters")
    public void shouldCheckpointRestoredOffsetsWhenClosingCleanDuringRestoring(
            final String groupProtocol, final boolean processingThreadsEnabled) throws Exception {

        final Properties streamsConfiguration = new Properties();
        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, applicationId);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
        streamsConfiguration.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        streamsConfiguration.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, 1);
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory(applicationId).getPath());
        streamsConfiguration.put(InternalConfig.PROCESSING_THREADS_ENABLED, processingThreadsEnabled);
        streamsConfiguration.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);
        streamsConfiguration.put(StreamsConfig.restoreConsumerPrefix(ConsumerConfig.MAX_POLL_RECORDS_CONFIG), 100);
        final String stateStoreName = "stateStore";

        purgeLocalStreamsState(streamsConfiguration);

        final int startKey = 1;
        final int endKey = 30001;
        final List<KeyValue<Integer, Integer>> recordBatch1 = IntStream.range(startKey, endKey - 1000).mapToObj(i -> KeyValue.pair(i, 0)).collect(Collectors.toList());
        IntegrationTestUtils.produceKeyValuesSynchronously(MULTI_PARTITION_INPUT_TOPIC,
            recordBatch1,
            TestUtils.producerConfig(CLUSTER.bootstrapServers(),
                IntegerSerializer.class,
                IntegerSerializer.class),
            CLUSTER.time);

        final StoreBuilder<KeyValueStore<Integer, String>> stateStore = Stores.keyValueStoreBuilder(
            Stores.persistentKeyValueStore(stateStoreName),
            Serdes.Integer(),
            Serdes.String()).withCachingEnabled();

        final int partitionToVerify = 0;
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicBoolean throwException = new AtomicBoolean(false);
        final TaskId task00 = new TaskId(0, partitionToVerify);
        final AtomicLong restoredOffsetsForPartition0 = new AtomicLong(0);
        final Topology topology = new Topology();
        topology
            .addSource("source", MULTI_PARTITION_INPUT_TOPIC)
            .addProcessor("processor", () -> new Processor<Integer, String, Integer, String>() {
                KeyValueStore<Integer, String> stateStore;
                ProcessorContext<Integer, String> context;

                @Override
                public void init(final ProcessorContext<Integer, String> context) {
                    Processor.super.init(context);
                    this.context = context;
                    stateStore = context.getStateStore(stateStoreName);
                }

                @Override
                public void process(final Record<Integer, String> record) {
                    context.recordMetadata().ifPresent(recordMetadata -> {
                        if (recordMetadata.partition() == partitionToVerify) {
                            if (throwException.compareAndSet(true, false)) {
                                throw new TaskCorruptedException(Collections.singleton(task00));
                            }
                            stateStore.put(record.key(), record.value());
                        } else {
                            stateStore.put(record.key(), record.value());
                        }
                    });
                }

                @Override
                public void close() {
                    Processor.super.close();
                }
            }, "source")
            .addStateStore(stateStore, "processor");

        final KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfiguration);
        kafkaStreams.setGlobalStateRestoreListener(new StateRestoreListener() {
            @Override
            public void onRestoreStart(final TopicPartition topicPartition,
                                       final String storeName,
                                       final long startingOffset,
                                       final long endingOffset) {}

            @Override
            public void onBatchRestored(final TopicPartition topicPartition,
                                        final String storeName,
                                        final long batchEndOffset,
                                        final long numRestored) {
                if (topicPartition.partition() == 0) {
                    restoredOffsetsForPartition0.set(batchEndOffset);
                    if (batchEndOffset > 100) {
                        latch.countDown();
                    }
                }
            }

            @Override
            public void onRestoreEnd(final TopicPartition topicPartition,
                                     final String storeName,
                                     final long totalRestored) {}
        });
        startApplicationAndWaitUntilRunning(Collections.singletonList(kafkaStreams), Duration.ofSeconds(60));
        ensureCommittedRecordsInTopicPartition(
            applicationId + "-" + stateStoreName + "-changelog",
            partitionToVerify,
            2000
        );
        throwException.set(true);
        final List<KeyValue<Integer, Integer>> recordBatch2 = IntStream.range(endKey - 1000, endKey).mapToObj(i -> KeyValue.pair(i, 0)).collect(Collectors.toList());
        IntegrationTestUtils.produceKeyValuesSynchronously(MULTI_PARTITION_INPUT_TOPIC,
            recordBatch2,
            TestUtils.producerConfig(CLUSTER.bootstrapServers(),
                IntegerSerializer.class,
                IntegerSerializer.class),
            CLUSTER.time);
        latch.await();
        kafkaStreams.close();
        waitForApplicationState(Collections.singletonList(kafkaStreams), KafkaStreams.State.NOT_RUNNING, Duration.ofSeconds(60));

        final File taskDir = Paths.get(
            streamsConfiguration.getProperty(StreamsConfig.STATE_DIR_CONFIG),
            streamsConfiguration.getProperty(StreamsConfig.APPLICATION_ID_CONFIG),
            task00.toString()
        ).toFile();
        assertTrue(taskDir.exists());
    }


    private final AtomicReference<String> transactionalProducerId = new AtomicReference<>();

    private class TestClientSupplier extends DefaultKafkaClientSupplier {
        @Override
        public Producer<byte[], byte[]> getProducer(final Map<String, Object> config) {
            transactionalProducerId.compareAndSet(null, (String) config.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

            return new KafkaProducer<>(config, new ByteArraySerializer(), new ByteArraySerializer());
        }
    }

    static final AtomicReference<TaskId> TASK_WITH_DATA = new AtomicReference<>();
    static final AtomicBoolean DID_REVOKE_IDLE_TASK = new AtomicBoolean(false);

    @ParameterizedTest
    @ValueSource(strings = {"classic", "streams"})
    public void shouldNotCommitActiveTasksWithPendingInputIfRevokedTaskDidNotMakeProgress(final String groupProtocol) throws Exception {
        // Reset static variables to ensure test isolation
        TASK_WITH_DATA.set(null);
        DID_REVOKE_IDLE_TASK.set(false);

        final AtomicBoolean requestCommit = new AtomicBoolean(false);

        final StreamsBuilder builder = new StreamsBuilder();
        builder.<Long, Long>stream(MULTI_PARTITION_INPUT_TOPIC)
            .process(() -> new ContextualProcessor<Long, Long, Long, Long>() {
                @Override
                public void process(final Record<Long, Long> record) {
                    if (!requestCommit.get()) {
                        if (TASK_WITH_DATA.get() != null) {
                            throw new IllegalStateException("Should only process single record using single task");
                        }
                        TASK_WITH_DATA.set(context().taskId());
                    }

                    context().forward(record.withValue(context().recordMetadata().get().offset()));

                    if (requestCommit.get()) {
                        context().commit();
                    }
                }
            })
            .to(SINGLE_PARTITION_OUTPUT_TOPIC);

        final Properties properties = new Properties();
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        properties.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, Integer.MAX_VALUE);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_RECORDS_CONFIG), 1);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.METADATA_MAX_AGE_CONFIG), "1000");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG), "earliest");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), MAX_POLL_INTERVAL_MS - 1);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG), MAX_POLL_INTERVAL_MS);
        properties.put(StreamsConfig.producerPrefix(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG), Integer.MAX_VALUE);
        properties.put(StreamsConfig.TASK_ASSIGNOR_CLASS_CONFIG, TestTaskAssignor.class.getName());
        properties.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);

        final Properties config = StreamsTestUtils.getStreamsConfig(
            applicationId,
            CLUSTER.bootstrapServers(),
            Serdes.LongSerde.class.getName(),
            Serdes.LongSerde.class.getName(),
            properties
        );


        try (final KafkaStreams streams = new KafkaStreams(builder.build(), config, new TestClientSupplier())) {
            startApplicationAndWaitUntilRunning(streams);

            // PHASE 1:
            // write single input record, and wait for it to get into output topic (uncommitted)
            // StreamThread-1 now has a task with progress, and one task w/o progress
            final List<KeyValue<Long, Long>> inputDataTask0 = Collections.singletonList(KeyValue.pair(1L, -1L));

            IntegrationTestUtils.produceKeyValuesSynchronously(
                MULTI_PARTITION_INPUT_TOPIC,
                inputDataTask0,
                TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class),
                CLUSTER.time
            );

            final List<KeyValue<Long, Long>> expectedUncommittedResultTask0 = Collections.singletonList(KeyValue.pair(1L, 0L));
            final List<KeyValue<Long, Long>> uncommittedRecordsBeforeRebalance = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, expectedUncommittedResultTask0.size(), null);
            checkResultPerKey(uncommittedRecordsBeforeRebalance, expectedUncommittedResultTask0, "The uncommitted records do not match what expected");

            // PHASE 2:
            // add second thread, to trigger rebalance
            // expect idle task to get revoked -- this should not trigger a TX commit
            streams.addStreamThread();
            if (groupProtocol.equals("classic")) {
                waitForCondition(DID_REVOKE_IDLE_TASK::get, "Idle Task was not revoked as expected.");
            }
            // best-effort sanity check (might pass and not detect issue in slow environments)
            try {
                readResult(SINGLE_PARTITION_OUTPUT_TOPIC, 1, "consumer", 10_000L);
                throw new Exception("Should not be able to read records, as they should have not been committed.");
            } catch (final AssertionError expected) {
                // swallow -- we expect to not be able to read uncommitted data, but time-out
            }

            // PHASE 3:
            // fence producer to abort pending TX of first input record
            // expect rebalancing and recovery until both input record are processed
            requestCommit.set(true);

            // produce into input topic to fence KS producer
            final List<KeyValue<Long, Long>> inputDataTask0Fencing = Collections.singletonList(KeyValue.pair(4L, -3L));

            final Properties producerConfigs = new Properties();
            producerConfigs.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalProducerId.get());
            IntegrationTestUtils.produceKeyValuesSynchronously(
                MULTI_PARTITION_INPUT_TOPIC,
                inputDataTask0Fencing,
                TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class, producerConfigs),
                CLUSTER.time,
                true
            );

            final List<KeyValue<Long, Long>> expectedUncommittedResultAfterError = Arrays.asList(KeyValue.pair(1L, 0L), KeyValue.pair(1L, 0L), KeyValue.pair(4L, 1L));
            final List<KeyValue<Long, Long>> uncommittedRecordsAfterError = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, expectedUncommittedResultAfterError.size(), null);
            checkResultPerKey(uncommittedRecordsAfterError, expectedUncommittedResultAfterError, "The committed records do not match what expected");
        }

        final List<KeyValue<Long, Long>> expectedFinalResult = Arrays.asList(KeyValue.pair(1L, 0L), KeyValue.pair(4L, 1L));
        final List<KeyValue<Long, Long>> finalResult = readResult(SINGLE_PARTITION_OUTPUT_TOPIC, 2, "committed-only-consumer");
        checkResultPerKey(finalResult, expectedFinalResult, "The committed records do not match what expected");
    }

    private void verifyOffsetsAreInCheckpoint(final int partition) throws IOException {
        final String stateStoreDir = stateTmpDir + File.separator + "appDir" + File.separator + applicationId + File.separator + "0_" + partition + File.separator;

        // Verify that the checkpointed offsets match exactly with max offset of the records in the changelog
        final OffsetCheckpoint checkpoint = new OffsetCheckpoint(new File(stateStoreDir + ".checkpoint"));
        final Map<TopicPartition, Long> checkpointedOffsets = checkpoint.read();
        checkpointedOffsets.forEach(this::verifyChangelogMaxRecordOffsetMatchesCheckpointedOffset);
    }

    private void verifyChangelogMaxRecordOffsetMatchesCheckpointedOffset(final TopicPartition tp, final long checkpointedOffset) {
        try (
            KafkaConsumer<String, String> consumer = new KafkaConsumer<>(
                consumerConfig(
                    CLUSTER.bootstrapServers(),
                    ByteArrayDeserializer.class,
                    ByteArrayDeserializer.class
                )
            )
        ) {
            final List<TopicPartition> partitions = Collections.singletonList(tp);
            consumer.assign(partitions);
            consumer.seekToEnd(partitions);
            final long topicEndOffset = consumer.position(tp);

            assertTrue(topicEndOffset >= checkpointedOffset,
                "changelog topic end " + topicEndOffset + " is less than checkpointed offset " + checkpointedOffset);

            consumer.seekToBeginning(partitions);

            Long maxRecordOffset = null;
            while (consumer.position(tp) != topicEndOffset) {
                final List<ConsumerRecord<String, String>> records = consumer.poll(Duration.ofMillis(0)).records(tp);
                if (!records.isEmpty()) {
                    maxRecordOffset = records.get(records.size() - 1).offset();
                }
            }

            assertEquals(maxRecordOffset, (Long) checkpointedOffset,
                "Checkpointed offset does not match end of changelog");
        }
    }

    private List<KeyValue<Long, Long>> prepareData(final long fromInclusive,
                                                   final long toExclusive,
                                                   final Long... keys) {
        final long dataSize = keys.length * (toExclusive - fromInclusive);
        final List<KeyValue<Long, Long>> data = new ArrayList<>((int) dataSize);

        for (final Long k : keys) {
            for (long v = fromInclusive; v < toExclusive; ++v) {
                data.add(new KeyValue<>(k, v));
            }
        }

        return data;
    }

    // the threads should no longer fail one thread one at a time
    private KafkaStreams getKafkaStreams(final String dummyHostName,
                                         final boolean withState,
                                         final String appDir,
                                         final int numberOfStreamsThreads,
                                         final String groupProtocol,
                                         final boolean processingThreadsEnabled) {
        commitRequested = new AtomicInteger(0);
        errorInjected = new AtomicBoolean(false);
        stallInjected = new AtomicBoolean(false);
        stallingHost = new AtomicReference<>();
        final StreamsBuilder builder = new StreamsBuilder();

        String[] storeNames = new String[0];
        if (withState) {
            storeNames = new String[] {storeName};
            final StoreBuilder<KeyValueStore<Long, Long>> storeBuilder = Stores
                .keyValueStoreBuilder(Stores.persistentKeyValueStore(storeName), Serdes.Long(), Serdes.Long())
                .withCachingEnabled();

            builder.addStateStore(storeBuilder);
        }

        final KStream<Long, Long> input = builder.stream(MULTI_PARTITION_INPUT_TOPIC);
        input.process(() -> new Processor<Long, Long, Long, Long>() {
                ProcessorContext<Long, Long> context;
                KeyValueStore<Long, Long> state = null;

                @Override
                public void init(final ProcessorContext<Long, Long> context) {
                    this.context = context;

                    if (withState) {
                        state = context.getStateStore(storeName);
                    }
                }

                @Override
                public void process(final Record<Long, Long> record) {
                    if (stallInjected.compareAndSet(true, false)) {
                        LOG.info(dummyHostName + " is executing the injected stall");
                        stallingHost.set(dummyHostName);
                        while (doStall) {
                            final Thread thread = Thread.currentThread();
                            if (thread.isInterrupted()) {
                                throw new RuntimeException("Detected we've been interrupted.");
                            }
                            if (!processingThreadsEnabled) {
                                if (!((StreamThread) thread).isRunning()) {
                                    throw new RuntimeException("Detected we've been interrupted.");
                                }
                            }
                            try {
                                Thread.sleep(100);
                            } catch (final InterruptedException e) {
                                throw new RuntimeException(e);
                            }
                        }
                    }

                    final long key = record.key();
                    final long value = record.value();

                    if ((value + 1) % 10 == 0) {
                        context.commit();
                        commitRequested.incrementAndGet();
                    }

                    if (state != null) {
                        Long sum = state.get(key);

                        if (sum == null) {
                            sum = value;
                        } else {
                            sum += value;
                        }
                        state.put(key, sum);
                    }


                    if (errorInjected.compareAndSet(true, false)) {
                        // only tries to fail once on one of the task
                        throw new RuntimeException("Injected test exception.");
                    }

                    if (state != null) {
                        context.forward(record.withValue(state.get(key)));
                    } else {
                        context.forward(record);
                    }
                }
            }, storeNames)
            .to(SINGLE_PARTITION_OUTPUT_TOPIC);

        stateTmpDir = TestUtils.tempDirectory().getPath() + File.separator;

        final Properties properties = new Properties();
        // Set commit interval to a larger value to avoid affection of controlled stream commit,
        // but not too large as we need to have a relatively low transaction timeout such
        // that it should help trigger the timed out transaction in time.
        final long commitIntervalMs = 20_000L;
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        properties.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, numberOfStreamsThreads);
        properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, commitIntervalMs);
        properties.put(StreamsConfig.producerPrefix(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG), (int) commitIntervalMs);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.METADATA_MAX_AGE_CONFIG), "1000");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG), "earliest");
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), MAX_POLL_INTERVAL_MS - 1);
        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG), MAX_POLL_INTERVAL_MS);
        properties.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        properties.put(StreamsConfig.STATE_DIR_CONFIG, stateTmpDir + appDir);
        properties.put(StreamsConfig.APPLICATION_SERVER_CONFIG, dummyHostName + ":2142");
        properties.put(InternalConfig.PROCESSING_THREADS_ENABLED, processingThreadsEnabled);
        properties.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);

        final Properties config = StreamsTestUtils.getStreamsConfig(
            applicationId,
            CLUSTER.bootstrapServers(),
            Serdes.LongSerde.class.getName(),
            Serdes.LongSerde.class.getName(),
            properties);

        final KafkaStreams streams = new KafkaStreams(builder.build(), config);

        streams.setUncaughtExceptionHandler(e -> {
            if (uncaughtException != null ||
                !(e instanceof StreamsException) ||
                !e.getCause().getMessage().equals("Injected test exception.")) {
                e.printStackTrace(System.err);
                hasUnexpectedError = true;
            }
            uncaughtException = e;
            return StreamsUncaughtExceptionHandler.StreamThreadExceptionResponse.REPLACE_THREAD;
        });
        return streams;
    }

    private void writeInputData(final List<KeyValue<Long, Long>> records) {
        IntegrationTestUtils.produceKeyValuesSynchronously(
            MULTI_PARTITION_INPUT_TOPIC,
            records,
            TestUtils.producerConfig(CLUSTER.bootstrapServers(), LongSerializer.class, LongSerializer.class),
            CLUSTER.time
        );
    }

    private List<KeyValue<Long, Long>> readResult(final String topic,
                                                  final int numberOfRecords,
                                                  final String groupId) throws Exception {
        return readResult(topic, numberOfRecords, LongDeserializer.class, LongDeserializer.class, groupId, DEFAULT_TIMEOUT);
    }

    private List<KeyValue<Long, Long>> readResult(final String topic,
                                                  final int numberOfRecords,
                                                  final String groupId,
                                                  final long timeout) throws Exception {
        return readResult(topic, numberOfRecords, LongDeserializer.class, LongDeserializer.class, groupId, timeout);
    }

    private <K, V> List<KeyValue<K, V>> readResult(final String topic,
                                                   final int numberOfRecords,
                                                   final Class<? extends Deserializer<K>> keyDeserializer,
                                                   final Class<? extends Deserializer<V>> valueDeserializer,
                                                   final String groupId,
                                                   final long timeout) throws Exception {
        if (groupId != null) {
            return IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived(
                TestUtils.consumerConfig(
                    CLUSTER.bootstrapServers(),
                    groupId,
                    keyDeserializer,
                    valueDeserializer,
                    Utils.mkProperties(Collections.singletonMap(
                        ConsumerConfig.ISOLATION_LEVEL_CONFIG,
                        IsolationLevel.READ_COMMITTED.toString()))),
                topic,
                numberOfRecords,
                timeout
            );
        }

        // read uncommitted
        return IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived(
            TestUtils.consumerConfig(CLUSTER.bootstrapServers(), keyDeserializer, valueDeserializer),
            topic,
            numberOfRecords,
            timeout
        );
    }

    private <K, V> void ensureCommittedRecordsInTopicPartition(final String topic,
                                                               final int partition,
                                                               final int numberOfRecords) throws Exception {
        final long timeoutMs = 2 * DEFAULT_TIMEOUT;
        final int maxTries = 10;
        final long deadline = System.currentTimeMillis() + timeoutMs;
        int tries = 0;
        while (true) {
            final List<ConsumerRecord<K, V>> consumerRecords = waitUntilMinRecordsReceived(
                TestUtils.consumerConfig(
                    CLUSTER.bootstrapServers(),
                    CONSUMER_GROUP_ID,
                    IntegerDeserializer.class,
                    IntegerDeserializer.class,
                    Utils.mkProperties(Collections.singletonMap(
                        ConsumerConfig.ISOLATION_LEVEL_CONFIG,
                        IsolationLevel.READ_COMMITTED.toString())
                    )
                ),
                topic,
                numberOfRecords,
                timeoutMs
            );
            ++tries;
            if (consumerRecords.stream().anyMatch(record -> record.partition() == partition)) {
                return;
            }
            if (tries >= maxTries) {
                throw new AssertionError("No committed records in topic " + topic
                    + ", partition " + partition + " after " + maxTries + " retries.");
            }
            final long now = System.currentTimeMillis();
            if (now > deadline) {
                throw new AssertionError("No committed records in topic " + topic
                    + ", partition " + partition + " after " + timeoutMs + " ms.");
            }
        }
    }

    private List<KeyValue<Long, Long>> computeExpectedResult(final List<KeyValue<Long, Long>> input) {
        final List<KeyValue<Long, Long>> expectedResult = new ArrayList<>(input.size());

        final HashMap<Long, Long> sums = new HashMap<>();

        for (final KeyValue<Long, Long> record : input) {
            Long sum = sums.get(record.key);
            if (sum == null) {
                sum = record.value;
            } else {
                sum += record.value;
            }
            sums.put(record.key, sum);
            expectedResult.add(new KeyValue<>(record.key, sum));
        }

        return expectedResult;
    }

    private Set<KeyValue<Long, Long>> getMaxPerKey(final List<KeyValue<Long, Long>> input) {
        final Set<KeyValue<Long, Long>> expectedResult = new HashSet<>(input.size());

        final HashMap<Long, Long> maxPerKey = new HashMap<>();

        for (final KeyValue<Long, Long> record : input) {
            final Long max = maxPerKey.get(record.key);
            if (max == null || record.value > max) {
                maxPerKey.put(record.key, record.value);
            }

        }

        for (final Map.Entry<Long, Long> max : maxPerKey.entrySet()) {
            expectedResult.add(new KeyValue<>(max.getKey(), max.getValue()));
        }

        return expectedResult;
    }

    private void verifyStateStore(final KafkaStreams streams,
                                  final Set<KeyValue<Long, Long>> expectedStoreContent,
                                  final String reason) {
        final StateQueryRequest<KeyValueIterator<Long, Long>> request =
                inStore(storeName).withQuery(RangeQuery.withNoBounds());

        final StateQueryResult<KeyValueIterator<Long, Long>> result =
                IntegrationTestUtils.iqv2WaitForResult(streams, request);

        for (final QueryResult<KeyValueIterator<Long, Long>> partitionResult: result.getPartitionResults().values()) {
            try (final KeyValueIterator<Long, Long> it = partitionResult.getResult()) {
                while (it.hasNext()) {
                    assertTrue(expectedStoreContent.remove(it.next()), reason);
                }
            }
        }

        assertTrue(expectedStoreContent.isEmpty(), reason);
    }
}
