KAFKA-9487: Follow-up PR of Kafka-9445

streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java

@@ -1171,11 +1171,9 @@ public <T> T store(final String storeName, final QueryableStoreType<T> queryable
 
     /**
      * Get a facade wrapping the local {@link StateStore} instances with the provided {@link StoreQueryParams}.
-     * StoreQueryParams need required parameters to be set, which are {@code storeName} and if
-     * type is accepted by the provided {@link QueryableStoreType#accepts(StateStore) queryableStoreType}.
      * The returned object can be used to query the {@link StateStore} instances.
      *
-     * @param storeQueryParams   to set the optional parameters to fetch type of stores user wants to fetch when a key is queried
+     * @param storeQueryParams   the parameteres used to fetch a type of queryable store
      * @return A facade wrapping the local {@link StateStore} instances
      * @throws InvalidStateStoreException if Kafka Streams is (re-)initializing or a store with {@code storeName} and
      * {@code queryableStoreType} doesn't exist

streams/src/main/java/org/apache/kafka/streams/StoreQueryParams.java

@@ -21,12 +21,11 @@
 import java.util.Objects;
 
 /**
- * Represents all the query options that a user can provide to state what kind of stores it is expecting.
- * The options would be whether a user would want to enable/disable stale stores
- * or whether it knows the list of partitions that it specifically wants to fetch.
- * If this information is not provided the default behavior is to fetch the stores for all the partitions
+ * Represents all the query options that a user can provide to control the kind of stores we would like to fetch.
+ * The options could indicate - whether we want stale stores to be included (or)
+ * we want to filter for only stores belonging to a partition.
+ * If no specific partition is specified the default behavior is to fetch the stores for all the partitions
  * available on that instance for that particular store name.
- * It contains a partition, which for a point queries can be populated from the {@link KeyQueryMetadata}.
  */
 public class StoreQueryParams<T> {
 
@@ -35,46 +34,42 @@ public class StoreQueryParams<T> {
     private final String storeName;
     private final QueryableStoreType<T> queryableStoreType;
 
-    private StoreQueryParams(final String storeName, final QueryableStoreType<T>  queryableStoreType) {
+    private StoreQueryParams(final String storeName, final QueryableStoreType<T>  queryableStoreType, final Integer partition, final boolean staleStores) {
         this.storeName = storeName;
         this.queryableStoreType = queryableStoreType;
+        this.partition = partition;
+        this.staleStores = staleStores;
     }
 
     public static <T> StoreQueryParams<T> fromNameAndType(final String storeName,
                                                           final QueryableStoreType<T>  queryableStoreType) {
-        return new StoreQueryParams<T>(storeName, queryableStoreType);
+        return new StoreQueryParams<T>(storeName, queryableStoreType, null, false);
     }
 
     /**
      * Set a specific partition that should be queried exclusively.
      *
      * @param partition   The specific integer partition to be fetched from the stores list by using {@link StoreQueryParams}.
      *
-     * @return String storeName
+     * @return StoreQueryParams a new {@code StoreQueryParams} instance configured with the specified partition
      */
     public StoreQueryParams<T> withPartition(final Integer partition) {
-        final StoreQueryParams<T> storeQueryParams = StoreQueryParams.fromNameAndType(this.storeName(), this.queryableStoreType());
-        storeQueryParams.partition = partition;
-        storeQueryParams.staleStores = this.staleStores;
-        return storeQueryParams;
+        return new StoreQueryParams<T>(this.storeName(), this.queryableStoreType(), partition, this.staleStores);
     }
 
     /**
      * Enable querying of stale state stores, i.e., allow to query active tasks during restore as well as standby tasks.
      *
-     * @return String storeName
+     * @return StoreQueryParams a new {@code StoreQueryParams} instance configured with serving from stale stores enabled
      */
     public StoreQueryParams<T> enableStaleStores() {
-        final StoreQueryParams<T> storeQueryParams = StoreQueryParams.fromNameAndType(this.storeName(), this.queryableStoreType());
-        storeQueryParams.partition = this.partition;
-        storeQueryParams.staleStores = true;
-        return storeQueryParams;
+        return new StoreQueryParams<T>(this.storeName(), this.queryableStoreType(), this.partition, true);
     }
 
     /**
-     * Get the store name for which key is queried by the user.
+     * Get the name of the state store that should be queried.
      *
-     * @return String storeName
+     * @return String state store name
      */
     public String storeName() {
         return storeName;
@@ -83,16 +78,16 @@ public String storeName() {
     /**
      * Get the queryable store type for which key is queried by the user.
      *
-     * @return QueryableStoreType queryableStoreType
+     * @return QueryableStoreType type of queryable store
      */
     public QueryableStoreType<T> queryableStoreType() {
         return queryableStoreType;
     }
 
     /**
-     * Get the partition to be used to fetch list of stores.
+     * Get the store partition that will be queried.
      * If the method returns {@code null}, it would mean that no specific partition has been requested,
-     * so all the local partitions for the store will be returned.
+     * so all the local partitions for the store will be queried.
      *
      * @return Integer partition
      */

streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java

@@ -1148,6 +1148,17 @@ public Map<String, List<String>> stateStoreNameToSourceTopics() {
         return results;
     }
 
+    public List<String> sourceTopicsForStore(final String storeName) {
+        final List<String> results = new ArrayList<>();
+        for (final Map.Entry<String, Set<String>> entry : stateStoreNameToSourceTopics.entrySet()) {
+            if (entry.getKey().equals(storeName)) {
+                results.addAll(maybeDecorateInternalSourceTopics(entry.getValue()));
+                return results;
+            }
+        }
+        return Collections.EMPTY_LIST;
+    }
+
     public synchronized Collection<Set<String>> copartitionGroups() {
         // compute transitive closures of copartitionGroups to relieve registering code to know all members
         // of a copartitionGroup at the same time

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java

@@ -109,7 +109,7 @@ public synchronized Collection<StreamsMetadata> getAllMetadataForStore(final Str
             return allMetadata;
         }
 
-        final List<String> sourceTopics = builder.stateStoreNameToSourceTopics().get(storeName);
+        final List<String> sourceTopics = builder.sourceTopicsForStore(storeName);
         if (sourceTopics == null) {
             return Collections.emptyList();
         }
@@ -408,7 +408,7 @@ private <K> StreamsMetadata getStreamsMetadataForKey(final String storeName,
     }
 
     private SourceTopicsInfo getSourceTopicsInfo(final String storeName) {
-        final List<String> sourceTopics = builder.stateStoreNameToSourceTopics().get(storeName);
+        final List<String> sourceTopics = builder.sourceTopicsForStore(storeName);
         if (sourceTopics == null || sourceTopics.isEmpty()) {
             return null;
         }

streams/src/main/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProvider.java

@@ -91,15 +91,15 @@ private TaskId createKeyTaskId(final String storeName, final Integer partition)
         if (partition == null) {
             return null;
         }
-        final List<String> sourceTopics = internalTopologyBuilder.stateStoreNameToSourceTopics().get(storeName);
+        final List<String> sourceTopics = internalTopologyBuilder.sourceTopicsForStore(storeName);
         final Set<String> sourceTopicsSet = sourceTopics.stream().collect(Collectors.toSet());
         final Map<Integer, InternalTopologyBuilder.TopicsInfo> topicGroups = internalTopologyBuilder.topicGroups();
         for (final Map.Entry<Integer, InternalTopologyBuilder.TopicsInfo> topicGroup : topicGroups.entrySet()) {
             if (topicGroup.getValue().sourceTopics.containsAll(sourceTopicsSet)) {
                 return new TaskId(topicGroup.getKey(), partition.intValue());
             }
         }
-        throw new InvalidStateStoreException("Cannot get state store " + storeName + " because the requested partition " + partition + "is" +
+        throw new InvalidStateStoreException("Cannot get state store " + storeName + " because the requested partition " + partition + "is " +
                                                 "not available on this instance");
     }
 }

streams/src/test/java/org/apache/kafka/streams/integration/StoreQueryIntegrationTest.java

@@ -28,7 +28,6 @@
 import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.StreamsConfig;
 import org.apache.kafka.streams.KeyValue;
-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;
@@ -91,7 +90,7 @@ public void after() {
     }
 
     @Test
-    public void shouldQueryAllActivePartitionStoresByDefault() throws Exception {
+    public void shouldQueryOnlyActivePartitionStoresByDefault() throws Exception {
         final int batch1NumMessages = 100;
         final int key = 1;
         final Semaphore semaphore = new Semaphore(0);
@@ -164,28 +163,28 @@ public void shouldQuerySpecificActivePartitionStores() throws Exception {
 
         //key doesn't belongs to this partition
         final int keyDontBelongPartition = (keyPartition == 0) ? 1 : 0;
+        final boolean kafkaStreams1IsActive;
+        if ((keyQueryMetadata.getActiveHost().port() % 2) == 1) {
+            kafkaStreams1IsActive = true;
+        } else {
+            kafkaStreams1IsActive = false;
+        }
+
         final QueryableStoreType<ReadOnlyKeyValueStore<Integer, Integer>> queryableStoreType = QueryableStoreTypes.keyValueStore();
         ReadOnlyKeyValueStore<Integer, Integer> store1 = null;
         ReadOnlyKeyValueStore<Integer, Integer> store2 = null;
-        try {
+        if (kafkaStreams1IsActive) {
             store1 = kafkaStreams1
                     .store(StoreQueryParams.fromNameAndType(TABLE_NAME, queryableStoreType).withPartition(keyPartition));
-        } catch (final InvalidStateStoreException exception) {
-        //Only one among kafkaStreams1 and kafkaStreams2 will contain the specific active store requested. The other will throw exception
-        }
-        try {
+        } else {
             store2 = kafkaStreams2
                     .store(StoreQueryParams.fromNameAndType(TABLE_NAME, queryableStoreType).withPartition(keyPartition));
-        } catch (final InvalidStateStoreException exception) {
-            //Only one among kafkaStreams1 and kafkaStreams2 will contain the specific active store requested. The other will throw exception
         }
-        final boolean kafkaStreams1IsActive;
-        if ((keyQueryMetadata.getActiveHost().port() % 2) == 1) {
-            kafkaStreams1IsActive = true;
+
+        if (kafkaStreams1IsActive) {
             assertThat(store1, is(notNullValue()));
             assertThat(store2, is(nullValue()));
         } else {
-            kafkaStreams1IsActive = false;
             assertThat(store2, is(notNullValue()));
             assertThat(store1, is(nullValue()));
         }
@@ -195,17 +194,12 @@ public void shouldQuerySpecificActivePartitionStores() throws Exception {
 
         ReadOnlyKeyValueStore<Integer, Integer> store3 = null;
         ReadOnlyKeyValueStore<Integer, Integer> store4 = null;
-        try {
+        if (!kafkaStreams1IsActive) {
             store3 = kafkaStreams1
                     .store(StoreQueryParams.fromNameAndType(TABLE_NAME, queryableStoreType).withPartition(keyDontBelongPartition));
-        } catch (final InvalidStateStoreException exception) {
-            //Only one among kafkaStreams1 and kafkaStreams2 will contain the specific active store requested. The other will throw exception
-        }
-        try {
+        } else {
             store4 = kafkaStreams2
                     .store(StoreQueryParams.fromNameAndType(TABLE_NAME, queryableStoreType).withPartition(keyDontBelongPartition));
-        } catch (final InvalidStateStoreException exception) {
-            //Only one among kafkaStreams1 and kafkaStreams2 will contain the specific active store requested. The other will throw exception
         }
 
         // Assert that key is not served when wrong specific partition is requested
@@ -246,8 +240,14 @@ public void shouldQueryAllStalePartitionStores() throws Exception {
                 .store(StoreQueryParams.fromNameAndType(TABLE_NAME, queryableStoreType).enableStaleStores());
 
         // Assert that both active and standby are able to query for a key
-        assertThat(store1.get(key), is(notNullValue()));
-        assertThat(store2.get(key), is(notNullValue()));
+        TestUtils.retryOnExceptionWithTimeout(100, 60 * 1000, () -> {
+            // Assert that after failover we have recovered to the last store write
+            assertThat(store1.get(key), is(equalTo(batch1NumMessages - 1)));
+        });
+        TestUtils.retryOnExceptionWithTimeout(100, 60 * 1000, () -> {
+            // Assert that after failover we have recovered to the last store write
+            assertThat(store2.get(key), is(equalTo(batch1NumMessages - 1)));
+        });
     }
 
     @Test

streams/src/test/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilderTest.java

@@ -257,8 +257,8 @@ public void shouldMapStateStoresToCorrectSourceTopics() {
         mapped.leftJoin(table, MockValueJoiner.TOSTRING_JOINER).groupByKey().count(Materialized.as("count"));
         builder.buildAndOptimizeTopology();
         builder.internalTopologyBuilder.rewriteTopology(new StreamsConfig(StreamsTestUtils.getStreamsConfig(APP_ID)));
-        assertEquals(Collections.singletonList("table-topic"), builder.internalTopologyBuilder.stateStoreNameToSourceTopics().get("table-store"));
-        assertEquals(Collections.singletonList(APP_ID + "-KSTREAM-MAP-0000000003-repartition"), builder.internalTopologyBuilder.stateStoreNameToSourceTopics().get("count"));
+        assertEquals(Collections.singletonList("table-topic"), builder.internalTopologyBuilder.sourceTopicsForStore("table-store"));
+        assertEquals(Collections.singletonList(APP_ID + "-KSTREAM-MAP-0000000003-repartition"), builder.internalTopologyBuilder.sourceTopicsForStore("count"));
     }
 
     @Test