chore: simplify and nano-optimise batch producer

Author: AlexeyRaga

src/KafkaFlow/Producers/BatchProduceContext.cs

@@ -12,17 +12,11 @@ namespace KafkaFlow.Producers;
 /// </summary>
 internal sealed class BatchProduceContext
 {
-    private readonly ConcurrentBag<BatchMessageEntry> _entries = new();
+    private readonly ConcurrentQueue<BatchMessageEntry> _entries = new();
 
     public IReadOnlyList<BatchMessageEntry> GetEntries() =>
         _entries.OrderBy(x => x.BatchIndex).ToList();
 
-    public IReadOnlyList<BatchMessageEntry> GetSuccessfulEntries() =>
-        _entries.Where(e => e.IsSuccess).OrderBy(e => e.BatchIndex).ToList();
-
-    public IReadOnlyList<BatchMessageEntry> GetFailedEntries() =>
-        _entries.Where(e => !e.IsSuccess).OrderBy(e => e.BatchIndex).ToList();
-
     /// <summary>
     /// Registers a message that completed middleware processing successfully.
     /// </summary>
@@ -32,7 +26,7 @@ public void Register(
         Message<byte[], byte[]> messageToProduce,
         IMessageContext context)
     {
-        _entries.Add(new BatchMessageEntry(batchIndex, item, true, messageToProduce, context, null));
+        _entries.Enqueue(new BatchMessageEntry(batchIndex, item, true, messageToProduce, context, null));
     }
 
     /// <summary>
@@ -46,7 +40,7 @@ public void RegisterFailure(int batchIndex, BatchProduceItem item, Exception exc
             Error = new Error(ErrorCode.Local_Fail, exception.Message),
             Status = PersistenceStatus.NotPersisted,
         };
-        _entries.Add(new BatchMessageEntry(batchIndex, item, false, null, null, exception));
+        _entries.Enqueue(new BatchMessageEntry(batchIndex, item, false, null, null, exception));
     }
 }
 

src/KafkaFlow/Producers/BatchProduceException.cs

@@ -4,7 +4,7 @@
 namespace KafkaFlow.Producers;
 
 /// <summary>
-/// Exception thrown by <see cref="M:BatchProduceExtension.BatchProduceAsync"/>
+/// Exception thrown by <see cref="M:MessageProducer.BatchProduceAsync"/>
 /// </summary>
 public class BatchProduceException : Exception
 {

src/KafkaFlow/Producers/MessageProducer.cs

@@ -211,17 +211,17 @@ public async Task<IReadOnlyCollection<BatchProduceItem>> BatchProduceAsync(
         // All items are now registered (either success with message, or failure with DeliveryReport set)
         await Task.WhenAll(middlewareTasks).ConfigureAwait(false);
 
-        // If instructed, throw if any message failed during middleware processing
-        var unsuccessful = batchContext.GetFailedEntries();
-        if (unsuccessful.Count > 0 && throwIfAnyProduceFail)
+        var allEntries = batchContext.GetEntries();
+        var successfulEntries = allEntries.Where(e => e.IsSuccess).ToList();
+
+        var errors = allEntries.Where(e => !e.IsSuccess).Select(e => e.Error).ToList();
+
+        if (errors.Count > 0 && throwIfAnyProduceFail)
         {
-            var exceptions = unsuccessful.Select(e => e.Error).ToList();
-            var aggregateException = new AggregateException("One or more messages failed during middleware processing", exceptions);
+            var aggregateException = new AggregateException("One or more messages failed during middleware processing", errors);
             throw new BatchProduceException(items, aggregateException);
         }
 
-        // Produce the messages that passed middleware
-        var successfulEntries = batchContext.GetSuccessfulEntries();
         if (successfulEntries.Count > 0)
         {
             try
@@ -247,7 +247,7 @@ public async Task<IReadOnlyCollection<BatchProduceItem>> BatchProduceAsync(
             }
         }
 
-        return batchContext.GetEntries().Select(e => e.Item).ToList();
+        return allEntries.Select(e => e.Item).ToList();
     }
 
     private async Task ExecuteMessageWithBatchContextAsync(
@@ -483,7 +483,13 @@ void DeliveryHandler(DeliveryReport<byte[], byte[]> report)
                     {
                         FillContextWithResultMetadata(entry.Context, report);
                     }
-
+                }
+                catch (Exception ex)
+                {
+                    _logHandler.Error("Batch Produce Delivery Handler Error", ex, new { Report = report });
+                }
+                finally
+                {
                     // Set the DeliveryReport directly on the item
                     entry.Item.DeliveryReport = report;
 
@@ -492,10 +498,6 @@ void DeliveryHandler(DeliveryReport<byte[], byte[]> report)
                         completionSource.TrySetResult(true);
                     }
                 }
