Kafka performance improvments and bug fixes

Author: cavus700

src/Motor.Extensions.Hosting.Kafka/KafkaMessageConsumer.cs

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Channels;
@@ -34,6 +33,7 @@ public sealed class KafkaMessageConsumer<TData> : IMessageConsumer<TData>, IDisp
     private readonly IMetricFamily<ISummary>? _consumerLagSummary;
     private readonly ILogger<KafkaMessageConsumer<TData>> _logger;
     private readonly IHostApplicationLifetime _applicationLifetime;
+    private readonly AsyncPolicy<ProcessedMessageStatus> _retryPolicy;
     private IConsumer<string?, byte[]>? _consumer;
     private readonly CancellationTokenSource _internalCts = new();
 
@@ -69,6 +69,13 @@ CloudEventFormatter cloudEventFormatter
             _options.MaxConcurrentMessages
         );
         _timer = new Timer(HandleCommitTimer);
+
+        _retryPolicy = Policy
+            .HandleResult<ProcessedMessageStatus>(status => status == ProcessedMessageStatus.TemporaryFailure)
+            .WaitAndRetryAsync(
+                _options.RetriesOnTemporaryFailure,
+                retryAttempt => _options.RetryBasePeriod * Math.Pow(2, retryAttempt)
+            );
     }
 
     public Func<
@@ -160,36 +167,41 @@ private void WriteLog(LogMessage logMessage)
             case SyslogLevel.Alert:
             case SyslogLevel.Critical:
                 _logger.LogCritical(
-                    $"{logMessage.Message} -(Facility: {{facility}}, Name: {{name}})",
+                    "{Message} -(Facility: {Facility}, Name: {Name})",
+                    logMessage.Message,
                     logMessage.Facility,
                     logMessage.Name
                 );
                 break;
             case SyslogLevel.Error:
                 _logger.LogError(
-                    $"{logMessage.Message} -(Facility: {{facility}}, Name: {{name}})",
+                    "{Message} -(Facility: {Facility}, Name: {Name})",
+                    logMessage.Message,
                     logMessage.Facility,
                     logMessage.Name
                 );
                 break;
             case SyslogLevel.Warning:
                 _logger.LogWarning(
-                    $"{logMessage.Message} -(Facility: {{facility}}, Name: {{name}})",
+                    "{Message} -(Facility: {Facility}, Name: {Name})",
+                    logMessage.Message,
                     logMessage.Facility,
                     logMessage.Name
                 );
                 break;
             case SyslogLevel.Notice:
             case SyslogLevel.Info:
                 _logger.LogInformation(
-                    $"{logMessage.Message} -(Facility: {{facility}}, Name: {{name}})",
+                    "{Message} -(Facility: {Facility}, Name: {Name})",
+                    logMessage.Message,
                     logMessage.Facility,
                     logMessage.Name
                 );
                 break;
             case SyslogLevel.Debug:
                 _logger.LogDebug(
-                    $"{logMessage.Message} -(Facility: {{facility}}, Name: {{name}})",
+                    "{Message} -(Facility: {Facility}, Name: {Name})",
+                    logMessage.Message,
                     logMessage.Facility,
                     logMessage.Name
                 );
@@ -201,26 +213,37 @@ private void WriteLog(LogMessage logMessage)
 
     private void WriteStatistics(string json)
     {
-        var partitionConsumerLags = JsonSerializer
-            .Deserialize<KafkaStatistics>(json)
-            ?.Topics?.Select(t => t.Value)
-            .SelectMany(t => t.Partitions ?? new Dictionary<string, KafkaStatisticsPartition>())
-            .Select(t => (Parition: t.Key.ToString(), t.Value.ConsumerLag));
-        if (partitionConsumerLags is null)
+        using var document = JsonDocument.Parse(json);
+        var root = document.RootElement;
+
+        if (!root.TryGetProperty("topics", out var topics))
         {
             return;
         }
 
-        foreach (var (partition, consumerLag) in partitionConsumerLags)
+        foreach (var topic in topics.EnumerateObject())
         {
-            var lag = consumerLag;
-            if (lag == -1)
+            if (!topic.Value.TryGetProperty("partitions", out var partitions))
             {
-                lag = 0;
+                continue;
             }
 
-            _consumerLagSummary?.WithLabels(_options.Topic, partition)?.Observe(lag);
-            _consumerLagGauge?.WithLabels(_options.Topic, partition)?.Set(lag);
+            foreach (var partition in partitions.EnumerateObject())
+            {
+                if (!partition.Value.TryGetProperty("consumer_lag", out var consumerLagElement))
+                {
+                    continue;
+                }
+
+                var lag = consumerLagElement.GetInt64();
+                if (lag == -1)
+                {
+                    lag = 0;
+                }
+
+                _consumerLagSummary?.WithLabels(_options.Topic, partition.Name)?.Observe(lag);
+                _consumerLagGauge?.WithLabels(_options.Topic, partition.Name)?.Set(lag);
+            }
         }
     }
 
