KAFKA-9445: Allow adding changes to allow serving from a specific partition

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

@@ -764,7 +764,7 @@ private KafkaStreams(final InternalTopologyBuilder internalTopologyBuilder,
                 delegatingStateRestoreListener,
                 i + 1);
             threadState.put(threads[i].getId(), threads[i].state());
-            storeProviders.add(new StreamThreadStateStoreProvider(threads[i]));
+            storeProviders.add(new StreamThreadStateStoreProvider(threads[i], internalTopologyBuilder));
         }
 
         final StreamStateListener streamStateListener = new StreamStateListener(threadState, globalThreadState);
@@ -1160,47 +1160,29 @@ public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
         return streamsMetadataState.getKeyQueryMetadataForKey(storeName, key, partitioner);
     }
 
+
     /**
-     * Get a facade wrapping the local {@link StateStore} instances with the provided {@code storeName} if the Store's
-     * type is accepted by the provided {@link QueryableStoreType#accepts(StateStore) queryableStoreType}.
-     * The returned object can be used to query the {@link StateStore} instances.
-     *
-     * Only permits queries on active replicas of the store (no standbys or restoring replicas).
-     * See {@link KafkaStreams#store(java.lang.String, org.apache.kafka.streams.state.QueryableStoreType, boolean)}
-     * for the option to set {@code includeStaleStores} to true and trade off consistency in favor of availability.
-     *
-     * @param storeName           name of the store to find
-     * @param queryableStoreType  accept only stores that are accepted by {@link QueryableStoreType#accepts(StateStore)}
-     * @param <T>                 return type
-     * @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
+     * @deprecated since 2.5 release; use {@link #store(StoreQueryParams)}  instead
      */
+    @Deprecated
     public <T> T store(final String storeName, final QueryableStoreType<T> queryableStoreType) {
-        return store(storeName, queryableStoreType, false);
+        return store(StoreQueryParams.fromNameAndType(storeName, queryableStoreType));
     }
 
     /**
-     * Get a facade wrapping the local {@link StateStore} instances with the provided {@code storeName} if the Store's
+     * 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 storeName           name of the store to find
-     * @param queryableStoreType  accept only stores that are accepted by {@link QueryableStoreType#accepts(StateStore)}
-     * @param includeStaleStores      If false, only permit queries on the active replica for a partition, and only if the
-     *                            task for that partition is running. I.e., the state store is not a standby replica,
-     *                            and it is not restoring from the changelog.
-     *                            If true, allow queries on standbys and restoring replicas in addition to active ones.
-     * @param <T>                 return type
+     * @param storeQueryParams   to set the optional parameters to fetch type of stores user wants to fetch when a key is queried
      * @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
      */
-    public <T> T store(final String storeName,
-                       final QueryableStoreType<T> queryableStoreType,
-                       final boolean includeStaleStores) {
+    public <T> T store(final StoreQueryParams<T> storeQueryParams) {
         validateIsRunningOrRebalancing();
-        return queryableStoreProvider.getStore(storeName, queryableStoreType, includeStaleStores);
+        return queryableStoreProvider.getStore(storeQueryParams);
     }
 
     /**

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

@@ -0,0 +1,138 @@
+/*
+ * 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;
+
+import org.apache.kafka.streams.state.QueryableStoreType;
+
+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
+ * 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> {
+
+    private Integer partition;
+    private boolean staleStores;
+    private final String storeName;
+    private final QueryableStoreType<T> queryableStoreType;
+
+    private StoreQueryParams(final String storeName, final QueryableStoreType<T>  queryableStoreType) {
+        this.storeName = storeName;
+        this.queryableStoreType = queryableStoreType;
+    }
+
+    public static <T> StoreQueryParams<T> fromNameAndType(final String storeName,
+                                                          final QueryableStoreType<T>  queryableStoreType) {
+        return new StoreQueryParams<T>(storeName, queryableStoreType);
+    }
+
+    /**
+     * 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
+     */
+    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;
+    }
+
+    /**
+     * Enable querying of stale state stores, i.e., allow to query active tasks during restore as well as standby tasks.
+     *
+     * @return String storeName
+     */
+    public StoreQueryParams<T> enableStaleStores() {
+        final StoreQueryParams<T> storeQueryParams = StoreQueryParams.fromNameAndType(this.storeName(), this.queryableStoreType());
+        storeQueryParams.partition = this.partition;
+        storeQueryParams.staleStores = true;
+        return storeQueryParams;
+    }
+
+    /**
+     * Get the store name for which key is queried by the user.
+     *
+     * @return String storeName
+     */
+    public String storeName() {
+        return storeName;
+    }
+
+    /**
+     * Get the queryable store type for which key is queried by the user.
+     *
+     * @return QueryableStoreType queryableStoreType
+     */
+    public QueryableStoreType<T> queryableStoreType() {
+        return queryableStoreType;
+    }
+
+    /**
+     * Get the partition to be used to fetch list of stores.
+     * 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.
+     *
+     * @return Integer partition
+     */
+    public Integer partition() {
+        return partition;
+    }
+
+    /**
+     * Get the flag staleStores. If {@code true}, include standbys and recovering stores along with running stores.
+     *
+     * @return boolean staleStores
+     */
+    public boolean staleStoresEnabled() {
+        return staleStores;
+    }
+
+    @Override
+    public boolean equals(final Object obj) {
+        if (!(obj instanceof StoreQueryParams)) {
+            return false;
+        }
+        final StoreQueryParams storeQueryParams = (StoreQueryParams) obj;
+        return Objects.equals(storeQueryParams.partition, partition)
+                && Objects.equals(storeQueryParams.staleStores, staleStores)
+                && Objects.equals(storeQueryParams.storeName, storeName)
+                && Objects.equals(storeQueryParams.queryableStoreType, queryableStoreType);
+    }
+
+    @Override
+    public String toString() {
+        return "StoreQueryParams {" +
+                "partition=" + partition +
+                ", staleStores=" + staleStores +
+                ", storeName=" + storeName +
+                ", queryableStoreType=" + queryableStoreType +
+                '}';
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(partition, staleStores, storeName, queryableStoreType);
+    }
+}

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

