<commit_before>/*
 * Copyright 2018-2022 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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.springframework.cloud.stream.binder.kafka.streams;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.KeyQueryMetadata;
import org.apache.kafka.streams.StoreQueryParameters;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.errors.InvalidStateStoreException;
import org.apache.kafka.streams.errors.UnknownStateStoreException;
import org.apache.kafka.streams.state.HostInfo;
import org.apache.kafka.streams.state.QueryableStoreType;
import org.apache.kafka.streams.state.StreamsMetadata;

import org.springframework.cloud.stream.binder.kafka.streams.properties.KafkaStreamsBinderConfigurationProperties;
import org.springframework.retry.RetryPolicy;
import org.springframework.retry.backoff.FixedBackOffPolicy;
import org.springframework.retry.policy.SimpleRetryPolicy;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.StringUtils;

/**
 * Services pertinent to the interactive query capabilities of Kafka Streams. This class
 * provides services such as querying for a particular store, which instance is hosting a
 * particular store etc. This is part of the public API of the kafka streams binder and
 * the users can inject this service in their applications to make use of it.
 *
 * @author <NAME>
 * @author <NAME>
 * @author <NAME>
 * @author <NAME>
 * @since 2.1.0
 */
public class InteractiveQueryService {

	private static final Log LOG = LogFactory.getLog(InteractiveQueryService.class);

	private final KafkaStreamsRegistry kafkaStreamsRegistry;

	private final KafkaStreamsBinderConfigurationProperties binderConfigurationProperties;

	/**
	 * Constructor for InteractiveQueryService.
	 * @param kafkaStreamsRegistry holding {@link KafkaStreamsRegistry}
	 * @param binderConfigurationProperties kafka Streams binder configuration properties
	 */
	public InteractiveQueryService(KafkaStreamsRegistry kafkaStreamsRegistry,
			KafkaStreamsBinderConfigurationProperties binderConfigurationProperties) {
		this.kafkaStreamsRegistry = kafkaStreamsRegistry;
		this.binderConfigurationProperties = binderConfigurationProperties;
	}

	/**
	 * Retrieve and return a queryable store by name created in the application.
	 * @param storeName name of the queryable store
	 * @param storeType type of the queryable store
	 * @param <T> generic queryable store
	 * @return queryable store.
	 */
	public <T> T getQueryableStore(String storeName, QueryableStoreType<T> storeType) {

		KafkaStreams contextSpecificKafkaStreams = getThreadContextSpecificKafkaStreams();

		return getRetryTemplate().execute(context -> {
			T store = null;
			Throwable throwable = null;
			if (contextSpecificKafkaStreams != null) {
				try {
					store = contextSpecificKafkaStreams.store(StoreQueryParameters.fromNameAndType(storeName, storeType));
				}
				catch (InvalidStateStoreException e) {
					// pass through..
					throwable = e;
				}
			}
			if (store != null) {
				return store;
			}
			if (contextSpecificKafkaStreams != null) {
				LOG.warn("Store (" + storeName + ") could not be found in Streams context, falling back to all known Streams instances");
			}

			for (KafkaStreams kafkaStreamApp : kafkaStreamsRegistry.getKafkaStreams()) {
				try {
					return getStateStoreFromKafkaStreams(kafkaStreamApp, storeName, storeType);
				}
				catch (Exception ex) {
					throwable = ex;
				}
			}
			throw new IllegalStateException("Error retrieving state store: " + storeName, throwable);
		});
	}

	private <T> T getStateStoreFromKafkaStreams(KafkaStreams kafkaStreams, String storeName, QueryableStoreType<T> storeType) {
		// Check KafkaStreams app knows about the state store
		T store = kafkaStreams.store(StoreQueryParameters.fromNameAndType(storeName, storeType));

		// Check KafkaStreams app actually has the state store
		if (kafkaStreams.streamsMetadataForStore(storeName).stream()
				.noneMatch((sm) -> sm.stateStoreNames().contains(storeName))) {
			throw new UnknownStateStoreException("Store (" + storeName + ") not available to Streams instance");
		}
		return store;
	}

	/**
	 * Retrieves the current {@link KafkaStreams} context if executing Thread is created by a Streams App (contains a matching application id in Thread's name).
	 *
	 * @return KafkaStreams instance associated with Thread
	 */
	private KafkaStreams getThreadContextSpecificKafkaStreams() {
		return this.kafkaStreamsRegistry.getKafkaStreams().stream()
				.filter(this::filterByThreadName).findAny().orElse(null);
	}

