GH-3242: Key StreamBridge function cache by binding name

core/spring-cloud-stream-integration-tests/src/test/java/org/springframework/cloud/stream/function/StreamBridgeTests.java

@@ -47,6 +47,8 @@
 import org.springframework.cloud.function.cloudevent.CloudEventMessageBuilder;
 import org.springframework.cloud.function.cloudevent.CloudEventMessageUtils;
 import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
+import org.springframework.cloud.stream.binder.BinderHeaders;
+import org.springframework.cloud.stream.binder.PartitionKeyExtractorStrategy;
 import org.springframework.cloud.stream.binder.test.InputDestination;
 import org.springframework.cloud.stream.binder.test.OutputDestination;
 import org.springframework.cloud.stream.binder.test.TestChannelBinderConfiguration;
@@ -487,6 +489,67 @@ void delayedSend() {
 		}
 	}
 
+	/*
+	 * Two bindings whose properties hash alike must still get their own function, which is why the
+	 * cache is keyed by value rather than by a hash of those properties. This pair collides under
+	 * Objects.hash(contentType, nativeEncoding, partitioned, partitionCount, bindingName): with a
+	 * cache keyed by that hash, the non-partitioned send picks up the function left partition-aware
+	 * by the previous send and fails with "Partition key cannot be null" as in GH-3242.
+	 */
+	@SuppressWarnings("rawtypes")
+	@Test
+	void partitionedBindingIsNotSharedWithHashCollidingBinding() throws Exception {
+		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(
+			TestChannelBinderConfiguration.getCompleteConfiguration(
+				PartitionKeyExtractorConfiguration.class)).web(WebApplicationType.NONE).run(
+			"--spring.cloud.stream.source=nonPartitioned",
+			"--spring.cloud.stream.bindings[A>].producer.partition-count=120",
+			"--spring.cloud.stream.bindings[A>].producer.partition-key-extractor-name=partitionKeyExtractor",
+			"--spring.cloud.stream.bindings.nonPartitioned-out-0.producer.partition-count=1",
+			"--spring.jmx.enabled=false")) {
+			StreamBridge streamBridge = context.getBean(StreamBridge.class);
+			Field field = ReflectionUtils.findField(StreamBridge.class, "streamBridgeFunctionCache");
+			Objects.requireNonNull(field).setAccessible(true);
+			Map functionCache = (Map) field.get(streamBridge);
+
+			streamBridge.send("A>", MessageBuilder.withPayload("partitioned").setHeader("partitionKey", "key").build());
+			streamBridge.send("nonPartitioned-out-0", MessageBuilder.withPayload("nonPartitioned").build());
+
+			assertThat(functionCache.size()).isEqualTo(2);
+
+			OutputDestination output = context.getBean(OutputDestination.class);
+			assertThat(output.receive(1000, "A>").getHeaders()
+				.containsKey(BinderHeaders.PARTITION_HEADER)).isTrue();
+			assertThat(output.receive(1000, "nonPartitioned-out-0").getHeaders()
+				.containsKey(BinderHeaders.PARTITION_HEADER)).isFalse();
+		}
+	}
+
+	// See https://github.com/spring-cloud/spring-cloud-stream/issues/3242
+	@Test
+	void test_3242() {
+		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(
+			TestChannelBinderConfiguration.getCompleteConfiguration(
+				PartitionKeyExtractorConfiguration.class)).web(WebApplicationType.NONE).run(
+			"--spring.cloud.stream.source=partitioned;nonPartitioned",
+			"--spring.cloud.stream.bindings.partitioned-out-0.producer.partition-count=7",
+			"--spring.cloud.stream.bindings.partitioned-out-0.producer.partition-key-extractor-name=partitionKeyExtractor",
+			"--spring.cloud.stream.bindings.nonPartitioned-out-0.producer.partition-count=1",
+			"--spring.jmx.enabled=false")) {
+			StreamBridge streamBridge = context.getBean(StreamBridge.class);
+
+			streamBridge.send("partitioned-out-0",
+				MessageBuilder.withPayload("partitioned").setHeader("partitionKey", "key").build());
+			streamBridge.send("nonPartitioned-out-0", MessageBuilder.withPayload("nonPartitioned").build());
+
+			OutputDestination output = context.getBean(OutputDestination.class);
+			assertThat(output.receive(1000, "partitioned-out-0").getHeaders()
+				.containsKey(BinderHeaders.PARTITION_HEADER)).isTrue();
+			assertThat(output.receive(1000, "nonPartitioned-out-0").getHeaders()
+				.containsKey(BinderHeaders.PARTITION_HEADER)).isFalse();
+		}
+	}
+
 	@Test
 	void withInterceptorsMatchedAgainstAllPatterns() {
 		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration
@@ -930,6 +993,16 @@ public static class EmptyConfiguration {
 
 	}
 
+	@EnableAutoConfiguration
+	public static class PartitionKeyExtractorConfiguration {
+
+		@Bean
+		public PartitionKeyExtractorStrategy partitionKeyExtractor() {
+			return message -> message.getHeaders().get("partitionKey");
+		}
+
+	}
+
 	@EnableAutoConfiguration
 	public static class EmptyConfigurationWithCustomConverters {
 

core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java

@@ -117,7 +117,7 @@ public final class StreamBridge implements StreamOperations, SmartInitializingSi
 
 	private final BindingService bindingService;
 
-	private final Map<Integer, FunctionInvocationWrapper> streamBridgeFunctionCache;
+	private final Map<StreamBridgeFunctionKey, FunctionInvocationWrapper> streamBridgeFunctionCache;
 
 	private final FunctionInvocationHelper<?> functionInvocationHelper;
 
@@ -196,7 +196,7 @@ public boolean send(String bindingName, @Nullable String binderName, Object data
 		ProducerProperties producerProperties = this.bindingServiceProperties.getProducerProperties(bindingName);
 		MessageChannel messageChannel = this.resolveDestination(bindingName, producerProperties, binderName);
 
-		Function functionToInvoke = this.getStreamBridgeFunction(outputContentType.toString(), producerProperties);
+		Function functionToInvoke = this.getStreamBridgeFunction(bindingName, outputContentType.toString(), producerProperties);
 
 		if (producerProperties != null && producerProperties.isPartitioned()) {
 			functionToInvoke = new PartitionAwareFunctionWrapper(functionToInvoke, this.applicationContext, producerProperties);
@@ -232,21 +232,12 @@ public boolean send(String bindingName, @Nullable String binderName, Object data
 		return messageChannel.send(resultMessage);
 	}
 
-	private int hashProducerProperties(ProducerProperties producerProperties, String outputContentType) {
-		int hash = outputContentType.hashCode()
-				+ Boolean.hashCode(producerProperties.isUseNativeEncoding())
-				+ Boolean.hashCode(producerProperties.isPartitioned())
-				+ producerProperties.getPartitionCount();
-
-		if (producerProperties.getPartitionKeyExpression() != null && producerProperties.getBindingName() != null) {
-			hash += producerProperties.getBindingName().hashCode();
-		}
-
-		return hash;
-	}
-
-	private FunctionInvocationWrapper getStreamBridgeFunction(String outputContentType, ProducerProperties producerProperties) {
-		int streamBridgeFunctionKey = this.hashProducerProperties(producerProperties, outputContentType);
+	private FunctionInvocationWrapper getStreamBridgeFunction(String bindingName, String outputContentType, ProducerProperties producerProperties) {
+		StreamBridgeFunctionKey streamBridgeFunctionKey = new StreamBridgeFunctionKey(outputContentType,
+				producerProperties.isUseNativeEncoding(),
+				producerProperties.isPartitioned(),
+				producerProperties.getPartitionCount(),
+				producerProperties.isPartitioned() ? bindingName : null);
 
 		return this.streamBridgeFunctionCache.computeIfAbsent(streamBridgeFunctionKey, key -> {
 			FunctionInvocationWrapper functionToInvoke = this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString());
@@ -393,6 +384,19 @@ private void closeChannelsGracefully() {
 		});
 	}
 
+	/*
+	 * Identifies the function cached for a send(..). A partitioned binding mutates the cached
+	 * function by setting the partition enhancer on it, so it must not share that function with
+	 * another binding; its binding name is therefore part of the key. The name is taken from the
+	 * send(..) argument, since ProducerProperties#getBindingName() is only populated for binders
+	 * that are not an ExtendedPropertiesBinder (see GH-3242). Non-partitioned bindings leave it
+	 * null and keep sharing a single function. Equality rather than a computed hash decides cache
+	 * hits, so two distinct bindings can never be conflated by a hash collision.
+	 */
+	private record StreamBridgeFunctionKey(String outputContentType, boolean useNativeEncoding,
+			boolean partitioned, int partitionCount, String bindingName) {
+	}
+
 	private static final class ContextPropagationHelper {
 		static ExecutorService wrap(ExecutorService executorService) {
 			return ContextExecutorService.wrap(executorService, () -> ContextSnapshotFactory.builder().build().captureAll());