@@ -1499,7 +1499,7 @@ public void addPredecessor(final TopologyDescription.Node predecessor) {
         @Override
         public String toString() {
             final String topicsString = topics == null ? topicPattern.toString() : topics.toString();
-            
+
             return "Source: " + name + " (topics: " + topicsString + ")\n      --> " + nodeNames(successors);
         }
 
@@ -1701,7 +1701,7 @@ public int hashCode() {
 
     public static class TopicsInfo {
         final Set<String> sinkTopics;
-        final Set<String> sourceTopics;
+        public final Set<String> sourceTopics;
         public final Map<String, InternalTopicConfig> stateChangelogTopics;
         public final Map<String, InternalTopicConfig> repartitionSourceTopics;
 

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import org.apache.kafka.streams.StoreQueryParams;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.state.QueryableStoreType;
@@ -41,29 +42,29 @@ public QueryableStoreProvider(final List<StreamThreadStateStoreProvider> storePr
      * Get a composite object wrapping the instances of the {@link StateStore} with the provided
      * storeName and {@link QueryableStoreType}
      *
-     * @param storeName          name of the store
-     * @param queryableStoreType accept stores passing {@link QueryableStoreType#accepts(StateStore)}
-     * @param includeStaleStores     if true, include standbys and recovering stores;
-     *                                        if false, only include running actives.
+     * @param storeQueryParams       if stateStoresEnabled is used i.e. staleStoresEnabled is true, include standbys and recovering stores;
+     *                                        if stateStoresDisabled i.e. staleStoresEnabled is false, only include running actives;
+     *                                        if partition is null then it fetches all local partitions on the instance;
+     *                                        if partition is set then it fetches a specific partition.
      * @param <T>                The expected type of the returned store
      * @return A composite object that wraps the store instances.
      */
-    public <T> T getStore(final String storeName,
-                          final QueryableStoreType<T> queryableStoreType,
-                          final boolean includeStaleStores) {
+    public <T> T getStore(final StoreQueryParams<T> storeQueryParams) {
+        final String storeName = storeQueryParams.storeName();
+        final QueryableStoreType<T> queryableStoreType = storeQueryParams.queryableStoreType();
         final List<T> globalStore = globalStoreProvider.stores(storeName, queryableStoreType);
         if (!globalStore.isEmpty()) {
             return queryableStoreType.create(globalStoreProvider, storeName);
         }
         final List<T> allStores = new ArrayList<>();
         for (final StreamThreadStateStoreProvider storeProvider : storeProviders) {
-            allStores.addAll(storeProvider.stores(storeName, queryableStoreType, includeStaleStores));
+            allStores.addAll(storeProvider.stores(storeQueryParams));
         }
         if (allStores.isEmpty()) {
             throw new InvalidStateStoreException("The state store, " + storeName + ", may have migrated to another instance.");
         }
         return queryableStoreType.create(
-            new WrappingStoreProvider(storeProviders, includeStaleStores),
+            new WrappingStoreProvider(storeProviders, storeQueryParams),
             storeName
         );
     }

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

@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import org.apache.kafka.streams.StoreQueryParams;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder;
 import org.apache.kafka.streams.processor.internals.StreamThread;
 import org.apache.kafka.streams.processor.internals.Task;
 import org.apache.kafka.streams.state.QueryableStoreType;
@@ -30,27 +32,36 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 public class StreamThreadStateStoreProvider {
 
     private final StreamThread streamThread;
+    private final InternalTopologyBuilder internalTopologyBuilder;
 
-    public StreamThreadStateStoreProvider(final StreamThread streamThread) {
+    public StreamThreadStateStoreProvider(final StreamThread streamThread,
+                                          final InternalTopologyBuilder internalTopologyBuilder) {
         this.streamThread = streamThread;
+        this.internalTopologyBuilder = internalTopologyBuilder;
     }
 
     @SuppressWarnings("unchecked")
-    public <T> List<T> stores(final String storeName,
-                              final QueryableStoreType<T> queryableStoreType,
-                              final boolean includeStaleStores) {
+    public <T> List<T> stores(final StoreQueryParams storeQueryParams) {
+        final String storeName = storeQueryParams.storeName();
+        final QueryableStoreType<T> queryableStoreType = storeQueryParams.queryableStoreType();
+        final TaskId keyTaskId = createKeyTaskId(storeName, storeQueryParams.partition());
         if (streamThread.state() == StreamThread.State.DEAD) {
             return Collections.emptyList();
         }
         final StreamThread.State state = streamThread.state();
-        if (includeStaleStores ? state.isAlive() : state == StreamThread.State.RUNNING) {
-            final Map<TaskId, ? extends Task> tasks = includeStaleStores ? streamThread.allTasks() : streamThread.activeTasks();
+        if (storeQueryParams.staleStoresEnabled() ? state.isAlive() : state == StreamThread.State.RUNNING) {
+            final Map<TaskId, ? extends Task> tasks = storeQueryParams.staleStoresEnabled() ? streamThread.allTasks() : streamThread.activeTasks();
             final List<T> stores = new ArrayList<>();
             for (final Task streamTask : tasks.values()) {
+                if (keyTaskId != null && !keyTaskId.equals(streamTask.id())) {
+                    continue;
+                }
                 final StateStore store = streamTask.getStore(storeName);
                 if (store != null && queryableStoreType.accepts(store)) {
                     if (!store.isOpen()) {
@@ -72,8 +83,23 @@ public <T> List<T> stores(final String storeName,
         } else {
             throw new InvalidStateStoreException("Cannot get state store " + storeName + " because the stream thread is " +
                                                      state + ", not RUNNING" +
-                                                     (includeStaleStores ? " or REBALANCING" : ""));
+                                                     (storeQueryParams.staleStoresEnabled() ? " or REBALANCING" : ""));
         }
     }
 
+    private TaskId createKeyTaskId(final String storeName, final Integer partition) {
+        if (partition == null) {
+            return null;
+        }
+        final List<String> sourceTopics = internalTopologyBuilder.stateStoreNameToSourceTopics().get(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" +
+                                                "not available on this instance");
+    }
 }

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import org.apache.kafka.streams.StoreQueryParams;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.state.QueryableStoreType;
 
@@ -28,20 +29,25 @@
 public class WrappingStoreProvider implements StateStoreProvider {
 
     private final List<StreamThreadStateStoreProvider> storeProviders;
-    private final boolean includeStaleStores;
+    private StoreQueryParams storeQueryParams;
 
     WrappingStoreProvider(final List<StreamThreadStateStoreProvider> storeProviders,
-                          final boolean includeStaleStores) {
+                          final StoreQueryParams storeQueryParams) {
         this.storeProviders = storeProviders;
-        this.includeStaleStores = includeStaleStores;
+        this.storeQueryParams = storeQueryParams;
+    }
+
+    //visible for testing
+    public void setStoreQueryParams(final StoreQueryParams storeQueryParams) {
+        this.storeQueryParams = storeQueryParams;
     }
 
     @Override
     public <T> List<T> stores(final String storeName,
                               final QueryableStoreType<T> queryableStoreType) {
         final List<T> allStores = new ArrayList<>();
         for (final StreamThreadStateStoreProvider provider : storeProviders) {
-            final List<T> stores = provider.stores(storeName, queryableStoreType, includeStaleStores);
+            final List<T> stores = provider.stores(storeQueryParams);
             allStores.addAll(stores);
         }
         if (allStores.isEmpty()) {