GH-2297: Concurrency in Reactor Kafka Binder (#2396)

* GH-2297: Concurrency in Reactor Kafka Binder

Resolves https://github.com/spring-cloud/spring-cloud-stream/issues/2297

* Remove unnecessary local var.

* Remove subscription cancellation; already handled by the super class.

binders/kafka-binder/spring-cloud-stream-binder-kafka-reactive/src/main/java/org/springframework/cloud/stream/binder/reactorkafka/ReactorKafkaBinder.java

@@ -16,7 +16,9 @@
 
 package org.springframework.cloud.stream.binder.reactorkafka;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicReference;
@@ -25,7 +27,6 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
-import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
@@ -75,7 +76,7 @@ public class ReactorKafkaBinder
 		implements
 		ExtendedPropertiesBinder<MessageChannel, KafkaConsumerProperties, KafkaProducerProperties> {
 
-	private static final Log log = LogFactory.getLog(ReactorKafkaBinder.class);
+	private static final Log logger = LogFactory.getLog(ReactorKafkaBinder.class);
 
 	private final KafkaBinderConfigurationProperties configurationProperties;
 
@@ -136,35 +137,40 @@ public class ReactorKafkaBinder
 				this.configurationProperties);
 		Assert.isInstanceOf(RecordMessageConverter.class, converter);
 		ReceiverOptions<Object, Object> opts = ReceiverOptions.create(configs)
-			.addAssignListener(parts -> System.out.println("Assigned: " + parts))
+			.addAssignListener(parts -> logger.info("Assigned: " + parts))
 			.subscription(Collections.singletonList(destination.getName()));
 
-		return new MessageProducerSupport() {
+		class ReactorMessageProducer extends MessageProducerSupport {
 
-			private final KafkaReceiver<Object, Object> receiver = KafkaReceiver.create(opts);
+			private final List<KafkaReceiver<Object, Object>> receivers = new ArrayList<>();
 
-			private volatile Subscription subscription;
+			ReactorMessageProducer() {
+				for (int i = 0; i < properties.getConcurrency(); i++) {
+					this.receivers.add(KafkaReceiver.create(opts));
+				}
+			}
 
 			@SuppressWarnings("unchecked")
 			@Override
 			protected void doStart() {
-				Flux<Message<Object>> flux = receiver
-						.receive()
-						.doOnSubscribe(subs -> this.subscription = subs)
-						.map(record -> (Message<Object>) ((RecordMessageConverter) converter)
-								.toMessage(record, null, null, null));
-				subscribeToPublisher(flux);
-			}
-
-			@Override
-			protected synchronized void doStop() {
-				if (this.subscription != null) {
-					this.subscription.cancel();
-					this.subscription = null;
+				List<Flux<Message<Object>>> fluxes = new ArrayList<>();
+				int concurrency = properties.getConcurrency();
+				for (int i = 0; i < concurrency; i++) {
+					fluxes.add(this.receivers.get(i)
+							.receive()
+							.map(record -> (Message<Object>) ((RecordMessageConverter) converter)
+									.toMessage(record, null, null, null)));
+				}
+				if (concurrency == 1) {
+					subscribeToPublisher(fluxes.get(0));
+				}
+				else {
+					subscribeToPublisher(Flux.merge(fluxes));
 				}
 			}
 
-		};
+		}
+		return new ReactorMessageProducer();
 	}
 
 	@Override

binders/kafka-binder/spring-cloud-stream-binder-kafka-reactive/src/test/java/org/springframework/cloud/stream/binder/reactorkafka/ReactorKafkaBinderTests.java

@@ -17,6 +17,8 @@
 package org.springframework.cloud.stream.binder.reactorkafka;
 
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
@@ -40,6 +42,7 @@ import org.springframework.integration.channel.FluxMessageChannel;
 import org.springframework.integration.support.MessageBuilder;
 import org.springframework.kafka.core.DefaultKafkaProducerFactory;
 import org.springframework.kafka.core.KafkaTemplate;
+import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.test.condition.EmbeddedKafkaCondition;
 import org.springframework.kafka.test.context.EmbeddedKafka;
 import org.springframework.kafka.test.utils.KafkaTestUtils;
@@ -55,7 +58,7 @@ import static org.mockito.Mockito.mock;
  * @since 4.0
  *
  */
-@EmbeddedKafka(topics = { "testC", "testP" })
+@EmbeddedKafka(topics = { "testC", "testC1", "testP" })
 public class ReactorKafkaBinderTests {
 
 	@SuppressWarnings({ "rawtypes", "unchecked" })
@@ -110,6 +113,72 @@ public class ReactorKafkaBinderTests {
 		pf.destroy();
 	}
 
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	@Test
+	void concurrency() throws Exception {
+		KafkaProperties kafkaProperties = new KafkaProperties();
+		kafkaProperties.setBootstrapServers(
+				Collections.singletonList(EmbeddedKafkaCondition.getBroker().getBrokersAsString()));
+		KafkaBinderConfigurationProperties binderProps = new KafkaBinderConfigurationProperties(kafkaProperties);
+		KafkaTopicProvisioner provisioner = new KafkaTopicProvisioner(binderProps, kafkaProperties, null);
+		ReactorKafkaBinder binder = new ReactorKafkaBinder(binderProps, provisioner);
+		binder.setApplicationContext(mock(GenericApplicationContext.class));
+
+		CountDownLatch subscriptionLatch = new CountDownLatch(1);
+		CountDownLatch messageLatch = new CountDownLatch(4);
+		Set<Integer> partitions = new HashSet<>();
+
+		FluxMessageChannel inbound = new FluxMessageChannel();
+		Subscriber<Message<?>> sub = new Subscriber<Message<?>>() {
+
+			@Override
+			public void onSubscribe(Subscription s) {
+				s.request(6);
+				subscriptionLatch.countDown();
+			}
+
+			@Override
+			public void onNext(Message<?> msg) {
+				partitions.add(msg.getHeaders().get(KafkaHeaders.RECEIVED_PARTITION, Integer.class));
+				messageLatch.countDown();
+			}
+
+			@Override
+			public void onError(Throwable t) {
+			}
+
+			@Override
+			public void onComplete() {
+			}
+
+		};
+		inbound.subscribe(sub);
+
+		KafkaConsumerProperties ext = new KafkaConsumerProperties();
+		ExtendedConsumerProperties<KafkaConsumerProperties> props =
+				new ExtendedConsumerProperties<KafkaConsumerProperties>(ext);
+		props.setConcurrency(2);
+
+		Binding<MessageChannel> consumer = binder.bindConsumer("testC1", "foo", inbound, props);
+
+		assertThat(subscriptionLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		DefaultKafkaProducerFactory pf =
+				new DefaultKafkaProducerFactory<>(KafkaTestUtils.producerProps(EmbeddedKafkaCondition.getBroker()));
+		KafkaTemplate kt = new KafkaTemplate<>(pf);
+		kt.send("testC1", 0, null, "foo").get(10, TimeUnit.SECONDS);
+		kt.send("testC1", 1, null, "bar").get(10, TimeUnit.SECONDS);
+		kt.send("testC1", 0, null, "baz").get(10, TimeUnit.SECONDS);
+		kt.send("testC1", 1, null, "qux").get(10, TimeUnit.SECONDS);
+		consumer.stop();
+		consumer.start();
+		kt.send("testC1", 0, null, "fiz").get(10, TimeUnit.SECONDS);
+		kt.send("testC1", 1, null, "buz").get(10, TimeUnit.SECONDS);
+		assertThat(messageLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(partitions).hasSize(2);
+		consumer.unbind();
+		pf.destroy();
+	}
+
 	@Test
 	void producerBinding() throws InterruptedException {
 		KafkaProperties kafkaProperties = new KafkaProperties();