Remove the ordering decorator and place it in Kafka source before the throttled commit strategy receive

Author: ozangunalp

smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/IncomingKafkaRecord.java

@@ -2,6 +2,7 @@
 
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -26,6 +27,7 @@ public class IncomingKafkaRecord<K, T> implements KafkaRecord<K, T>, MetadataInj
     private final KafkaCommitHandler commitHandler;
     private final KafkaFailureHandler onNack;
     private final T payload;
+    private Runnable afterProcessing;
 
     public IncomingKafkaRecord(ConsumerRecord<K, T> record,
             String channel,
@@ -125,17 +127,23 @@ public BiFunction<Throwable, Metadata, CompletionStage<Void>> getNackWithMetadat
 
     @Override
     public CompletionStage<Void> ack(Metadata metadata) {
-        return commitHandler.handle(this).subscribeAsCompletionStage();
+        CompletableFuture<Void> ack = commitHandler.handle(this).subscribeAsCompletionStage();
+        return afterProcessing == null ? ack : ack.thenRun(afterProcessing);
     }
 
     @Override
     public CompletionStage<Void> nack(Throwable reason, Metadata metadata) {
-        return onNack.handle(this, reason, metadata).subscribeAsCompletionStage();
+        CompletableFuture<Void> nack = onNack.handle(this, reason, metadata).subscribeAsCompletionStage();
+        return afterProcessing == null ? nack : nack.thenRun(afterProcessing);
     }
 
     @Override
     public synchronized void injectMetadata(Object metadata) {
         this.metadata = this.metadata.with(metadata);
     }
 
+    public void afterProcessing(Runnable runnable) {
+        this.afterProcessing = runnable;
+    }
+
 }

smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/impl/KafkaSource.java

@@ -28,6 +28,7 @@
 import io.smallrye.reactive.messaging.ClientCustomizer;
 import io.smallrye.reactive.messaging.health.HealthReport;
 import io.smallrye.reactive.messaging.kafka.*;
+import io.smallrye.reactive.messaging.kafka.api.IncomingKafkaRecordMetadata;
 import io.smallrye.reactive.messaging.kafka.commit.ContextHolder;
 import io.smallrye.reactive.messaging.kafka.commit.KafkaCommitHandler;
 import io.smallrye.reactive.messaging.kafka.fault.KafkaDeadLetterQueue;
@@ -36,6 +37,7 @@
 import io.smallrye.reactive.messaging.kafka.health.KafkaSourceHealth;
 import io.smallrye.reactive.messaging.kafka.tracing.KafkaOpenTelemetryInstrumenter;
 import io.smallrye.reactive.messaging.kafka.tracing.KafkaTrace;
+import io.smallrye.reactive.messaging.providers.helpers.PausableMulti;
 import io.vertx.core.impl.ContextInternal;
 import io.vertx.core.impl.VertxInternal;
 import io.vertx.mutiny.core.Vertx;
@@ -166,14 +168,43 @@ public KafkaSource(Vertx vertx,
                 log.unableToReadRecord(topics, t);
                 reportFailure(t, false);
             });
