diff --git a/core/spring-cloud-stream-integration-tests/src/test/java/org/springframework/cloud/stream/function/StreamBridgeTests.java b/core/spring-cloud-stream-integration-tests/src/test/java/org/springframework/cloud/stream/function/StreamBridgeTests.java index c47923cd5..5b05fe521 100644 --- a/core/spring-cloud-stream-integration-tests/src/test/java/org/springframework/cloud/stream/function/StreamBridgeTests.java +++ b/core/spring-cloud-stream-integration-tests/src/test/java/org/springframework/cloud/stream/function/StreamBridgeTests.java @@ -22,11 +22,15 @@ import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.UUID; +import java.util.concurrent.BrokenBarrierException; +import java.util.concurrent.CyclicBarrier; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeoutException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; + import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -321,6 +325,42 @@ class StreamBridgeTests { } } + /* + * This test verifies that when a partition key expression is set, then scst_partition is always set, even in + * concurrent scenarios using function binding. + * See https://github.com/spring-cloud/spring-cloud-stream/issues/2961 for more details + */ + @Test + void scstPartitionAlwaysSetEvenInConcurrentScenariosWithFunctions() throws Exception { + try (ConfigurableApplicationContext context = new SpringApplicationBuilder( + TestChannelBinderConfiguration.getCompleteConfiguration(FunctionPartitionConfiguration.class)).web( + WebApplicationType.NONE).run("--spring.cloud.stream.function.definition=concurrentFunction", + "--spring.cloud.stream.bindings.concurrentFunction-out-0.producer.partition-count=3", + "--spring.cloud.stream.bindings.concurrentFunction-out-0.producer.partition-key-expression=headers['partitionKey']", + "--spring.jmx.enabled=false")) { + CyclicBarrier barrier = context.getBean(CyclicBarrier.class); + StreamBridge streamBridge = context.getBean(StreamBridge.class); + + Thread otherThread = new Thread(() -> streamBridge.send("concurrentFunction-in-0", "wait")); + otherThread.start(); + barrier.await(5, TimeUnit.SECONDS); // wait for thread to be started and function called + streamBridge.send("concurrentFunction-in-0", "passThrough"); + barrier.await(5, TimeUnit.SECONDS); // notifies thread to continue function + otherThread.join(); + + int messagesWithoutScstPartition = 2; // total messages + OutputDestination output = context.getBean(OutputDestination.class); + Message message = output.receive(1000, "concurrentFunction-out-0"); + while (message != null) { + if (message.getHeaders().containsKey("scst_partition")) { + messagesWithoutScstPartition--; + } + message = output.receive(1000, "concurrentFunction-out-0"); + } + assertThat(messagesWithoutScstPartition).isEqualTo(0); + } + } + @Test void withOutputContentTypeWildCardBindings() { try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration @@ -866,6 +906,31 @@ class StreamBridgeTests { } } + @EnableAutoConfiguration + public static class FunctionPartitionConfiguration { + + @Bean + public CyclicBarrier cyclicBarrierFunction() { + return new CyclicBarrier(2); + } + @Bean + public Function> concurrentFunction(StreamBridge bridge, BindingServiceProperties properties, CyclicBarrier cyclicBarrierFunction) { + return s -> { + if (s.startsWith("wait")) { + try { + cyclicBarrierFunction.await(5, TimeUnit.SECONDS); // wait for notifying main thread to send other event + cyclicBarrierFunction.await(5, TimeUnit.SECONDS); // wait for other event been sent + } catch (BrokenBarrierException | InterruptedException | TimeoutException e) { + throw new RuntimeException(e); + } + } + return MessageBuilder.withPayload(s) + .setHeader("partitionKey", s.length()) + .build(); + }; + } + } + @EnableAutoConfiguration public static class InterceptorConfiguration { @Bean