GH-2245 Ensure DefaultPartitionInterceptor is added to partitioned destinations managed by StreamBridge

Resolves #2245

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();
+		}
+	}
+
+}
+

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<String, FunctionInvocationWrapper> 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<Function<Object, Object>> 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<String, DirectWithAttributesChannel> channels = applicationContext.getBeansOfType(DirectWithAttributesChannel.class);
 		for (Entry<String, DirectWithAttributesChannel> 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);
 		}