diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binding/DefaultPartitioningInterceptor.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binding/DefaultPartitioningInterceptor.java new file mode 100644 index 000000000..a0fdcde8b --- /dev/null +++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binding/DefaultPartitioningInterceptor.java @@ -0,0 +1,68 @@ +/* + * Copyright 2021-2021 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.binding; + +import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; +import org.springframework.cloud.stream.binder.BinderHeaders; +import org.springframework.cloud.stream.binder.PartitionHandler; +import org.springframework.cloud.stream.config.BindingProperties; +import org.springframework.integration.expression.ExpressionUtils; +import org.springframework.integration.support.MessageBuilder; +import org.springframework.messaging.Message; +import org.springframework.messaging.MessageChannel; +import org.springframework.messaging.support.ChannelInterceptor; + +/** + * + * @author Oleg Zhurakousky + * @since 3.1 + * + */ +public class DefaultPartitioningInterceptor implements ChannelInterceptor { + + private final PartitionHandler partitionHandler; + + public DefaultPartitioningInterceptor(BindingProperties bindingProperties, ConfigurableListableBeanFactory beanFactory) { + this.partitionHandler = new PartitionHandler( + ExpressionUtils.createStandardEvaluationContext(beanFactory), + bindingProperties.getProducer(), beanFactory); + } + + public void setPartitionCount(int partitionCount) { + this.partitionHandler.setPartitionCount(partitionCount); + } + + @Override + public Message preSend(Message message, MessageChannel channel) { + if (!message.getHeaders().containsKey(BinderHeaders.PARTITION_OVERRIDE)) { + int partition = this.partitionHandler.determinePartition(message); + return MessageBuilder + .fromMessage(message) + .setHeader(BinderHeaders.PARTITION_HEADER, partition).build(); + } + else { + return MessageBuilder + .fromMessage(message) + .setHeader(BinderHeaders.PARTITION_HEADER, + message.getHeaders() + .get(BinderHeaders.PARTITION_OVERRIDE)) + .removeHeader(BinderHeaders.PARTITION_OVERRIDE).build(); + } + } + +} + diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java index 68831974f..01ddedc87 100644 --- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java +++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java @@ -16,6 +16,7 @@ package org.springframework.cloud.stream.function; +import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; import java.util.Map.Entry; @@ -36,6 +37,8 @@ import org.springframework.cloud.stream.binder.BinderFactory; import org.springframework.cloud.stream.binder.ProducerProperties; import org.springframework.cloud.stream.binding.BinderAwareChannelResolver.NewDestinationBindingCallback; import org.springframework.cloud.stream.binding.BindingService; +import org.springframework.cloud.stream.binding.DefaultPartitioningInterceptor; +import org.springframework.cloud.stream.config.BindingProperties; import org.springframework.cloud.stream.config.BindingServiceProperties; import org.springframework.cloud.stream.messaging.DirectWithAttributesChannel; import org.springframework.context.ConfigurableApplicationContext; @@ -89,6 +92,8 @@ public final class StreamBridge implements SmartInitializingSingleton { private final BindingService bindingService; + private final Map streamBridgeFunctionCache; + /** * * @param functionCatalog instance of {@link FunctionCatalog} @@ -116,6 +121,7 @@ public final class StreamBridge implements SmartInitializingSingleton { return remove; } }; + this.streamBridgeFunctionCache = new HashMap<>(); } /** @@ -201,8 +207,7 @@ public final class StreamBridge implements SmartInitializingSingleton { ProducerProperties producerProperties = this.bindingServiceProperties.getProducerProperties(bindingName); SubscribableChannel messageChannel = this.resolveDestination(bindingName, producerProperties, binderName); - Function functionToInvoke = this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString()); - ((FunctionInvocationWrapper) functionToInvoke).setSkipOutputConversion(producerProperties.isUseNativeEncoding()); + Function functionToInvoke = this.getStreamBridgeFunction(outputContentType.toString(), producerProperties); if (producerProperties != null && producerProperties.isPartitioned()) { functionToInvoke = new PartitionAwareFunctionWrapper(functionToInvoke, this.applicationContext, producerProperties); @@ -215,12 +220,25 @@ public final class StreamBridge implements SmartInitializingSingleton { return messageChannel.send(resultMessage); } + private synchronized FunctionInvocationWrapper getStreamBridgeFunction(String outputContentType, ProducerProperties producerProperties) { + if (StringUtils.hasText(outputContentType) && this.streamBridgeFunctionCache.containsKey(outputContentType)) { + return this.streamBridgeFunctionCache.get(outputContentType); + } + else { + FunctionInvocationWrapper functionToInvoke = this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString()); + this.streamBridgeFunctionCache.put(outputContentType, functionToInvoke); + functionToInvoke.setSkipOutputConversion(producerProperties.isUseNativeEncoding()); + return functionToInvoke; + } + } + @Override public void afterSingletonsInstantiated() { if (this.initialized) { return; } FunctionRegistration> fr = new FunctionRegistration<>(v -> v, STREAM_BRIDGE_FUNC_NAME); + fr.getProperties().put("singleton", "false"); this.functionRegistry.register(fr.type(FunctionType.from(Object.class).to(Object.class).message())); Map channels = applicationContext.getBeansOfType(DirectWithAttributesChannel.class); for (Entry channelEntry : channels.entrySet()) { @@ -255,6 +273,11 @@ public final class StreamBridge implements SmartInitializingSingleton { this.bindingService.bindProducer(messageChannel, destinationName, false, binder); this.channelCache.put(destinationName, messageChannel); + if (producerProperties.isPartitioned()) { + BindingProperties bindingProperties = this.bindingServiceProperties.getBindingProperties(destinationName); + ((AbstractMessageChannel) messageChannel) + .addInterceptor(new DefaultPartitioningInterceptor(bindingProperties, this.applicationContext.getBeanFactory())); + } this.addInterceptors((AbstractMessageChannel) messageChannel, destinationName); }