Implement batch disposition

Signed-off-by: Gabriele Santomaggio <[email protected]>

Author: Gsantomaggio

RabbitMQ.AMQP.Client/IConsumer.cs

@@ -136,6 +136,6 @@ public interface IBatchContext : IContext
         /// Get the current number of message contexts in the batch context.
         /// @return current number of message contexts in the batch
         /// </summary>
-        int Size();
+        int Count();
     }
 }

RabbitMQ.AMQP.Client/Impl/DeliveryContext.cs

@@ -159,7 +159,12 @@ public IBatchContext Batch()
         }
     }
 
-    internal class BatchDeliveryContext : IBatchContext
+    ///<summary>
+    /// BatchDeliveryContext is a client side helper class that allows
+    /// accumulating multiple message contexts and settling them at once.
+    /// It is thread-safe and can be used from multiple threads.
+    /// </summary>
+    public class BatchDeliveryContext : IBatchContext
     {
         private readonly List<IContext> _contexts = new();
         private readonly SemaphoreSlim _semaphore = new(1, 1);
@@ -265,6 +270,6 @@ public void Add(IContext context)
             _contexts.Add(context);
         }
 
-        public int Size() => _contexts.Count;
+        public int Count() => _contexts.Count;
     }
 }

RabbitMQ.AMQP.Client/PublicAPI.Unshipped.txt

@@ -126,7 +126,7 @@ RabbitMQ.AMQP.Client.IBackOffDelayPolicy.IsActive() -> bool
 RabbitMQ.AMQP.Client.IBackOffDelayPolicy.Reset() -> void
 RabbitMQ.AMQP.Client.IBatchContext
 RabbitMQ.AMQP.Client.IBatchContext.Add(RabbitMQ.AMQP.Client.IContext! context) -> void
-RabbitMQ.AMQP.Client.IBatchContext.Size() -> int
+RabbitMQ.AMQP.Client.IBatchContext.Count() -> int
 RabbitMQ.AMQP.Client.IBindingSpecification
 RabbitMQ.AMQP.Client.IBindingSpecification.Argument(string! key, object! value) -> RabbitMQ.AMQP.Client.IBindingSpecification!
 RabbitMQ.AMQP.Client.IBindingSpecification.Arguments(System.Collections.Generic.Dictionary<string!, object!>! arguments) -> RabbitMQ.AMQP.Client.IBindingSpecification!