-
-            Multi<IncomingKafkaRecord<K, V>> incomingMulti = multi.onItem().transformToUni(rec -> {
-                IncomingKafkaRecord<K, V> record = new IncomingKafkaRecord<>(rec, channel, index, commitHandler,
-                        failureHandler, isCloudEventEnabled, isTracingEnabled);
+            Multi<IncomingKafkaRecord<K, V>> incomingMulti = multi.onItem()
+                    .transform(rec -> new IncomingKafkaRecord<>(rec, channel, index, commitHandler,
+                            failureHandler, isCloudEventEnabled, isTracingEnabled));
+
+            // preserve order by key if configured
+            ProcessingOrder processingOrder = ProcessingOrder.of(config.getThrottledProcessingOrder());
+            incomingMulti = switch (processingOrder) {
+                case UNORDERED -> incomingMulti;
+                case ORDERED_BY_KEY -> incomingMulti.group()
+                        .by(message -> TopicPartitionKey.ofKey(
+                                message.getMetadata(IncomingKafkaRecordMetadata.class).get().getRecord()))
+                        .onItem().transformToMulti(g -> {
+                            PausableMulti<IncomingKafkaRecord<K, V>> pausable = new PausableMulti<>(g, false);
+                            return pausable.invoke(rec -> {
+                                pausable.pause();
+                                rec.afterProcessing(pausable::resume);
+                            });
+                        })
+                        .merge(128);
+                case ORDERED_BY_PARTITION -> incomingMulti.group()
+                        .by(message -> TopicPartitionKey.ofPartition(
+                                message.getMetadata(IncomingKafkaRecordMetadata.class).get().getRecord()))
+                        .onItem().transformToMulti(g -> {
+                            PausableMulti<IncomingKafkaRecord<K, V>> pausable = new PausableMulti<>(g, false);
+                            return pausable.invoke(rec -> {
+                                pausable.pause();
+                                rec.afterProcessing(pausable::resume);
+                            });
+                        })
+                        .merge(128);
+            };
+
+            incomingMulti = incomingMulti.onItem().transformToUni(record -> {
                 if ((failureHandler instanceof KafkaDeadLetterQueue)
-                        && rec.headers() != null
-                        && rec.headers().lastHeader(DESERIALIZATION_FAILURE_DLQ) != null) {
-                    Header reasonMsgHeader = rec.headers().lastHeader(DESERIALIZATION_FAILURE_REASON);
+                        && record.getHeaders() != null
+                        && record.getHeaders().lastHeader(DESERIALIZATION_FAILURE_DLQ) != null) {
+                    Header reasonMsgHeader = record.getHeaders().lastHeader(DESERIALIZATION_FAILURE_REASON);
                     String message = reasonMsgHeader != null ? new String(reasonMsgHeader.value()) : null;
                     RecordDeserializationException reason = new RecordDeserializationException(
                             TopicPartitions.getTopicPartition(record), record.getOffset(), message, null);

smallrye-reactive-messaging-kafka/src/main/java/io/smallrye/reactive/messaging/kafka/impl/KeyBasedOrderingDecorator.java

@@ -10,7 +10,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
-import java.util.stream.Stream;
 
 import jakarta.enterprise.context.ApplicationScoped;
 
@@ -26,19 +25,20 @@
 
 public class ThrottledConcurrencyTest extends KafkaCompanionTestBase {
 
+    int partitions = 10;
+    int concurrency = 10;
+    int recordsParPartition = 100;
+
     @Test
     public void testUnorderedParallelProcessing() {
         // Create topic and produce messages with different keys
-        companion.topics().createAndWait(topic, 5);
+        companion.topics().createAndWait(topic, partitions);
 
         // Produce messages to different partitions
         companion.produceStrings()
-                .fromRecords(IntStream.range(0, 10).boxed().flatMap(i -> Stream.of(
-                        new ProducerRecord<>(topic, 0, "A", "A-" + i),
-                        new ProducerRecord<>(topic, 1, "B", "B-" + i),
-                        new ProducerRecord<>(topic, 2, "C", "C-" + i),
-                        new ProducerRecord<>(topic, 3, "D", "D-" + i),
-                        new ProducerRecord<>(topic, 4, "E", "E-" + i))).toList())
+                .fromRecords(IntStream.range(0, recordsParPartition).boxed().flatMap(i -> IntStream.range(0, partitions).boxed()
+                        .map(p -> new ProducerRecord<>(topic, p, "key-" + p, "value-" + i)))
+                        .toList())
                 .awaitCompletion();
 
         MapBasedConfig config = kafkaConfig("mp.messaging.incoming.unordered")
@@ -49,18 +49,20 @@ public void testUnorderedParallelProcessing() {
                 .with("key.deserializer", StringDeserializer.class.getName())
                 .with("commit-strategy", "throttled")
                 .withPrefix("")
-                .with("smallrye.messaging.worker.my-pool.max-concurrency", 5);
+                .with("smallrye.messaging.worker.my-pool.max-concurrency", concurrency);
 
         UnorderedParallelConsumer app = runApplication(config, UnorderedParallelConsumer.class);
         long start = System.currentTimeMillis();
         System.out.println("Started processing messages " + start);
 
         // Wait for all messages to be processed
-        await().atMost(1, TimeUnit.MINUTES).untilAsserted(() -> assertThat(app.received()).hasSize(50));
+        await().atMost(2, TimeUnit.MINUTES)
+                .untilAsserted(() -> assertThat(app.received())
+                        .hasSize(partitions * recordsParPartition));
         long duration = System.currentTimeMillis() - start;
         System.out.println("Processing duration: " + duration + " ms");
 
-        assertThat(app.maxConcurrency()).isEqualTo(5);
+        assertThat(app.maxConcurrency()).isEqualTo(concurrency);
     }
 
     @ApplicationScoped
@@ -93,16 +95,13 @@ public int maxConcurrency() {
     @Test
     public void testOrderedByKeyProcessing() {
         // Create topic and produce messages with different keys
-        companion.topics().createAndWait(topic, 5);
+        companion.topics().createAndWait(topic, partitions);
 
         // Produce messages to different keys
         companion.produceStrings()
-                .fromRecords(IntStream.range(0, 10).boxed().flatMap(i -> Stream.of(
-                        new ProducerRecord<>(topic, 0, "A", "A-" + i),
-                        new ProducerRecord<>(topic, 1, "B", "B-" + i),
-                        new ProducerRecord<>(topic, 2, "C", "C-" + i),
-                        new ProducerRecord<>(topic, 3, "D", "D-" + i),
-                        new ProducerRecord<>(topic, 4, "E", "E-" + i))).toList())
+                .fromRecords(IntStream.range(0, recordsParPartition).boxed().flatMap(i -> IntStream.range(0, partitions).boxed()
+                        .map(p -> new ProducerRecord<>(topic, p, "key-" + p, "value-" + i)))
+                        .toList())
                 .awaitCompletion();
 
         MapBasedConfig config = kafkaConfig("mp.messaging.incoming.key-ordered")
@@ -112,25 +111,27 @@ public void testOrderedByKeyProcessing() {
                 .with("value.deserializer", StringDeserializer.class.getName())
                 .with("key.deserializer", StringDeserializer.class.getName())
                 .with("throttled.processing-order", "ordered_by_key")
-                .with("commit-strategy", "throttled");
+                .with("commit-strategy", "throttled")
+                .withPrefix("")
+                .with("smallrye.messaging.worker.my-pool.max-concurrency", concurrency);
 
         OrderedByKeyParallelConsumer app = runApplication(config, OrderedByKeyParallelConsumer.class);
         long start = System.currentTimeMillis();
         System.out.println("Started processing messages " + start);
 
         // Wait for all messages to be processed
-        await().atMost(1, TimeUnit.MINUTES).untilAsserted(() -> assertThat(app.received()).hasSize(50));
+        await().atMost(2, TimeUnit.MINUTES)
+                .untilAsserted(() -> assertThat(app.received())
+                        .hasSize(partitions * recordsParPartition));
         long duration = System.currentTimeMillis() - start;
         System.out.println("Processing duration: " + duration + " ms");
 
         // With ordered-by-key, messages with the same key should be processed sequentially
         // So max concurrent processing per key should be 1
-        assertThat(app.keyMaxCounter("A")).hasValue(1);
-        assertThat(app.keyMaxCounter("B")).hasValue(1);
-        assertThat(app.keyMaxCounter("C")).hasValue(1);
-        assertThat(app.keyMaxCounter("D")).hasValue(1);
-        assertThat(app.keyMaxCounter("E")).hasValue(1);
-        assertThat(app.maxConcurrency()).isEqualTo(5);
+        for (int i = 0; i < partitions; i++) {
+            assertThat(app.keyMaxCounter("key-" + i)).hasValue(1);
+        }
+        assertThat(app.maxConcurrency()).isEqualTo(partitions);
     }
 
     @ApplicationScoped
@@ -177,16 +178,13 @@ public int maxConcurrency() {
     @Test
     public void testOrderedByPartitionProcessing() {
         // Create topic with 2 partitions
-        companion.topics().createAndWait(topic, 5);
+        companion.topics().createAndWait(topic, partitions);
 
         // Produce messages to different partitions
         companion.produceStrings()
-                .fromRecords(IntStream.range(0, 10).boxed().flatMap(i -> Stream.of(
-                        new ProducerRecord<>(topic, 0, "key-" + i, "p0-" + i),
-                        new ProducerRecord<>(topic, 1, "key-" + i, "p1-" + i),
-                        new ProducerRecord<>(topic, 2, "key-" + i, "p2-" + i),
-                        new ProducerRecord<>(topic, 3, "key-" + i, "p3-" + i),
-                        new ProducerRecord<>(topic, 4, "key-" + i, "p4-" + i))).toList())
+                .fromRecords(IntStream.range(0, recordsParPartition).boxed().flatMap(i -> IntStream.range(0, partitions).boxed()
+                        .map(p -> new ProducerRecord<>(topic, p, "key-" + p, "value-" + i)))
+                        .toList())
                 .awaitCompletion();
 
         MapBasedConfig config = kafkaConfig("mp.messaging.incoming.partition-ordered")
@@ -196,25 +194,27 @@ public void testOrderedByPartitionProcessing() {
                 .with("value.deserializer", StringDeserializer.class.getName())
                 .with("key.deserializer", StringDeserializer.class.getName())
                 .with("throttled.processing-order", "ordered_by_partition")
-                .with("commit-strategy", "throttled");
+                .with("commit-strategy", "throttled")
+                .withPrefix("")
+                .with("smallrye.messaging.worker.my-pool.max-concurrency", concurrency);
 
         OrderedByPartitionParallelConsumer app = runApplication(config, OrderedByPartitionParallelConsumer.class);
         long start = System.currentTimeMillis();
         System.out.println("Started processing messages " + start);
 
         // Wait for all messages to be processed
-        await().atMost(1, TimeUnit.MINUTES).untilAsserted(() -> assertThat(app.received()).hasSize(50));
+        await().atMost(2, TimeUnit.MINUTES)
+                .untilAsserted(() -> assertThat(app.received())
+                        .hasSize(partitions * recordsParPartition));
         long duration = System.currentTimeMillis() - start;
         System.out.println("Processing duration: " + duration + " ms");
 
         // With ordered-by-partition, messages from the same partition should be processed sequentially
         // So max concurrent processing per partition should be 1
-        assertThat(app.partitionMaxCounter(0)).hasValue(1);
-        assertThat(app.partitionMaxCounter(1)).hasValue(1);
-        assertThat(app.partitionMaxCounter(2)).hasValue(1);
-        assertThat(app.partitionMaxCounter(3)).hasValue(1);
-        assertThat(app.partitionMaxCounter(4)).hasValue(1);
-        assertThat(app.maxConcurrency()).isEqualTo(5);
+        for (int i = 0; i < partitions; i++) {
+            assertThat(app.partitionMaxCounter(i)).hasValue(1);
+        }
+        assertThat(app.maxConcurrency()).isEqualTo(partitions);
     }
 
     @ApplicationScoped

smallrye-reactive-messaging-kafka/src/test/java/io/smallrye/reactive/messaging/kafka/ThrottledConcurrencyTest.java

@@ -26,7 +26,6 @@
 import io.smallrye.reactive.messaging.kafka.fault.KafkaFailureHandler;
 import io.smallrye.reactive.messaging.kafka.fault.KafkaIgnoreFailure;
 import io.smallrye.reactive.messaging.kafka.impl.KafkaClientServiceImpl;
-import io.smallrye.reactive.messaging.kafka.impl.KeyBasedOrderingDecorator;
 import io.smallrye.reactive.messaging.kafka.impl.TopicPartitions;
 import io.smallrye.reactive.messaging.kafka.reply.KafkaRequestReplyFactory;
 import io.smallrye.reactive.messaging.kafka.reply.UUIDCorrelationIdHandler;
@@ -115,7 +114,6 @@ public void initWeld() {
         weld.addBeanClass(ObservationDecorator.class);
         weld.addBeanClass(OutgoingObservationDecorator.class);
         weld.addBeanClass(PausableChannelDecorator.class);
-        weld.addBeanClass(KeyBasedOrderingDecorator.class);
         weld.disableDiscovery();
     }