GH-3242: Key StreamBridge function cache by binding name

StreamBridge caches the FunctionInvocationWrapper used by send(..) under the
producer properties. The binding name only took part in that key when
partitionKeyExpression and ProducerProperties#getBindingName() were both set,
and getBindingName() is never populated on the instance StreamBridge reads:
BindingService#bindProducer populates it on the extended copy it creates for an
ExtendedPropertiesBinder, not on the original returned by
BindingServiceProperties#getProducerProperties.

A partitioned binding therefore shared its cached function with another binding,
and since the partition enhancer is left on the cached function after a send
that produced a partition header, the next send on the other binding failed with
IllegalArgumentException: Partition key cannot be null.

Take the binding name from the send(..) argument, which is always available, and
include it for partitioned bindings. Key the cache by a record of the five
properties rather than by their computed int hash, so bindings are told apart by
equality and no hash collision can make two of them share a function. Bindings
that are not partitioned keep sharing a cached function as before.

Signed-off-by: kdelay <kdelay20@gmail.com>
Resolves #3244

Author: kdelay

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
@@ -933,6 +996,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;
 
@@ -197,7 +197,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);
@@ -233,21 +233,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());
@@ -394,6 +385,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());