@@ -514,6 +514,16 @@ RabbitMQ.AMQP.Client.Impl.AmqpStreamSpecification.InitialClusterSize(int initial
 RabbitMQ.AMQP.Client.Impl.AmqpStreamSpecification.MaxAge(System.TimeSpan maxAge) -> RabbitMQ.AMQP.Client.IStreamSpecification!
 RabbitMQ.AMQP.Client.Impl.AmqpStreamSpecification.MaxSegmentSizeBytes(RabbitMQ.AMQP.Client.ByteCapacity! maxSegmentSize) -> RabbitMQ.AMQP.Client.IStreamSpecification!
 RabbitMQ.AMQP.Client.Impl.AmqpStreamSpecification.Queue() -> RabbitMQ.AMQP.Client.IQueueSpecification!
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Accept() -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Add(RabbitMQ.AMQP.Client.IContext! context) -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Batch() -> RabbitMQ.AMQP.Client.IBatchContext!
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.BatchDeliveryContext() -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Count() -> int
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Discard() -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Discard(System.Collections.Generic.Dictionary<string!, object!>! annotations) -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Requeue() -> void
+RabbitMQ.AMQP.Client.Impl.BatchDeliveryContext.Requeue(System.Collections.Generic.Dictionary<string!, object!>! annotations) -> void
 RabbitMQ.AMQP.Client.Impl.BindingSpecification
 RabbitMQ.AMQP.Client.Impl.BindingSpecification.ArgsToMap() -> Amqp.Types.Map!
 RabbitMQ.AMQP.Client.Impl.BindingSpecification.BindingSpecification() -> void

Tests/Consumer/ConsumerDispositionTests.cs

@@ -2,6 +2,7 @@
 // and the Mozilla Public License, version 2.0.
 // Copyright (c) 2017-2024 Broadcom. All Rights Reserved. The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
 
+using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
 using RabbitMQ.AMQP.Client;
@@ -31,7 +32,7 @@ public async Task BatchAcceptDisposition()
                 {
                     Assert.NotNull(batch);
                     batch.Add(context);
-                    if (batch.Size() != batchSize)
+                    if (batch.Count() != batchSize)
                     {
                         return Task.CompletedTask;
                     }
@@ -70,7 +71,7 @@ public async Task BatchDiscardDisposition()
                 {
                     Assert.NotNull(batch);
                     batch.Add(context);
-                    if (batch.Size() != batchSize)
+                    if (batch.Count() != batchSize)
                     {
                         return Task.CompletedTask;
                     }
@@ -109,7 +110,7 @@ public async Task BatchDiscardAnnotationDisposition()
                 {
                     Assert.NotNull(batch);
                     batch.Add(context);
-                    if (batch.Size() != batchSize)
+                    if (batch.Count() != batchSize)
                     {
                         return Task.CompletedTask;
                     }
@@ -151,7 +152,7 @@ public async Task BatchRequeueDisposition()
                 {
                     Assert.NotNull(batch);
                     batch.Add(context);
-                    if (batch.Size() != batchSize)
+                    if (batch.Count() != batchSize)
                     {
                         return Task.CompletedTask;
                     }
@@ -188,7 +189,7 @@ public async Task BatchRequeueAnnotationsDisposition()
                 {
                     Assert.NotNull(batch);
                     batch.Add(context);
-                    if (batch.Size() != batchSize)
+                    if (batch.Count() != batchSize)
                     {
                         return Task.CompletedTask;
                     }
@@ -198,23 +199,68 @@ public async Task BatchRequeueAnnotationsDisposition()
 
                     const string annotationKey1 = "x-opt-annotation1-key";
                     const string annotationValue1 = "annotation1-value";
-
+                    Assert.Equal(batchSize, batch.Count());
                     batch.Requeue(new Dictionary<string, object>()
                     {
                         { annotationKey, annotationValue }, { annotationKey1, annotationValue1 }
                     });
+                    Assert.Equal(0, batch.Count());
+                    
                     tcs.SetResult(true);
 
                     return Task.CompletedTask;
                 })
                 .BuildAndStartAsync();
 
             Assert.NotNull(consumer);
-            await tcs.Task;
+            await tcs.Task.WaitAsync(TimeSpan.FromSeconds(20));
             await consumer.CloseAsync();
             await WaitUntilQueueMessageCount(queueSpec, 18);
+            await queueSpec.DeleteAsync();
+        }
+
+        [Fact]
+        public async Task MixBatchAcceptAndDiscardDisposition()
+        {
+            Assert.NotNull(_connection);
+            Assert.NotNull(_management);
+
+            IQueueSpecification queueSpec = _management.Queue().Name(_queueName);
+            await queueSpec.DeclareAsync();
+            const int batchSize = 18;
+            await PublishAsync(queueSpec, batchSize * 2);
+            BatchDeliveryContext batch = new();
+            TaskCompletionSource<bool> tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            bool acceptNext = true; // Flag to alternate between accept and discard
+            IConsumer consumer = await _connection.ConsumerBuilder()
+                .Queue(queueSpec)
+                .MessageHandler((context, _) =>
+                {
+                    Assert.NotNull(batch);
+                    batch.Add(context);
+                    if (batch.Count() == batchSize && acceptNext)
+                    {
+                        Assert.Equal(batchSize, batch.Count());
+                        batch.Accept();
+                        acceptNext = false; // Switch to discard next
+                    }
+                    else if (batch.Count() == batchSize && !acceptNext)
+                    {
+                        Assert.Equal(batchSize, batch.Count());
+                        batch.Discard();
+                        tcs.SetResult(true);
+                    }
+                    return Task.CompletedTask;
+                })
+                .BuildAndStartAsync();
+
+            Assert.NotNull(consumer);
+            await tcs.Task;
+
             Assert.Equal(0, consumer.UnsettledMessageCount);
+            await WaitUntilQueueMessageCount(queueSpec, 0);
             await queueSpec.DeleteAsync();
+            await consumer.CloseAsync();
         }
     }
 }