<commit_before>/*
 * 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 kafka.utils.MockTime;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.serialization.LongDeserializer;
import org.apache.kafka.common.serialization.LongSerializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.KafkaStreamsTest;
import org.apache.kafka.streams.StreamsBuilder;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.KeyQueryMetadata;
import org.apache.kafka.streams.KafkaStreams.State;
import org.apache.kafka.streams.LagInfo;
import org.apache.kafka.streams.errors.InvalidStateStoreException;
import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
import org.apache.kafka.streams.kstream.Consumed;
import org.apache.kafka.streams.kstream.KGroupedStream;
import org.apache.kafka.streams.kstream.KStream;
import org.apache.kafka.streams.kstream.KTable;
import org.apache.kafka.streams.kstream.Materialized;
import org.apache.kafka.streams.kstream.Predicate;
import org.apache.kafka.streams.kstream.Produced;
import org.apache.kafka.streams.kstream.TimeWindows;
import org.apache.kafka.streams.kstream.ValueMapper;
import org.apache.kafka.streams.state.QueryableStoreTypes;
import org.apache.kafka.streams.state.ReadOnlyKeyValueStore;
import org.apache.kafka.streams.state.StreamsMetadata;
import org.apache.kafka.streams.state.ReadOnlyWindowStore;
import org.apache.kafka.streams.state.WindowStoreIterator;
import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.KeyValueStore;
import org.apache.kafka.test.IntegrationTest;
import org.apache.kafka.test.MockMapper;
import org.apache.kafka.test.NoRetryException;
import org.apache.kafka.test.TestUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.TestName;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.io.StringReader;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.time.Duration.ofMillis;
import static java.time.Duration.ofSeconds;
import static java.time.Instant.ofEpochMilli;
import static org.apache.kafka.common.utils.Utils.mkSet;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.safeUniqueTestName;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitForApplicationState;
import static org.apache.kafka.test.StreamsTestUtils.startKafkaStreamsAndWaitForRunningState;
import static org.apache.kafka.test.TestUtils.retryOnExceptionWithTimeout;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.core.IsEqual.equalTo;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

@Category({IntegrationTest.class})
public class QueryableStateIntegrationTest {
    private static final Logger log = LoggerFactory.getLogger(QueryableStateIntegrationTest.class);

    private static final long DEFAULT_TIMEOUT_MS = 120 * 1000;

    private static final int NUM_BROKERS = 1;

    @ClassRule
    public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(NUM_BROKERS);
    private static final int STREAM_THREE_PARTITIONS = 4;
    private final MockTime mockTime = CLUSTER.time;
    private String streamOne = "stream-one";
    private String streamTwo = "stream-two";
    private String streamThree = "stream-three";
    private String streamConcurrent = "stream-concurrent";
    private String outputTopic = "output";
    private String outputTopicConcurrent = "output-concurrent";
    private String outputTopicConcurrentWindowed = "output-concurrent-windowed";
    private String outputTopicThree = "output-three";
    // sufficiently large window size such that everything falls into 1 window
    private static final long WINDOW_SIZE = TimeUnit.MILLISECONDS.convert(2, TimeUnit.DAYS);
    private static final int STREAM_TWO_PARTITIONS = 2;
    private static final int NUM_REPLICAS = NUM_BROKERS;
    private Properties streamsConfiguration;
    private List<String> inputValues;
    private Set<String> inputValuesKeys;
    private KafkaStreams kafkaStreams;
    private Comparator<KeyValue<String, String>> stringComparator;
    private Comparator<KeyValue<String, Long>> stringLongComparator;

    private void createTopics() throws Exception {
        final String safeTestName = safeUniqueTestName(getClass(), testName);
        streamOne = streamOne + "-" + safeTestName;
        streamConcurrent = streamConcurrent + "-" + safeTestName;
        streamThree = streamThree + "-" + safeTestName;
        outputTopic = outputTopic + "-" + safeTestName;
        outputTopicConcurrent = outputTopicConcurrent + "-" + safeTestName;
        outputTopicConcurrentWindowed = outputTopicConcurrentWindowed + "-" + safeTestName;
        outputTopicThree = outputTopicThree + "-" + safeTestName;
        streamTwo = streamTwo + "-" + safeTestName;
        CLUSTER.createTopics(streamOne, streamConcurrent);
        CLUSTER.createTopic(streamTwo, STREAM_TWO_PARTITIONS, NUM_REPLICAS);
        CLUSTER.createTopic(streamThree, STREAM_THREE_PARTITIONS, 1);
        CLUSTER.createTopics(outputTopic, outputTopicConcurrent, outputTopicConcurrentWindowed, outputTopicThree);
    }

    /**
     * Try to read inputValues from {@code resources/QueryableStateIntegrationTest/inputValues.txt}, which might be useful
     * for larger scale testing. In case of exception, for instance if no such file can be read, return a small list
     * which satisfies all the prerequisites of the tests.
     */
    private List<String> getInputValues() {
        List<String> input = new ArrayList<>();
        final ClassLoader classLoader = getClass().getClassLoader();
        final String fileName = "QueryableStateIntegrationTest" + File.separator + "inputValues.txt";
        try (final BufferedReader reader = new BufferedReader(
            new FileReader(Objects.requireNonNull(classLoader.getResource(fileName)).getFile()))) {

            for (String line = reader.readLine(); line != null; line = reader.readLine()) {
                input.add(line);
            }
        } catch (final Exception e) {
            log.warn("Unable to read '{}{}{}'. Using default inputValues list", "resources", File.separator, fileName);
            input = Arrays.asList(
                        "hello world",
                        "all streams lead to kafka",
                        "streams",
                        "kafka streams",
                        "the cat in the hat",
                        "green eggs and ham",
                        "that Sam i am",
                        "up the creek without a paddle",
                        "run forest run",
                        "a tank full of gas",
                        "eat sleep rave repeat",
                        "one jolly sailor",
                        "king of the world");

        }
        return input;
    }

    @Rule
    public TestName testName = new TestName();

    @Before
    public void before() throws Exception {
        createTopics();
        streamsConfiguration = new Properties();
        final String safeTestName = safeUniqueTestName(getClass(), testName);

        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, "app-" + safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100);
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");

        stringComparator = Comparator.comparing((KeyValue<String, String> o) -> o.key).thenComparing(o -> o.value);
        stringLongComparator = Comparator.comparing((KeyValue<String, Long> o) -> o.key).thenComparingLong(o -> o.value);
        inputValues = getInputValues();
        inputValuesKeys = new HashSet<>();
        for (final String sentence : inputValues) {
            final String[] words = sentence.split("\\W+");
            Collections.addAll(inputValuesKeys, words);
        }
    }

    @After
    public void shutdown() throws Exception {
        if (kafkaStreams != null) {
            kafkaStreams.close(ofSeconds(30));
        }
        IntegrationTestUtils.purgeLocalStreamsState(streamsConfiguration);
    }

    /**
     * Creates a typical word count topology
     */
    private KafkaStreams createCountStream(final String inputTopic,
                                           final String outputTopic,
                                           final String windowOutputTopic,
                                           final String storeName,
                                           final String windowStoreName,
                                           final Properties streamsConfiguration) {
        final StreamsBuilder builder = new StreamsBuilder();
        final Serde<String> stringSerde = Serdes.String();
        final KStream<String, String> textLines = builder.stream(inputTopic, Consumed.with(stringSerde, stringSerde));

        final KGroupedStream<String, String> groupedByWord = textLines
            .flatMapValues((ValueMapper<String, Iterable<String>>) value -> Arrays.asList(value.split("\\W+")))
            .groupBy(MockMapper.selectValueMapper());

        // Create a State Store for the all time word count
        groupedByWord
            .count(Materialized.as(storeName + "-" + inputTopic))
            .toStream()
            .to(outputTopic, Produced.with(Serdes.String(), Serdes.Long()));

        // Create a Windowed State Store that contains the word count for every 1 minute
        groupedByWord
            .windowedBy(TimeWindows.of(ofMillis(WINDOW_SIZE)))
            .count(Materialized.as(windowStoreName + "-" + inputTopic))
            .toStream((key, value) -> key.key())
            .to(windowOutputTopic, Produced.with(Serdes.String(), Serdes.Long()));

        return new KafkaStreams(builder.build(), streamsConfiguration);
    }


    private void verifyOffsetLagFetch(final List<KafkaStreams> streamsList,
        final Set<String> stores,
        final List<Integer> partitionsPerStreamsInstance) {
        for (int i = 0; i < streamsList.size(); i++) {
            final Map<String, Map<Integer, LagInfo>> localLags = streamsList.get(i).allLocalStorePartitionLags();
            final int expectedPartitions = partitionsPerStreamsInstance.get(i);
            assertThat(localLags.values().stream().mapToInt(Map::size).sum(), equalTo(expectedPartitions));
            if (expectedPartitions > 0) {
                assertThat(localLags.keySet(), equalTo(stores));
            }
        }
    }

    @Deprecated // using KafkaStreams#metadataForKey
    private void verifyAllKVKeys(final List<KafkaStreams> streamsList,
                                 final KafkaStreams streams,
                                 final KafkaStreamsTest.StateListenerStub stateListener,
                                 final Set<String> keys,
                                 final String storeName,
                                 final long timeout,
                                 final boolean pickInstanceByPort) throws Exception {
        retryOnExceptionWithTimeout(timeout, () -> {
            final List<String> noMetadataKeys = new ArrayList<>();
            final List<String> nullStoreKeys = new ArrayList<>();
            final List<String> nullValueKeys = new ArrayList<>();
            final Map<String, Exception> exceptionalKeys = new TreeMap<>();
            final StringSerializer serializer = new StringSerializer();

            for (final String key: keys) {
                try {
                    final KeyQueryMetadata queryMetadata = streams.queryMetadataForKey(storeName, key, serializer);
                    final StreamsMetadata metadata = streams.metadataForKey(storeName, key, serializer);
                    if (queryMetadata == null || queryMetadata.equals(KeyQueryMetadata.NOT_AVAILABLE)) {
                        noMetadataKeys.add(key);
                        continue;
                    }
                    assertThat(metadata.hostInfo(), equalTo(queryMetadata.getActiveHost()));
                    if (!pickInstanceByPort) {
                        assertThat("Should have standbys to query from", !queryMetadata.getStandbyHosts().isEmpty());
                    }

                    final int index = queryMetadata.getActiveHost().port();
                    final KafkaStreams streamsWithKey = pickInstanceByPort ? streamsList.get(index) : streams;
                    final ReadOnlyKeyValueStore<String, Long> store =
                        IntegrationTestUtils.getStore(storeName, streamsWithKey, true, QueryableStoreTypes.keyValueStore());
                    if (store == null) {
                        nullStoreKeys.add(key);
                        continue;
                    }

                    if (store.get(key) == null) {
                        nullValueKeys.add(key);
                    }
                } catch (final InvalidStateStoreException e) {
                    if (stateListener.mapStates.get(KafkaStreams.State.REBALANCING) < 1) {
                        throw new NoRetryException(new AssertionError(
                            String.format("Received %s for key %s and expected at least one rebalancing state, but had none",
                                e.getClass().getName(), key)));
                    }
                } catch (final Exception e) {
                    exceptionalKeys.put(key, e);
                }
            }

            assertNoKVKeyFailures(storeName, timeout, noMetadataKeys, nullStoreKeys, nullValueKeys, exceptionalKeys);
        });
    }

    @Deprecated // using KafkaStreams#metadataForKey
    private void verifyAllWindowedKeys(final List<KafkaStreams> streamsList,
                                       final KafkaStreams streams,
                                       final KafkaStreamsTest.StateListenerStub stateListenerStub,
                                       final Set<String> keys,
                                       final String storeName,
                                       final Long from,
                                       final Long to,
                                       final long timeout,
                                       final boolean pickInstanceByPort) throws Exception {
        retryOnExceptionWithTimeout(timeout, () -> {
            final List<String> noMetadataKeys = new ArrayList<>();
            final List<String> nullStoreKeys = new ArrayList<>();
            final List<String> nullValueKeys = new ArrayList<>();
            final Map<String, Exception> exceptionalKeys = new TreeMap<>();
            final StringSerializer serializer = new StringSerializer();

            for (final String key: keys) {
                try {
                    final KeyQueryMetadata queryMetadata = streams.queryMetadataForKey(storeName, key, serializer);
                    final StreamsMetadata metadata = streams.metadataForKey(storeName, key, serializer);
                    if (queryMetadata == null || queryMetadata.equals(KeyQueryMetadata.NOT_AVAILABLE)) {
                        noMetadataKeys.add(key);
                        continue;
                    }
                    assertThat(metadata.hostInfo(), equalTo(queryMetadata.getActiveHost()));
                    if (pickInstanceByPort) {
                        assertThat(queryMetadata.getStandbyHosts().size(), equalTo(0));
                    } else {
                        assertThat("Should have standbys to query from", !queryMetadata.getStandbyHosts().isEmpty());
                    }

                    final int index = queryMetadata.getActiveHost().port();
                    final KafkaStreams streamsWithKey = pickInstanceByPort ? streamsList.get(index) : streams;
                    final ReadOnlyWindowStore<String, Long> store =
                        IntegrationTestUtils.getStore(storeName, streamsWithKey, true, QueryableStoreTypes.windowStore());
                    if (store == null) {
                        nullStoreKeys.add(key);
                        continue;
                    }

                    if (store.fetch(key, ofEpochMilli(from), ofEpochMilli(to)) == null) {
                        nullValueKeys.add(key);
                    }
                } catch (final InvalidStateStoreException e) {
                    // there must have been at least one rebalance state
                    if (stateListenerStub.mapStates.get(KafkaStreams.State.REBALANCING) < 1) {
                        throw new NoRetryException(new AssertionError(
                            String.format("Received %s for key %s and expected at least one rebalancing state, but had none",
                                e.getClass().getName(), key)));
                    }
                } catch (final Exception e) {
                    exceptionalKeys.put(key, e);
                }
            }

            assertNoKVKeyFailures(storeName, timeout, noMetadataKeys, nullStoreKeys, nullValueKeys, exceptionalKeys);
        });
    }

    private void assertNoKVKeyFailures(final String storeName,
                                       final long timeout,
                                       final List<String> noMetadataKeys,
                                       final List<String> nullStoreKeys,
                                       final List<String> nullValueKeys,
                                       final Map<String, Exception> exceptionalKeys) throws IOException {
        final StringBuilder reason = new StringBuilder();
        reason.append(String.format("Not all keys are available for store %s in %d ms", storeName, timeout));
        if (!noMetadataKeys.isEmpty()) {
            reason.append("\n    * No metadata is available for these keys: ").append(noMetadataKeys);
        }
        if (!nullStoreKeys.isEmpty()) {
            reason.append("\n    * No store is available for these keys: ").append(nullStoreKeys);
        }
        if (!nullValueKeys.isEmpty()) {
            reason.append("\n    * No value is available for these keys: ").append(nullValueKeys);
        }
        if (!exceptionalKeys.isEmpty()) {
            reason.append("\n    * Exceptions were raised for the following keys: ");
            for (final Entry<String, Exception> entry : exceptionalKeys.entrySet()) {
                reason.append(String.format("\n        %s:", entry.getKey()));

                final ByteArrayOutputStream baos = new ByteArrayOutputStream();
                final Exception exception = entry.getValue();

                exception.printStackTrace(new PrintStream(baos));
                try (final BufferedReader reader = new BufferedReader(new StringReader(baos.toString()))) {
                    String line = reader.readLine();
                    while (line != null) {
                        reason.append("\n            ").append(line);
                        line = reader.readLine();
                    }
                }
            }
        }

        assertThat(reason.toString(),
            noMetadataKeys.isEmpty() && nullStoreKeys.isEmpty() && nullValueKeys.isEmpty() && exceptionalKeys.isEmpty());
    }

    @Test
    public void shouldBeAbleToQueryDuringRebalance() throws Exception {
        final int numThreads = STREAM_TWO_PARTITIONS;
        final List<KafkaStreams> streamsList = new ArrayList<>(numThreads);
        final List<KafkaStreamsTest.StateListenerStub> listeners = new ArrayList<>(numThreads);

        final ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        // create stream threads
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";
        for (int i = 0; i < numThreads; i++) {
            final Properties props = (Properties) streamsConfiguration.clone();
            props.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory("shouldBeAbleToQueryDuringRebalance-" + i).getPath());
            props.put(StreamsConfig.APPLICATION_SERVER_CONFIG, "localhost:" + i);
            props.put(StreamsConfig.CLIENT_ID_CONFIG, "instance-" + i);
            final KafkaStreams streams =
                createCountStream(streamThree, outputTopicThree, outputTopicConcurrentWindowed, storeName, windowStoreName, props);
            final KafkaStreamsTest.StateListenerStub listener = new KafkaStreamsTest.StateListenerStub();
            streams.setStateListener(listener);
            listeners.add(listener);
            streamsList.add(streams);
        }
        startApplicationAndWaitUntilRunning(streamsList, Duration.ofSeconds(60));

        final Set<String> stores = mkSet(storeName + "-" + streamThree, windowStoreName + "-" + streamThree);
        verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 4));

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicThree, 1);

            for (int i = 0; i < streamsList.size(); i++) {
                verifyAllKVKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    storeName + "-" + streamThree,
                    DEFAULT_TIMEOUT_MS,
                    true);
                verifyAllWindowedKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    windowStoreName + "-" + streamThree,
                    0L,
                    WINDOW_SIZE,
                    DEFAULT_TIMEOUT_MS,
                    true);
            }
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 4));

            // kill N-1 threads
            for (int i = 1; i < streamsList.size(); i++) {
                final Duration closeTimeout = Duration.ofSeconds(60);
                assertThat(String.format("Streams instance %s did not close in %d ms", i, closeTimeout.toMillis()),
                    streamsList.get(i).close(closeTimeout));
            }

            waitForApplicationState(streamsList.subList(1, numThreads), State.NOT_RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 0));

            // It's not enough to assert that the first instance is RUNNING because it is possible
            // for the above checks to succeed while the instance is in a REBALANCING state.
            waitForApplicationState(streamsList.subList(0, 1), State.RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 0));

            // Even though the closed instance(s) are now in NOT_RUNNING there is no guarantee that
            // the running instance is aware of this, so we must run our follow up queries with
            // enough time for the shutdown to be detected.

            // query from the remaining thread
            verifyAllKVKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                storeName + "-" + streamThree,
                DEFAULT_TIMEOUT_MS,
                true);
            verifyAllWindowedKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                windowStoreName + "-" + streamThree,
                0L,
                WINDOW_SIZE,
                DEFAULT_TIMEOUT_MS,
                true);
            retryOnExceptionWithTimeout(DEFAULT_TIMEOUT_MS, () -> verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0)));
        } finally {
            for (final KafkaStreams streams : streamsList) {
                streams.close();
            }
        }
    }

    @Test
    public void shouldBeAbleQueryStandbyStateDuringRebalance() throws Exception {
        final int numThreads = STREAM_TWO_PARTITIONS;
        final List<KafkaStreams> streamsList = new ArrayList<>(numThreads);
        final List<KafkaStreamsTest.StateListenerStub> listeners = new ArrayList<>(numThreads);

        final ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        // create stream threads
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";
        final Set<String> stores = mkSet(storeName + "-" + streamThree, windowStoreName + "-" + streamThree);
        for (int i = 0; i < numThreads; i++) {
            final Properties props = (Properties) streamsConfiguration.clone();
            props.put(StreamsConfig.APPLICATION_SERVER_CONFIG, "localhost:" + i);
            props.put(StreamsConfig.CLIENT_ID_CONFIG, "instance-" + i);
            props.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
            props.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory("shouldBeAbleQueryStandbyStateDuringRebalance-" + i).getPath());
            final KafkaStreams streams =
                createCountStream(streamThree, outputTopicThree, outputTopicConcurrentWindowed, storeName, windowStoreName, props);
            final KafkaStreamsTest.StateListenerStub listener = new KafkaStreamsTest.StateListenerStub();
            streams.setStateListener(listener);
            listeners.add(listener);
            streamsList.add(streams);
        }
        startApplicationAndWaitUntilRunning(streamsList, ofSeconds(60));
        verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 8));

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicThree, 1);

            // Ensure each thread can serve all keys by itself; i.e standby replication works.
            for (int i = 0; i < streamsList.size(); i++) {
                verifyAllKVKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    storeName + "-" + streamThree,
                    DEFAULT_TIMEOUT_MS,
                    false);
                verifyAllWindowedKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    windowStoreName + "-" + streamThree,
                    0L,
                    WINDOW_SIZE,
                    DEFAULT_TIMEOUT_MS,
                    false);
            }
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 8));

            // kill N-1 threads
            for (int i = 1; i < streamsList.size(); i++) {
                final Duration closeTimeout = Duration.ofSeconds(60);
                assertThat(String.format("Streams instance %s did not close in %d ms", i, closeTimeout.toMillis()),
                    streamsList.get(i).close(closeTimeout));
            }

            waitForApplicationState(streamsList.subList(1, numThreads), State.NOT_RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0));

            // Now, confirm that all the keys are still queryable on the remaining thread, regardless of the state
            verifyAllKVKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                storeName + "-" + streamThree,
                DEFAULT_TIMEOUT_MS,
                false);
            verifyAllWindowedKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                windowStoreName + "-" + streamThree,
                0L,
                WINDOW_SIZE,
                DEFAULT_TIMEOUT_MS,
                false);

            retryOnExceptionWithTimeout(DEFAULT_TIMEOUT_MS, () -> verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0)));
        } finally {
            for (final KafkaStreams streams : streamsList) {
                streams.close();
            }
        }
    }

    @Test
    public void shouldAllowConcurrentAccesses() throws Exception {
        final int numIterations = 500000;
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";

        // send one round of records first to populate the stores
        ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        producerRunnable = new ProducerRunnable(streamConcurrent, inputValues, numIterations - 1);
        final Thread producerThread = new Thread(producerRunnable);
        kafkaStreams = createCountStream(
            streamConcurrent,
            outputTopicConcurrent,
            outputTopicConcurrentWindowed,
            storeName,
            windowStoreName,
            streamsConfiguration);

        startKafkaStreamsAndWaitForRunningState(kafkaStreams);

        producerThread.start();

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicConcurrent, 1);
            
            final ReadOnlyKeyValueStore<String, Long> keyValueStore =
                IntegrationTestUtils.getStore(storeName + "-" + streamConcurrent, kafkaStreams, QueryableStoreTypes.keyValueStore());

            final ReadOnlyWindowStore<String, Long> windowStore =<commit_msg>HOTFIX: set correct numIterations in shouldAllowConcurrentAccesses<commit_after>/*
 * 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 kafka.utils.MockTime;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.serialization.LongDeserializer;
import org.apache.kafka.common.serialization.LongSerializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.KafkaStreamsTest;
import org.apache.kafka.streams.StreamsBuilder;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.KeyQueryMetadata;
import org.apache.kafka.streams.KafkaStreams.State;
import org.apache.kafka.streams.LagInfo;
import org.apache.kafka.streams.errors.InvalidStateStoreException;
import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
import org.apache.kafka.streams.kstream.Consumed;
import org.apache.kafka.streams.kstream.KGroupedStream;
import org.apache.kafka.streams.kstream.KStream;
import org.apache.kafka.streams.kstream.KTable;
import org.apache.kafka.streams.kstream.Materialized;
import org.apache.kafka.streams.kstream.Predicate;
import org.apache.kafka.streams.kstream.Produced;
import org.apache.kafka.streams.kstream.TimeWindows;
import org.apache.kafka.streams.kstream.ValueMapper;
import org.apache.kafka.streams.state.QueryableStoreTypes;
import org.apache.kafka.streams.state.ReadOnlyKeyValueStore;
import org.apache.kafka.streams.state.StreamsMetadata;
import org.apache.kafka.streams.state.ReadOnlyWindowStore;
import org.apache.kafka.streams.state.WindowStoreIterator;
import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.KeyValueStore;
import org.apache.kafka.test.IntegrationTest;
import org.apache.kafka.test.MockMapper;
import org.apache.kafka.test.NoRetryException;
import org.apache.kafka.test.TestUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.TestName;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.io.StringReader;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.time.Duration.ofMillis;
import static java.time.Duration.ofSeconds;
import static java.time.Instant.ofEpochMilli;
import static org.apache.kafka.common.utils.Utils.mkSet;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.safeUniqueTestName;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitForApplicationState;
import static org.apache.kafka.test.StreamsTestUtils.startKafkaStreamsAndWaitForRunningState;
import static org.apache.kafka.test.TestUtils.retryOnExceptionWithTimeout;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.core.IsEqual.equalTo;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

@Category({IntegrationTest.class})
public class QueryableStateIntegrationTest {
    private static final Logger log = LoggerFactory.getLogger(QueryableStateIntegrationTest.class);

    private static final long DEFAULT_TIMEOUT_MS = 120 * 1000;

    private static final int NUM_BROKERS = 1;

    @ClassRule
    public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(NUM_BROKERS);
    private static final int STREAM_THREE_PARTITIONS = 4;
    private final MockTime mockTime = CLUSTER.time;
    private String streamOne = "stream-one";
    private String streamTwo = "stream-two";
    private String streamThree = "stream-three";
    private String streamConcurrent = "stream-concurrent";
    private String outputTopic = "output";
    private String outputTopicConcurrent = "output-concurrent";
    private String outputTopicConcurrentWindowed = "output-concurrent-windowed";
    private String outputTopicThree = "output-three";
    // sufficiently large window size such that everything falls into 1 window
    private static final long WINDOW_SIZE = TimeUnit.MILLISECONDS.convert(2, TimeUnit.DAYS);
    private static final int STREAM_TWO_PARTITIONS = 2;
    private static final int NUM_REPLICAS = NUM_BROKERS;
    private Properties streamsConfiguration;
    private List<String> inputValues;
    private Set<String> inputValuesKeys;
    private KafkaStreams kafkaStreams;
    private Comparator<KeyValue<String, String>> stringComparator;
    private Comparator<KeyValue<String, Long>> stringLongComparator;

    private void createTopics() throws Exception {
        final String safeTestName = safeUniqueTestName(getClass(), testName);
        streamOne = streamOne + "-" + safeTestName;
        streamConcurrent = streamConcurrent + "-" + safeTestName;
        streamThree = streamThree + "-" + safeTestName;
        outputTopic = outputTopic + "-" + safeTestName;
        outputTopicConcurrent = outputTopicConcurrent + "-" + safeTestName;
        outputTopicConcurrentWindowed = outputTopicConcurrentWindowed + "-" + safeTestName;
        outputTopicThree = outputTopicThree + "-" + safeTestName;
        streamTwo = streamTwo + "-" + safeTestName;
        CLUSTER.createTopics(streamOne, streamConcurrent);
        CLUSTER.createTopic(streamTwo, STREAM_TWO_PARTITIONS, NUM_REPLICAS);
        CLUSTER.createTopic(streamThree, STREAM_THREE_PARTITIONS, 1);
        CLUSTER.createTopics(outputTopic, outputTopicConcurrent, outputTopicConcurrentWindowed, outputTopicThree);
    }

    /**
     * Try to read inputValues from {@code resources/QueryableStateIntegrationTest/inputValues.txt}, which might be useful
     * for larger scale testing. In case of exception, for instance if no such file can be read, return a small list
     * which satisfies all the prerequisites of the tests.
     */
    private List<String> getInputValues() {
        List<String> input = new ArrayList<>();
        final ClassLoader classLoader = getClass().getClassLoader();
        final String fileName = "QueryableStateIntegrationTest" + File.separator + "inputValues.txt";
        try (final BufferedReader reader = new BufferedReader(
            new FileReader(Objects.requireNonNull(classLoader.getResource(fileName)).getFile()))) {

            for (String line = reader.readLine(); line != null; line = reader.readLine()) {
                input.add(line);
            }
        } catch (final Exception e) {
            log.warn("Unable to read '{}{}{}'. Using default inputValues list", "resources", File.separator, fileName);
            input = Arrays.asList(
                        "hello world",
                        "all streams lead to kafka",
                        "streams",
                        "kafka streams",
                        "the cat in the hat",
                        "green eggs and ham",
                        "that Sam i am",
                        "up the creek without a paddle",
                        "run forest run",
                        "a tank full of gas",
                        "eat sleep rave repeat",
                        "one jolly sailor",
                        "king of the world");

        }
        return input;
    }

    @Rule
    public TestName testName = new TestName();

    @Before
    public void before() throws Exception {
        createTopics();
        streamsConfiguration = new Properties();
        final String safeTestName = safeUniqueTestName(getClass(), testName);

        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, "app-" + safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100);
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");

        stringComparator = Comparator.comparing((KeyValue<String, String> o) -> o.key).thenComparing(o -> o.value);
        stringLongComparator = Comparator.comparing((KeyValue<String, Long> o) -> o.key).thenComparingLong(o -> o.value);
        inputValues = getInputValues();
        inputValuesKeys = new HashSet<>();
        for (final String sentence : inputValues) {
            final String[] words = sentence.split("\\W+");
            Collections.addAll(inputValuesKeys, words);
        }
    }

    @After
    public void shutdown() throws Exception {
        if (kafkaStreams != null) {
            kafkaStreams.close(ofSeconds(30));
        }
        IntegrationTestUtils.purgeLocalStreamsState(streamsConfiguration);
    }

    /**
     * Creates a typical word count topology
     */
    private KafkaStreams createCountStream(final String inputTopic,
                                           final String outputTopic,
                                           final String windowOutputTopic,
                                           final String storeName,
                                           final String windowStoreName,
                                           final Properties streamsConfiguration) {
        final StreamsBuilder builder = new StreamsBuilder();
        final Serde<String> stringSerde = Serdes.String();
        final KStream<String, String> textLines = builder.stream(inputTopic, Consumed.with(stringSerde, stringSerde));

        final KGroupedStream<String, String> groupedByWord = textLines
            .flatMapValues((ValueMapper<String, Iterable<String>>) value -> Arrays.asList(value.split("\\W+")))
            .groupBy(MockMapper.selectValueMapper());

        // Create a State Store for the all time word count
        groupedByWord
            .count(Materialized.as(storeName + "-" + inputTopic))
            .toStream()
            .to(outputTopic, Produced.with(Serdes.String(), Serdes.Long()));

        // Create a Windowed State Store that contains the word count for every 1 minute
        groupedByWord
            .windowedBy(TimeWindows.of(ofMillis(WINDOW_SIZE)))
            .count(Materialized.as(windowStoreName + "-" + inputTopic))
            .toStream((key, value) -> key.key())
            .to(windowOutputTopic, Produced.with(Serdes.String(), Serdes.Long()));

        return new KafkaStreams(builder.build(), streamsConfiguration);
    }


    private void verifyOffsetLagFetch(final List<KafkaStreams> streamsList,
        final Set<String> stores,
        final List<Integer> partitionsPerStreamsInstance) {
        for (int i = 0; i < streamsList.size(); i++) {
            final Map<String, Map<Integer, LagInfo>> localLags = streamsList.get(i).allLocalStorePartitionLags();
            final int expectedPartitions = partitionsPerStreamsInstance.get(i);
            assertThat(localLags.values().stream().mapToInt(Map::size).sum(), equalTo(expectedPartitions));
            if (expectedPartitions > 0) {
                assertThat(localLags.keySet(), equalTo(stores));
            }
        }
    }

    @Deprecated // using KafkaStreams#metadataForKey
    private void verifyAllKVKeys(final List<KafkaStreams> streamsList,
                                 final KafkaStreams streams,
                                 final KafkaStreamsTest.StateListenerStub stateListener,
                                 final Set<String> keys,
                                 final String storeName,
                                 final long timeout,
                                 final boolean pickInstanceByPort) throws Exception {
        retryOnExceptionWithTimeout(timeout, () -> {
            final List<String> noMetadataKeys = new ArrayList<>();
            final List<String> nullStoreKeys = new ArrayList<>();
            final List<String> nullValueKeys = new ArrayList<>();
            final Map<String, Exception> exceptionalKeys = new TreeMap<>();
            final StringSerializer serializer = new StringSerializer();

            for (final String key: keys) {
                try {
                    final KeyQueryMetadata queryMetadata = streams.queryMetadataForKey(storeName, key, serializer);
                    final StreamsMetadata metadata = streams.metadataForKey(storeName, key, serializer);
                    if (queryMetadata == null || queryMetadata.equals(KeyQueryMetadata.NOT_AVAILABLE)) {
                        noMetadataKeys.add(key);
                        continue;
                    }
                    assertThat(metadata.hostInfo(), equalTo(queryMetadata.getActiveHost()));
                    if (!pickInstanceByPort) {
                        assertThat("Should have standbys to query from", !queryMetadata.getStandbyHosts().isEmpty());
                    }

                    final int index = queryMetadata.getActiveHost().port();
                    final KafkaStreams streamsWithKey = pickInstanceByPort ? streamsList.get(index) : streams;
                    final ReadOnlyKeyValueStore<String, Long> store =
                        IntegrationTestUtils.getStore(storeName, streamsWithKey, true, QueryableStoreTypes.keyValueStore());
                    if (store == null) {
                        nullStoreKeys.add(key);
                        continue;
                    }

                    if (store.get(key) == null) {
                        nullValueKeys.add(key);
                    }
                } catch (final InvalidStateStoreException e) {
                    if (stateListener.mapStates.get(KafkaStreams.State.REBALANCING) < 1) {
                        throw new NoRetryException(new AssertionError(
                            String.format("Received %s for key %s and expected at least one rebalancing state, but had none",
                                e.getClass().getName(), key)));
                    }
                } catch (final Exception e) {
                    exceptionalKeys.put(key, e);
                }
            }

            assertNoKVKeyFailures(storeName, timeout, noMetadataKeys, nullStoreKeys, nullValueKeys, exceptionalKeys);
        });
    }

    @Deprecated // using KafkaStreams#metadataForKey
    private void verifyAllWindowedKeys(final List<KafkaStreams> streamsList,
                                       final KafkaStreams streams,
                                       final KafkaStreamsTest.StateListenerStub stateListenerStub,
                                       final Set<String> keys,
                                       final String storeName,
                                       final Long from,
                                       final Long to,
                                       final long timeout,
                                       final boolean pickInstanceByPort) throws Exception {
        retryOnExceptionWithTimeout(timeout, () -> {
            final List<String> noMetadataKeys = new ArrayList<>();
            final List<String> nullStoreKeys = new ArrayList<>();
            final List<String> nullValueKeys = new ArrayList<>();
            final Map<String, Exception> exceptionalKeys = new TreeMap<>();
            final StringSerializer serializer = new StringSerializer();

            for (final String key: keys) {
                try {
                    final KeyQueryMetadata queryMetadata = streams.queryMetadataForKey(storeName, key, serializer);
                    final StreamsMetadata metadata = streams.metadataForKey(storeName, key, serializer);
                    if (queryMetadata == null || queryMetadata.equals(KeyQueryMetadata.NOT_AVAILABLE)) {
                        noMetadataKeys.add(key);
                        continue;
                    }
                    assertThat(metadata.hostInfo(), equalTo(queryMetadata.getActiveHost()));
                    if (pickInstanceByPort) {
                        assertThat(queryMetadata.getStandbyHosts().size(), equalTo(0));
                    } else {
                        assertThat("Should have standbys to query from", !queryMetadata.getStandbyHosts().isEmpty());
                    }

                    final int index = queryMetadata.getActiveHost().port();
                    final KafkaStreams streamsWithKey = pickInstanceByPort ? streamsList.get(index) : streams;
                    final ReadOnlyWindowStore<String, Long> store =
                        IntegrationTestUtils.getStore(storeName, streamsWithKey, true, QueryableStoreTypes.windowStore());
                    if (store == null) {
                        nullStoreKeys.add(key);
                        continue;
                    }

                    if (store.fetch(key, ofEpochMilli(from), ofEpochMilli(to)) == null) {
                        nullValueKeys.add(key);
                    }
                } catch (final InvalidStateStoreException e) {
                    // there must have been at least one rebalance state
                    if (stateListenerStub.mapStates.get(KafkaStreams.State.REBALANCING) < 1) {
                        throw new NoRetryException(new AssertionError(
                            String.format("Received %s for key %s and expected at least one rebalancing state, but had none",
                                e.getClass().getName(), key)));
                    }
                } catch (final Exception e) {
                    exceptionalKeys.put(key, e);
                }
            }

            assertNoKVKeyFailures(storeName, timeout, noMetadataKeys, nullStoreKeys, nullValueKeys, exceptionalKeys);
        });
    }

    private void assertNoKVKeyFailures(final String storeName,
                                       final long timeout,
                                       final List<String> noMetadataKeys,
                                       final List<String> nullStoreKeys,
                                       final List<String> nullValueKeys,
                                       final Map<String, Exception> exceptionalKeys) throws IOException {
        final StringBuilder reason = new StringBuilder();
        reason.append(String.format("Not all keys are available for store %s in %d ms", storeName, timeout));
        if (!noMetadataKeys.isEmpty()) {
            reason.append("\n    * No metadata is available for these keys: ").append(noMetadataKeys);
        }
        if (!nullStoreKeys.isEmpty()) {
            reason.append("\n    * No store is available for these keys: ").append(nullStoreKeys);
        }
        if (!nullValueKeys.isEmpty()) {
            reason.append("\n    * No value is available for these keys: ").append(nullValueKeys);
        }
        if (!exceptionalKeys.isEmpty()) {
            reason.append("\n    * Exceptions were raised for the following keys: ");
            for (final Entry<String, Exception> entry : exceptionalKeys.entrySet()) {
                reason.append(String.format("\n        %s:", entry.getKey()));

                final ByteArrayOutputStream baos = new ByteArrayOutputStream();
                final Exception exception = entry.getValue();

                exception.printStackTrace(new PrintStream(baos));
                try (final BufferedReader reader = new BufferedReader(new StringReader(baos.toString()))) {
                    String line = reader.readLine();
                    while (line != null) {
                        reason.append("\n            ").append(line);
                        line = reader.readLine();
                    }
                }
            }
        }

        assertThat(reason.toString(),
            noMetadataKeys.isEmpty() && nullStoreKeys.isEmpty() && nullValueKeys.isEmpty() && exceptionalKeys.isEmpty());
    }

    @Test
    public void shouldBeAbleToQueryDuringRebalance() throws Exception {
        final int numThreads = STREAM_TWO_PARTITIONS;
        final List<KafkaStreams> streamsList = new ArrayList<>(numThreads);
        final List<KafkaStreamsTest.StateListenerStub> listeners = new ArrayList<>(numThreads);

        final ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        // create stream threads
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";
        for (int i = 0; i < numThreads; i++) {
            final Properties props = (Properties) streamsConfiguration.clone();
            props.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory("shouldBeAbleToQueryDuringRebalance-" + i).getPath());
            props.put(StreamsConfig.APPLICATION_SERVER_CONFIG, "localhost:" + i);
            props.put(StreamsConfig.CLIENT_ID_CONFIG, "instance-" + i);
            final KafkaStreams streams =
                createCountStream(streamThree, outputTopicThree, outputTopicConcurrentWindowed, storeName, windowStoreName, props);
            final KafkaStreamsTest.StateListenerStub listener = new KafkaStreamsTest.StateListenerStub();
            streams.setStateListener(listener);
            listeners.add(listener);
            streamsList.add(streams);
        }
        startApplicationAndWaitUntilRunning(streamsList, Duration.ofSeconds(60));

        final Set<String> stores = mkSet(storeName + "-" + streamThree, windowStoreName + "-" + streamThree);
        verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 4));

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicThree, 1);

            for (int i = 0; i < streamsList.size(); i++) {
                verifyAllKVKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    storeName + "-" + streamThree,
                    DEFAULT_TIMEOUT_MS,
                    true);
                verifyAllWindowedKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    windowStoreName + "-" + streamThree,
                    0L,
                    WINDOW_SIZE,
                    DEFAULT_TIMEOUT_MS,
                    true);
            }
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 4));

            // kill N-1 threads
            for (int i = 1; i < streamsList.size(); i++) {
                final Duration closeTimeout = Duration.ofSeconds(60);
                assertThat(String.format("Streams instance %s did not close in %d ms", i, closeTimeout.toMillis()),
                    streamsList.get(i).close(closeTimeout));
            }

            waitForApplicationState(streamsList.subList(1, numThreads), State.NOT_RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 0));

            // It's not enough to assert that the first instance is RUNNING because it is possible
            // for the above checks to succeed while the instance is in a REBALANCING state.
            waitForApplicationState(streamsList.subList(0, 1), State.RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(4, 0));

            // Even though the closed instance(s) are now in NOT_RUNNING there is no guarantee that
            // the running instance is aware of this, so we must run our follow up queries with
            // enough time for the shutdown to be detected.

            // query from the remaining thread
            verifyAllKVKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                storeName + "-" + streamThree,
                DEFAULT_TIMEOUT_MS,
                true);
            verifyAllWindowedKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                windowStoreName + "-" + streamThree,
                0L,
                WINDOW_SIZE,
                DEFAULT_TIMEOUT_MS,
                true);
            retryOnExceptionWithTimeout(DEFAULT_TIMEOUT_MS, () -> verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0)));
        } finally {
            for (final KafkaStreams streams : streamsList) {
                streams.close();
            }
        }
    }

    @Test
    public void shouldBeAbleQueryStandbyStateDuringRebalance() throws Exception {
        final int numThreads = STREAM_TWO_PARTITIONS;
        final List<KafkaStreams> streamsList = new ArrayList<>(numThreads);
        final List<KafkaStreamsTest.StateListenerStub> listeners = new ArrayList<>(numThreads);

        final ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        // create stream threads
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";
        final Set<String> stores = mkSet(storeName + "-" + streamThree, windowStoreName + "-" + streamThree);
        for (int i = 0; i < numThreads; i++) {
            final Properties props = (Properties) streamsConfiguration.clone();
            props.put(StreamsConfig.APPLICATION_SERVER_CONFIG, "localhost:" + i);
            props.put(StreamsConfig.CLIENT_ID_CONFIG, "instance-" + i);
            props.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
            props.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory("shouldBeAbleQueryStandbyStateDuringRebalance-" + i).getPath());
            final KafkaStreams streams =
                createCountStream(streamThree, outputTopicThree, outputTopicConcurrentWindowed, storeName, windowStoreName, props);
            final KafkaStreamsTest.StateListenerStub listener = new KafkaStreamsTest.StateListenerStub();
            streams.setStateListener(listener);
            listeners.add(listener);
            streamsList.add(streams);
        }
        startApplicationAndWaitUntilRunning(streamsList, ofSeconds(60));
        verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 8));

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicThree, 1);

            // Ensure each thread can serve all keys by itself; i.e standby replication works.
            for (int i = 0; i < streamsList.size(); i++) {
                verifyAllKVKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    storeName + "-" + streamThree,
                    DEFAULT_TIMEOUT_MS,
                    false);
                verifyAllWindowedKeys(
                    streamsList,
                    streamsList.get(i),
                    listeners.get(i),
                    inputValuesKeys,
                    windowStoreName + "-" + streamThree,
                    0L,
                    WINDOW_SIZE,
                    DEFAULT_TIMEOUT_MS,
                    false);
            }
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 8));

            // kill N-1 threads
            for (int i = 1; i < streamsList.size(); i++) {
                final Duration closeTimeout = Duration.ofSeconds(60);
                assertThat(String.format("Streams instance %s did not close in %d ms", i, closeTimeout.toMillis()),
                    streamsList.get(i).close(closeTimeout));
            }

            waitForApplicationState(streamsList.subList(1, numThreads), State.NOT_RUNNING, Duration.ofSeconds(60));
            verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0));

            // Now, confirm that all the keys are still queryable on the remaining thread, regardless of the state
            verifyAllKVKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                storeName + "-" + streamThree,
                DEFAULT_TIMEOUT_MS,
                false);
            verifyAllWindowedKeys(
                streamsList,
                streamsList.get(0),
                listeners.get(0),
                inputValuesKeys,
                windowStoreName + "-" + streamThree,
                0L,
                WINDOW_SIZE,
                DEFAULT_TIMEOUT_MS,
                false);

            retryOnExceptionWithTimeout(DEFAULT_TIMEOUT_MS, () -> verifyOffsetLagFetch(streamsList, stores, Arrays.asList(8, 0)));
        } finally {
            for (final KafkaStreams streams : streamsList) {
                streams.close();
            }
        }
    }

    @Test
    public void shouldAllowConcurrentAccesses() throws Exception {
        final int numIterations = 500000;
        final String storeName = "word-count-store";
        final String windowStoreName = "windowed-word-count-store";

        // send one round of records first to populate the stores
        ProducerRunnable producerRunnable = new ProducerRunnable(streamThree, inputValues, 1);
        producerRunnable.run();

        producerRunnable = new ProducerRunnable(streamConcurrent, inputValues, numIterations);
        final Thread producerThread = new Thread(producerRunnable);
        kafkaStreams = createCountStream(
            streamConcurrent,
            outputTopicConcurrent,
            outputTopicConcurrentWindowed,
            storeName,
            windowStoreName,
            streamsConfiguration);

        startKafkaStreamsAndWaitForRunningState(kafkaStreams);

        producerThread.start();

        try {
            waitUntilAtLeastNumRecordProcessed(outputTopicConcurrent, 1);
            
            final ReadOnlyKeyValueStore<String, Long> keyValueStore =
                IntegrationTestUtils.getStore(storeName + "-" + streamConcurrent, kafkaStreams, QueryableStoreTypes.keyValueStore());

            final ReadOnlyWindowStore<String, Long> windowStore =
{'fraction_non_alphanumeric': 0.08829603863534015, 'fraction_numerical': 0.0030902060654999456, 'mean_word_length': 2.2671097356805845, 'pattern_counts': {'":': 0, '<': 143, '<?xml version=': 0, '>': 147, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 58, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6546038', 'n_tokens_mistral': 16523, 'n_tokens_neox': 15991, 'n_words': 3702}