From 5558f7f7fc695f49c5e802b46b288cf5ab2d6d58 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Mon, 22 Apr 2019 16:13:49 -0400 Subject: [PATCH] Custom state stores with functional Kafka Streams Introducing the ability to provide custom state stores as regular Spring beans and use them in the functional model of Kafka Streams binding. Resolves #642 --- .../KafkaStreamsFunctionProcessor.java | 35 +++- ...kaStreamsBinderWordCountFunctionTests.java | 1 - .../KafkaStreamsFunctionStateStoreTests.java | 158 ++++++++++++++++++ 3 files changed, 189 insertions(+), 5 deletions(-) create mode 100644 spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsFunctionStateStoreTests.java diff --git a/spring-cloud-stream-binder-kafka-streams/src/main/java/org/springframework/cloud/stream/binder/kafka/streams/KafkaStreamsFunctionProcessor.java b/spring-cloud-stream-binder-kafka-streams/src/main/java/org/springframework/cloud/stream/binder/kafka/streams/KafkaStreamsFunctionProcessor.java index 0677de1e2..ec7f15ad5 100644 --- a/spring-cloud-stream-binder-kafka-streams/src/main/java/org/springframework/cloud/stream/binder/kafka/streams/KafkaStreamsFunctionProcessor.java +++ b/spring-cloud-stream-binder-kafka-streams/src/main/java/org/springframework/cloud/stream/binder/kafka/streams/KafkaStreamsFunctionProcessor.java @@ -40,6 +40,7 @@ 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.state.KeyValueStore; +import org.apache.kafka.streams.state.StoreBuilder; import org.springframework.beans.BeansException; import org.springframework.beans.factory.BeanInitializationException; @@ -48,6 +49,7 @@ import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; import org.springframework.beans.factory.support.BeanDefinitionBuilder; import org.springframework.beans.factory.support.BeanDefinitionRegistry; import org.springframework.cloud.function.context.FunctionCatalog; +import org.springframework.cloud.function.core.FluxedConsumer; import org.springframework.cloud.function.core.FluxedFunction; import org.springframework.cloud.stream.binder.ConsumerProperties; import org.springframework.cloud.stream.binder.kafka.properties.KafkaConsumerProperties; @@ -67,6 +69,7 @@ import org.springframework.kafka.core.CleanupConfig; import org.springframework.messaging.MessageHeaders; import org.springframework.messaging.support.MessageBuilder; import org.springframework.util.Assert; +import org.springframework.util.CollectionUtils; import org.springframework.util.StringUtils; /** @@ -151,11 +154,22 @@ public class KafkaStreamsFunctionProcessor implements ApplicationContextAware { Object[] adaptedInboundArguments = adaptAndRetrieveInboundArguments(stringResolvableTypeMap, functionName); try { if (resolvableType.getRawClass() != null && resolvableType.getRawClass().equals(Consumer.class)) { + //TOOD: Investigate why looking up by Consumer returns null Consumer consumer = functionCatalog.lookup(Consumer.class, functionName); - consumer.accept(adaptedInboundArguments[0]); + if (consumer == null) { + FluxedConsumer fluxedConsumer = functionCatalog.lookup(FluxedConsumer.class, functionName); + Assert.isTrue(fluxedConsumer != null, + "No corresponding consumer beans found in the catalog"); + Object target = fluxedConsumer.getTarget(); + if (Consumer.class.isAssignableFrom(target.getClass())) { + consumer = (Consumer) target; + } + } + if (consumer != null) { + consumer.accept(adaptedInboundArguments[0]); + } } else { - Function function = functionCatalog.lookup(Function.class, functionName); Object target = null; if (function instanceof FluxedFunction) { @@ -179,11 +193,9 @@ public class KafkaStreamsFunctionProcessor implements ApplicationContextAware { final Iterator iterator = outputs.iterator(); if (result.getClass().isArray()) { - final int length = ((Object[]) result).length; String[] methodAnnotatedOutboundNames = new String[length]; - for (int j = 0; j < length; j++) { if (iterator.hasNext()) { final String next = iterator.next(); @@ -382,6 +394,21 @@ public class KafkaStreamsFunctionProcessor implements ApplicationContextAware { BindingProperties bindingProperties, StreamsBuilder streamsBuilder, Serde keySerde, Serde valueSerde, Topology.AutoOffsetReset autoOffsetReset) { + try { + final Map storeBuilders = applicationContext.getBeansOfType(StoreBuilder.class); + if (!CollectionUtils.isEmpty(storeBuilders)) { + storeBuilders.values().forEach(storeBuilder -> { + streamsBuilder.addStateStore(storeBuilder); + if (LOG.isInfoEnabled()) { + LOG.info("state store " + storeBuilder.name() + " added to topology"); + } + }); + } + } + catch (Exception e) { + // Pass through. + } + String[] bindingTargets = StringUtils .commaDelimitedListToStringArray(this.bindingServiceProperties.getBindingDestination(inboundName)); diff --git a/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsBinderWordCountFunctionTests.java b/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsBinderWordCountFunctionTests.java index 7a8eae7f7..175f04e4d 100644 --- a/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsBinderWordCountFunctionTests.java +++ b/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsBinderWordCountFunctionTests.java @@ -185,5 +185,4 @@ public class KafkaStreamsBinderWordCountFunctionTests { new Date(key.window().start()), new Date(key.window().end())))); } } - } diff --git a/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsFunctionStateStoreTests.java b/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsFunctionStateStoreTests.java new file mode 100644 index 000000000..616c1c7be --- /dev/null +++ b/spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/function/KafkaStreamsFunctionStateStoreTests.java @@ -0,0 +1,158 @@ +/* + * Copyright 2019-2019 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.function; + +import java.util.Map; + +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.kstream.KStream; +import org.apache.kafka.streams.processor.Processor; +import org.apache.kafka.streams.processor.ProcessorContext; +import org.apache.kafka.streams.processor.ProcessorSupplier; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.StoreBuilder; +import org.apache.kafka.streams.state.Stores; +import org.apache.kafka.streams.state.WindowStore; +import org.junit.ClassRule; +import org.junit.Test; + +import org.springframework.boot.SpringApplication; +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.cloud.stream.annotation.EnableBinding; +import org.springframework.cloud.stream.annotation.Input; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.kafka.core.DefaultKafkaProducerFactory; +import org.springframework.kafka.core.KafkaTemplate; +import org.springframework.kafka.test.EmbeddedKafkaBroker; +import org.springframework.kafka.test.rule.EmbeddedKafkaRule; +import org.springframework.kafka.test.utils.KafkaTestUtils; + +import static org.assertj.core.api.Assertions.assertThat; + +public class KafkaStreamsFunctionStateStoreTests { + + @ClassRule + public static EmbeddedKafkaRule embeddedKafkaRule = new EmbeddedKafkaRule(1, true, + "counts"); + + private static EmbeddedKafkaBroker embeddedKafka = embeddedKafkaRule.getEmbeddedKafka(); + + @Test + public void testKafkaStreamsFuncionWithMultipleStateStores() throws Exception { + SpringApplication app = new SpringApplication(StateStoreTestApplication.class); + app.setWebApplicationType(WebApplicationType.NONE); + + try (ConfigurableApplicationContext context = app.run("--server.port=0", + "--spring.jmx.enabled=false", + "--spring.cloud.stream.bindings.input.destination=words", + "--spring.cloud.stream.kafka.streams.default.consumer.application-id=basic-word-count-1", + "--spring.cloud.stream.kafka.streams.binder.configuration.commit.interval.ms=1000", + "--spring.cloud.stream.kafka.streams.binder.configuration.default.key.serde" + + "=org.apache.kafka.common.serialization.Serdes$StringSerde", + "--spring.cloud.stream.kafka.streams.binder.configuration.default.value.serde" + + "=org.apache.kafka.common.serialization.Serdes$StringSerde", + "--spring.cloud.stream.kafka.streams.binder.brokers=" + embeddedKafka.getBrokersAsString())) { + receiveAndValidate(context); + } + } + + private void receiveAndValidate(ConfigurableApplicationContext context) throws Exception { + Map senderProps = KafkaTestUtils.producerProps(embeddedKafka); + DefaultKafkaProducerFactory pf = new DefaultKafkaProducerFactory<>(senderProps); + try { + KafkaTemplate template = new KafkaTemplate<>(pf, true); + template.setDefaultTopic("words"); + template.sendDefault("foobar"); + Thread.sleep(2000L); + StateStoreTestApplication processorApplication = context + .getBean(StateStoreTestApplication.class); + + KeyValueStore state1 = processorApplication.state1; + assertThat(processorApplication.processed).isTrue(); + assertThat(state1 != null).isTrue(); + assertThat(state1.name()).isEqualTo("my-store"); + WindowStore state2 = processorApplication.state2; + assertThat(state2 != null).isTrue(); + assertThat(state2.name()).isEqualTo("other-store"); + assertThat(state2.persistent()).isTrue(); + } + finally { + pf.destroy(); + } + } + + @EnableBinding(KStreamProcessorX.class) + @EnableAutoConfiguration + static class StateStoreTestApplication { + + KeyValueStore state1; + WindowStore state2; + + boolean processed; + + @Bean + public java.util.function.Consumer> process() { + return input -> + input.process((ProcessorSupplier) () -> new Processor() { + @Override + @SuppressWarnings("unchecked") + public void init(ProcessorContext context) { + state1 = (KeyValueStore) context.getStateStore("my-store"); + state2 = (WindowStore) context.getStateStore("other-store"); + } + + @Override + public void process(Object key, String value) { + processed = true; + } + + @Override + public void close() { + if (state1 != null) { + state1.close(); + } + if (state2 != null) { + state2.close(); + } + } + }, "my-store", "other-store"); + } + + @Bean + public StoreBuilder myStore() { + return Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore("my-store"), Serdes.Long(), + Serdes.Long()); + } + + @Bean + public StoreBuilder otherStore() { + return Stores.windowStoreBuilder( + Stores.persistentWindowStore("other-store", + 1L, 3, 3L, false), Serdes.Long(), + Serdes.Long()); + } + } + + interface KStreamProcessorX { + @Input("input") + KStream input(); + } + +}