	/**
	 * Checks if the supplied {@link KafkaStreams} instance belongs to the calling Thread by matching the Thread's name with the Streams Application Id.
	 *
	 * @param streams {@link KafkaStreams} instance to filter
	 * @return true if Streams Instance is associated with Thread
	 */
	private boolean filterByThreadName(KafkaStreams streams) {<commit_msg>Smarter state store retrieval in InteractiveQueryService<commit_after>/*
 * Copyright 2018-2022 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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.springframework.cloud.stream.binder.kafka.streams;

import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.KeyQueryMetadata;
import org.apache.kafka.streams.StoreQueryParameters;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.errors.InvalidStateStoreException;
import org.apache.kafka.streams.errors.UnknownStateStoreException;
import org.apache.kafka.streams.processor.internals.TopologyMetadata;
import org.apache.kafka.streams.state.HostInfo;
import org.apache.kafka.streams.state.QueryableStoreType;
import org.apache.kafka.streams.state.StreamsMetadata;

import org.springframework.cloud.stream.binder.kafka.streams.properties.KafkaStreamsBinderConfigurationProperties;
import org.springframework.retry.RetryPolicy;
import org.springframework.retry.backoff.FixedBackOffPolicy;
import org.springframework.retry.policy.SimpleRetryPolicy;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;

/**
 * Services pertinent to the interactive query capabilities of Kafka Streams. This class
 * provides services such as querying for a particular store, which instance is hosting a
 * particular store etc. This is part of the public API of the kafka streams binder and
 * the users can inject this service in their applications to make use of it.
 *
 * @author <NAME>
 * @author <NAME>
 * @author <NAME>
 * @author <NAME>
 * @author <NAME>
 * @since 2.1.0
 */
public class InteractiveQueryService {

	private static final Log LOG = LogFactory.getLog(InteractiveQueryService.class);

	private final KafkaStreamsRegistry kafkaStreamsRegistry;

	private final KafkaStreamsBinderConfigurationProperties binderConfigurationProperties;

	private final Field topologyMetadataField;

	/**
	 * Constructor for InteractiveQueryService.
	 * @param kafkaStreamsRegistry holding {@link KafkaStreamsRegistry}
	 * @param binderConfigurationProperties kafka Streams binder configuration properties
	 */
	public InteractiveQueryService(KafkaStreamsRegistry kafkaStreamsRegistry,
			KafkaStreamsBinderConfigurationProperties binderConfigurationProperties) {
		this.kafkaStreamsRegistry = kafkaStreamsRegistry;
		this.binderConfigurationProperties = binderConfigurationProperties;
		this.topologyMetadataField = ReflectionUtils.findField(KafkaStreams.class, "topologyMetadata");
		this.topologyMetadataField.setAccessible(true);
	}

	/**
	 * Retrieve and return a queryable store by name created in the application.
	 * @param storeName name of the queryable store
	 * @param storeType type of the queryable store
	 * @param <T> generic queryable store
	 * @return queryable store.
	 */
	public <T> T getQueryableStore(String storeName, QueryableStoreType<T> storeType) {

		KafkaStreams contextSpecificKafkaStreams = getThreadContextSpecificKafkaStreams();

		final StoreQueryParameters<T> storeQueryParams = StoreQueryParameters.fromNameAndType(storeName, storeType);

		return getRetryTemplate().execute(context -> {
			T store = null;
			Throwable throwable = null;
			if (contextSpecificKafkaStreams != null) {
				try {
					store = contextSpecificKafkaStreams.store(storeQueryParams);
				}
				catch (InvalidStateStoreException e) {
					throwable = e;
				}
			}
			if (store != null) {
				return store;
			}
			if (contextSpecificKafkaStreams != null) {
				LOG.warn("Store (" + storeName + ") could not be found in Streams context, falling back to all known Streams instances");
			}

			// Find all apps that know about the store
			Map<KafkaStreams, T> candidateStores = new HashMap<>();
			for (KafkaStreams kafkaStreamApp : kafkaStreamsRegistry.getKafkaStreams()) {
				try {
					candidateStores.put(kafkaStreamApp, kafkaStreamApp.store(storeQueryParams));
				}
				catch (Exception ex) {
					throwable = ex;
				}
			}

			// Store exists in a single app - no further resolution required
			if (candidateStores.size() == 1) {
				return candidateStores.values().stream().findFirst().get();
			}

			// If the store is in multiple streams apps - discard any apps that do not actually have the store
			if (candidateStores.size() > 1) {

				candidateStores = candidateStores.entrySet().stream()
						.filter((e) -> storeActuallyAvailable(e.getKey(), storeName))
						.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

				if (candidateStores.size() == 1) {
					return candidateStores.values().stream().findFirst().get();
				}

				throwable = (candidateStores.size() == 0) ?
						new UnknownStateStoreException("Store (" + storeName + ") not available to Streams instance") :
						new InvalidStateStoreException("Store (" + storeName + ") available to more than one Streams instance");

			}
			throw new IllegalStateException("Error retrieving state store: " + storeName, throwable);
		});
	}

	private boolean storeActuallyAvailable(KafkaStreams kafkaStreams, String storeName) {
		try {
			TopologyMetadata topologyMetadata = (TopologyMetadata) ReflectionUtils.getField(topologyMetadataField, kafkaStreams);
			return !topologyMetadata.sourceTopicsForStore(storeName, null).isEmpty();
		}
		catch (Exception ex) {
			LOG.error("Unable to determine if store (" + storeName + ") is available to the app due to: " + ex.getMessage(), ex);
		}
		return false;
	}

	/**
	 * Retrieves the current {@link KafkaStreams} context if executing Thread is created by a Streams App (contains a matching application id in Thread's name).
	 *
	 * @return KafkaStreams instance associated with Thread
	 */
	private KafkaStreams getThreadContextSpecificKafkaStreams() {
		return this.kafkaStreamsRegistry.getKafkaStreams().stream()
				.filter(this::filterByThreadName).findAny().orElse(null);
	}

	/**
	 * Checks if the supplied {@link KafkaStreams} instance belongs to the calling Thread by matching the Thread's name with the Streams Application Id.
	 *
	 * @param streams {@link KafkaStreams} instance to filter
	 * @return true if Streams Instance is associated with Thread
	 */
	private boolean filterByThreadName(KafkaStreams streams) {
{'fraction_non_alphanumeric': 0.08720350042220004, 'fraction_numerical': 0.002609963921086973, 'mean_word_length': 4.844773441004935, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 28, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4633845', 'n_tokens_mistral': 3903, 'n_tokens_neox': 3532, 'n_words': 1048}