-                catch (Exception ex)
-                {
-                    _logHandler.Error("Batch Produce Delivery Handler Error", ex, new { Report = report });
-                }
             }
 
             try

tests/KafkaFlow.UnitTests/BatchProduce/BatchProduceContextTests.cs

@@ -88,64 +88,6 @@ public void GetEntries_ShouldReturnAllEntriesSortedByBatchIndex()
         entries[2].BatchIndex.Should().Be(2);
     }
 
-    [TestMethod]
-    public void GetSuccessfulEntries_ShouldReturnOnlySuccessfulEntriesSorted()
-    {
-        // Arrange
-        var context = new BatchProduceContext();
-        var item = CreateBatchProduceItem("test-topic");
-        var message = CreateMessage();
-        var messageContext = Mock.Of<IMessageContext>();
-
-        context.Register(0, item, message, messageContext);
-        context.RegisterFailure(1, item, new Exception("Error 1"));
-        context.Register(2, item, message, messageContext);
-        context.RegisterFailure(3, item, new Exception("Error 2"));
-        context.Register(4, item, message, messageContext);
-
-        // Act
-        var successful = context.GetSuccessfulEntries();
-
-        // Assert
-        successful.Should().HaveCount(3);
-        successful[0].BatchIndex.Should().Be(0);
-        successful[0].IsSuccess.Should().BeTrue();
-        successful[1].BatchIndex.Should().Be(2);
-        successful[1].IsSuccess.Should().BeTrue();
-        successful[2].BatchIndex.Should().Be(4);
-        successful[2].IsSuccess.Should().BeTrue();
-    }
-
-    [TestMethod]
-    public void GetFailedEntries_ShouldReturnOnlyFailedEntriesSorted()
-    {
-        // Arrange
-        var context = new BatchProduceContext();
-        var item = CreateBatchProduceItem("test-topic");
-        var message = CreateMessage();
-        var messageContext = Mock.Of<IMessageContext>();
-        var error1 = new InvalidOperationException("Error 1");
-        var error2 = new ArgumentException("Error 2");
-
-        context.Register(0, item, message, messageContext);
-        context.RegisterFailure(1, item, error1);
-        context.Register(2, item, message, messageContext);
-        context.RegisterFailure(3, item, error2);
-        context.Register(4, item, message, messageContext);
-
-        // Act
-        var failed = context.GetFailedEntries();
-
-        // Assert
-        failed.Should().HaveCount(2);
-        failed[0].BatchIndex.Should().Be(1);
-        failed[0].IsSuccess.Should().BeFalse();
-        failed[0].Error.Should().BeSameAs(error1);
-        failed[1].BatchIndex.Should().Be(3);
-        failed[1].IsSuccess.Should().BeFalse();
-        failed[1].Error.Should().BeSameAs(error2);
-    }
-
     [TestMethod]
     public void Register_ShouldBeThreadSafe()
     {
@@ -189,7 +131,7 @@ public void RegisterFailure_ShouldBeThreadSafe()
         Task.WaitAll(tasks);
 
         // Assert
-        var entries = context.GetFailedEntries();
+        var entries = context.GetEntries().Where(e => !e.IsSuccess).ToList();
         entries.Should().HaveCount(100);
         entries.Select(e => e.BatchIndex).Should().BeEquivalentTo(Enumerable.Range(0, 100));
     }
@@ -222,14 +164,12 @@ public void MixedOperations_ShouldBeThreadSafe()
 
         // Assert
         var allEntries = context.GetEntries();
-        var successful = context.GetSuccessfulEntries();
-        var failed = context.GetFailedEntries();
+        var successful = allEntries.Where(e => e.IsSuccess).ToList();
+        var failed = allEntries.Where(e => !e.IsSuccess).ToList();
 
         allEntries.Should().HaveCount(100);
         successful.Should().HaveCount(50);
         failed.Should().HaveCount(50);
-        successful.All(e => e.IsSuccess).Should().BeTrue();
-        failed.All(e => !e.IsSuccess).Should().BeTrue();
     }
 
     private static Message<byte[], byte[]> CreateMessage()