@@ -241,13 +264,7 @@ CancellationToken token
             );
             var cloudEvent = KafkaMessageToCloudEvent(msg.Message);
 
-            var retryPolicy = Policy
-                .HandleResult<ProcessedMessageStatus>(status => status == ProcessedMessageStatus.TemporaryFailure)
-                .WaitAndRetryAsync(
-                    _options.RetriesOnTemporaryFailure,
-                    retryAttempt => _options.RetryBasePeriod * Math.Pow(2, retryAttempt)
-                );
-            var status = await retryPolicy.ExecuteAsync(
+            var status = await _retryPolicy.ExecuteAsync(
                 (cancellationToken) => ConsumeCallbackAsync!.Invoke(cloudEvent, cancellationToken),
                 token
             );
@@ -272,6 +289,7 @@ CancellationToken token
     private readonly Timer _timer;
     private readonly object _commitLock = new();
     private bool _pendingCommit;
+    private int _messagesSinceLastCommit;
 
     private async Task ExecuteCommitLoopAsync(CancellationToken cancellationToken)
     {
@@ -281,7 +299,7 @@ private async Task ExecuteCommitLoopAsync(CancellationToken cancellationToken)
         {
             try
             {
-                var result = await PeekAndAwaitProcessedMessages(cancellationToken);
+                var result = await ReadAndAwaitProcessedMessage(cancellationToken);
 
                 if (IsIrrecoverableFailure(result.ProcessedMessageStatus))
                 {
@@ -290,16 +308,16 @@ private async Task ExecuteCommitLoopAsync(CancellationToken cancellationToken)
                     break;
                 }
 
-                // Remove message from channel, when Task is successfully completed
-                await _processedMessages.Reader.ReadAsync(cancellationToken);
-
                 lock (_commitLock)
                 {
                     _consumer?.StoreOffset(result.ConsumeResult);
                     _pendingCommit = true;
+                    _messagesSinceLastCommit++;
                 }
 
-                if ((result.ConsumeResult.Offset.Value + 1) % _options.CommitPeriod == 0)
+                // Use message count since last commit instead of offset-based check.
+                // This works correctly across multiple partitions with non-contiguous offsets.
+                if (_messagesSinceLastCommit >= _options.CommitPeriod)
                 {
                     Commit();
                     RestartCommitTimer();
@@ -314,17 +332,11 @@ private async Task ExecuteCommitLoopAsync(CancellationToken cancellationToken)
         StopCommitTimer();
     }
 
-    private async Task<ConsumeResultAndProcessedMessageStatus> PeekAndAwaitProcessedMessages(
+    private async Task<ConsumeResultAndProcessedMessageStatus> ReadAndAwaitProcessedMessage(
         CancellationToken cancellationToken
     )
     {
-        await _processedMessages.Reader.WaitToReadAsync(cancellationToken);
-
-        if (!_processedMessages.Reader.TryPeek(out var consumeAndProcessTask))
-        {
-            throw new InvalidOperationException("Awaited channel data has been removed by another consumer");
-        }
-
+        var consumeAndProcessTask = await _processedMessages.Reader.ReadAsync(cancellationToken);
         return await consumeAndProcessTask;
     }
 
@@ -338,6 +350,7 @@ private void Commit()
             }
 
             _pendingCommit = false;
+            _messagesSinceLastCommit = 0;
             try
             {
                 _consumer?.Commit();
@@ -386,7 +399,7 @@ private bool IsIrrecoverableFailure(ProcessedMessageStatus status)
             default:
                 _logger.LogCritical(
                     LogEvents.UnknownProcessedMessageStatus,
-                    "Unknown processed message status {status}",
+                    "Unknown processed message status {Status}",
                     status
                 );
                 return true;

src/Motor.Extensions.Hosting.Kafka/KafkaMessagePublisher.cs

@@ -5,6 +5,7 @@
 using CloudNative.CloudEvents.Extensions;
 using CloudNative.CloudEvents.Kafka;
 using Confluent.Kafka;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Motor.Extensions.Hosting.Abstractions;
 using Motor.Extensions.Hosting.CloudEvents;
@@ -19,24 +20,54 @@ public class KafkaMessagePublisher<TOutput> : IRawMessagePublisher<TOutput>, IDi
     private readonly IProducer<string?, byte[]> _producer;
     private readonly KafkaPublisherOptions<TOutput> _options;
     private readonly PublisherOptions _publisherOptions;
+    private readonly ILogger<KafkaMessagePublisher<TOutput>> _logger;
 
     public KafkaMessagePublisher(
         IOptions<KafkaPublisherOptions<TOutput>> options,
         CloudEventFormatter cloudEventFormatter,
-        IOptions<PublisherOptions> publisherOptions
+        IOptions<PublisherOptions> publisherOptions,
+        ILogger<KafkaMessagePublisher<TOutput>> logger
     )
     {
         _cloudEventFormatter = cloudEventFormatter;
         _options = options.Value ?? throw new ArgumentNullException(nameof(options));
         _publisherOptions = publisherOptions.Value ?? throw new ArgumentNullException(nameof(publisherOptions));
+        _logger = logger;
         _producer = new ProducerBuilder<string?, byte[]>(_options).Build();
     }
 
-    public async Task PublishMessageAsync(MotorCloudEvent<byte[]> motorCloudEvent, CancellationToken token = default)
+    public Task PublishMessageAsync(MotorCloudEvent<byte[]> motorCloudEvent, CancellationToken token = default)
     {
         var topic = motorCloudEvent.GetKafkaTopic() ?? _options.Topic;
         var message = CloudEventToKafkaMessage(motorCloudEvent);
-        await _producer.ProduceAsync(topic, message, token);
+
+        // Use Produce with a TaskCompletionSource for pipelining instead of
+        // awaiting ProduceAsync per message. This allows librdkafka to batch
+        // multiple messages into a single broker request, significantly
+        // improving throughput.
+        var tcs = new TaskCompletionSource<DeliveryResult<string?, byte[]>>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        try
+        {
+            _producer.Produce(topic, message, deliveryReport =>
+            {
+                if (deliveryReport.Error.IsError)
+                {
+                    tcs.SetException(new ProduceException<string?, byte[]>(
+                        deliveryReport.Error, deliveryReport));
+                }
+                else
+                {
+                    tcs.SetResult(deliveryReport);
+                }
+            });
+        }
+        catch (ProduceException<string?, byte[]> ex)
+        {
+            tcs.SetException(ex);
+        }
+
+        return tcs.Task;
     }
 
     public Message<string?, byte[]> CloudEventToKafkaMessage(MotorCloudEvent<byte[]> motorCloudEvent)
@@ -57,6 +88,14 @@ public async Task PublishMessageAsync(MotorCloudEvent<byte[]> motorCloudEvent, C
 
     public void Dispose()
     {
+        try
+        {
+            _producer.Flush(TimeSpan.FromSeconds(10));
+        }
+        catch (Exception e)
+        {
+            _logger.LogWarning(e, "Error flushing producer during dispose");
+        }
         _producer.Dispose();
     }
 }

src/Motor.Extensions.Hosting.Kafka/LogEvents.cs

@@ -14,7 +14,7 @@ public static class LogEvents
     public static readonly EventId UnknownProcessedMessageStatus = new(7, nameof(UnknownProcessedMessageStatus));
 
     public static readonly EventId MessageHandlingUnexpectedException = new(
-        7,
+        8,
         nameof(MessageHandlingUnexpectedException)
     );
 }

src/Motor.Extensions.Hosting.Kafka/Options/KafkaPublisherOptions.cs

@@ -4,5 +4,14 @@ namespace Motor.Extensions.Hosting.Kafka.Options;
 
 public class KafkaPublisherOptions<T> : ProducerConfig
 {
+    public KafkaPublisherOptions()
+    {
+        // Allow librdkafka to batch messages for better throughput.
+        // LingerMs controls how long to wait for additional messages before sending a batch.
+        LingerMs ??= 5;
+        // Enable Snappy compression for better network throughput with low CPU overhead.
+        CompressionType ??= Confluent.Kafka.CompressionType.Snappy;
+    }
+
     public string? Topic { get; set; }
 }

test/Motor.Extensions.Hosting.Kafka_IntegrationTest/KafkaExtensionTests.cs

@@ -8,6 +8,7 @@
 using CloudNative.CloudEvents.SystemTextJson;
 using Confluent.Kafka;
 using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Moq;
 using Motor.Extensions.Hosting.Abstractions;
@@ -519,7 +520,8 @@ private KafkaMessagePublisher<T> GetPublisher<T>(string topic)
     {
         var options = Options.Create(GetPublisherConfig<T>(topic));
         var publisherOptions = Options.Create(new PublisherOptions());
-        return new KafkaMessagePublisher<T>(options, new JsonEventFormatter(), publisherOptions);
+        return new KafkaMessagePublisher<T>(options, new JsonEventFormatter(), publisherOptions,
+            Mock.Of<ILogger<KafkaMessagePublisher<T>>>());
     }
 
     private KafkaPublisherOptions<T> GetPublisherConfig<T>(string topic)

test/Motor.Extensions.Hosting.Kafka_UnitTest/KafkaMessageTests.cs

@@ -96,7 +96,8 @@ private static KafkaMessagePublisher<TData> GetKafkaPublisher<TData>(
         return new KafkaMessagePublisher<TData>(
             Options.Create(options),
             new JsonEventFormatter(),
-            Options.Create(new PublisherOptions { CloudEventFormat = format })
+            Options.Create(new PublisherOptions { CloudEventFormat = format }),
+            Mock.Of<ILogger<KafkaMessagePublisher<TData>>>()
         );
     }