Merge pull request #123 from patrikwlund/blocking-consumer-delivery

odalet · web-flow · commit 241d00450775 · 2022-12-14T17:15:49.000+01:00
Allow blocking/synchronous delivery to consumers
diff --git a/src/AddUp.FakeRabbitMQ.Tests/FakeModelBasicTests.cs b/src/AddUp.FakeRabbitMQ.Tests/FakeModelBasicTests.cs
@@ -3,6 +3,7 @@
 using System.Linq;
 using System.Text;
 using System.Threading;
+using FluentAssertions;
 using RabbitMQ.Client;
 using RabbitMQ.Client.Events;
 using RabbitMQ.Client.Exceptions;
@@ -425,6 +426,64 @@ public void BasicPublish_after_BasicConsume_does_not_deadlock()
             }
         }
 
+        [Fact]
+        public void BasicPublish_after_BasicConsume_with_BlockingDelivery_is_synchronous()
+        {
+            var server = new RabbitServer { BlockingConsumerDelivery = true };
+            using var model = new FakeModel(server);
+
+            model.ExchangeDeclare("my_exchange", ExchangeType.Direct);
+            model.QueueDeclare("my_queue");
+            model.ExchangeBind("my_queue", "my_exchange", null);
+
+            var trackedValue = new AsyncLocal<string>
+            {
+                Value = "initial"
+            };
+
+            var consumer = new EventingBasicConsumer(model);
+            consumer.Received += (_, _) =>
+            {
+                trackedValue.Value = "from_consumer";
+            };
+
+            model.BasicConsume("my_queue", false, consumer);
+
+            var encodedMessage = Encoding.ASCII.GetBytes("hello world!");
+            model.BasicPublish("my_exchange", null, model.CreateBasicProperties(), encodedMessage);
+
+            trackedValue.Value.Should().Be("from_consumer");
+        }
+
+        [Fact]
+        public void BasicPublish_before_BasicConsume_with_BlockingDelivery_is_synchronous()
+        {
+            var server = new RabbitServer { BlockingConsumerDelivery = true };
+            using var model = new FakeModel(server);
+
+            model.ExchangeDeclare("my_exchange", ExchangeType.Direct);
+            model.QueueDeclare("my_queue");
+            model.ExchangeBind("my_queue", "my_exchange", null);
+
+            var encodedMessage = Encoding.ASCII.GetBytes("hello world!");
+            model.BasicPublish("my_exchange", null, model.CreateBasicProperties(), encodedMessage);
+
+            var trackedValue = new AsyncLocal<string>
+            {
+                Value = "initial"
+            };
+
+            var consumer = new EventingBasicConsumer(model);
+            consumer.Received += (_, _) =>
+            {
+                trackedValue.Value = "from_consumer";
+            };
+
+            model.BasicConsume("my_queue", false, consumer);
+
+            trackedValue.Value.Should().Be("from_consumer");
+        }
+
         [Fact]
         public void BasicPublishBatch_publishes_messages()
         {
diff --git a/src/AddUp.FakeRabbitMQ/ConsumerDeliveryQueue.cs b/src/AddUp.FakeRabbitMQ/ConsumerDeliveryQueue.cs
@@ -0,0 +1,170 @@
+﻿using System.Threading.Channels;
+using System;
+using System.Threading.Tasks;
+using System.Threading;
+using RabbitMQ.Client.Events;
+
+namespace AddUp.RabbitMQ.Fakes
+{
+    internal abstract class ConsumerDeliveryQueue
+    {
+        protected readonly FakeModel model;
+        private readonly Action<CallbackExceptionEventArgs> onDeliveryException;
+
+        protected ConsumerDeliveryQueue(FakeModel model, Action<CallbackExceptionEventArgs> onDeliveryException)
+        {
+            this.model = model;
+            this.onDeliveryException = onDeliveryException;
+        }
+
+        public abstract void Deliver(Action deliveryAction);
+
+        protected void ExecuteDelivery(Action deliveryAction)
+        {
+            try
+            {
+                if (!model.IsOpen) return;
+                deliveryAction();
+            }
+            catch (Exception ex)
+            {
+                var callbackArgs = CallbackExceptionEventArgs.Build(ex, "");
+                onDeliveryException(callbackArgs);
+            }
+        }
+
+        /// <summary>
+        /// Marks the queue as complete, meaning no more new deliveries will be accepted.
+        /// </summary>
+        public abstract void Complete();
+
+        /// <summary>
+        /// Wait for any remaining queues deliveries to finish.
+        /// </summary>
+        public abstract void WaitForCompletion();
+
+        public static ConsumerDeliveryQueue Create(FakeModel model, bool blockingDelivery, Action<CallbackExceptionEventArgs> onDeliveryException)
+        {
+            return blockingDelivery
+                ? (ConsumerDeliveryQueue) new BlockingDeliveryQueue(model, onDeliveryException)
+                : new NonBlockingDeliveryQueue(model, onDeliveryException);
+        }
+    }
+
+    internal class NonBlockingDeliveryQueue : ConsumerDeliveryQueue
+    {
+        private readonly Task deliveriesTask;
+        private readonly AsyncLocal<bool> isDeliveriesTask = new AsyncLocal<bool>();
+        private readonly Channel<Action> deliveries = Channel.CreateUnbounded<Action>(new UnboundedChannelOptions
+        {
+            SingleReader = true,
+            SingleWriter = false,
+        });
+
+        public NonBlockingDeliveryQueue(FakeModel model, Action<CallbackExceptionEventArgs> onDeliveryException)
+            : base(model, onDeliveryException)
+        {
+            deliveriesTask = Task.Run(HandleDeliveries);
+        }
+
+
+        public override void Deliver(Action deliveryAction)
+        {
+            _ = deliveries.Writer.TryWrite(deliveryAction);
+        }
+
+        /// <summary>
+        /// Rabbit docs states that each connection is backed by a single background thread:
+        /// 
+        /// https://www.rabbitmq.com/dotnet-api-guide.html#concurrency-thread-usage
+        /// 
+        /// However, this is not actually true, it's backed by a Task:
+        /// 
+        /// https://github.com/rabbitmq/rabbitmq-dotnet-client/blob/65dd5f92dda130ec35b4ad6fe7bc54dbcb1637fd/projects/RabbitMQ.Client/client/impl/ConsumerWorkService.cs#L81
+        /// 
+        /// FakeModels aren't aware of their connection, so in order to emulate this, just
+        /// run a task that handles deliveries per task. It's necessary to match RabbitMQ
+        /// semantics as running delivery callbacks synchronously can cause deadlocks in
+        /// code under test.
+        /// </summary>
+        private async Task HandleDeliveries()
+        {
+            try
+            {
+                isDeliveriesTask.Value = true;
+                while (await deliveries.Reader.WaitToReadAsync().ConfigureAwait(false))
+                {
+                    while (deliveries.Reader.TryRead(out var delivery))
+                    {
+                        ExecuteDelivery(delivery);
+                    }
+                }
+            }
+            catch
+            {
+                // Swallow exceptions so FakeModel.Close() doesn't have to deal with it.
+            }
+        }
+
+        public override void Complete()
+        {
+            _ = deliveries.Writer.TryComplete();
+        }
+
+        public override void WaitForCompletion()
+        {
+            // It's possible that we can end up calling Close on a model from within the delivery handler.
+            // If this is the case, we must not wait on it to complete as this will deadlock!
+            if (!isDeliveriesTask.Value)
+                deliveriesTask.Wait();
+        }
+    }
+
+    internal class BlockingDeliveryQueue : ConsumerDeliveryQueue
+    {
+        private static readonly TimeSpan DeliveryWaitTimeout = TimeSpan.FromMinutes(1);
+
+        private readonly SemaphoreSlim deliveryLock = new SemaphoreSlim(1);
+
+        private volatile bool notAcceptingNewDeliveries = false;
+
+        public BlockingDeliveryQueue(FakeModel model, Action<CallbackExceptionEventArgs> onDeliveryException)
+            : base(model, onDeliveryException)
+        {
+        }
+
+        public override void Deliver(Action deliveryAction)
+        {
+            if (notAcceptingNewDeliveries)
+                return;
+
+            try
+            {
+                deliveryLock.Wait(DeliveryWaitTimeout);
+
+                ExecuteDelivery(deliveryAction);
+            }
+            finally
+            {
+                deliveryLock.Release();
+            }
+        }
+
+        public override void Complete()
+        {
+            notAcceptingNewDeliveries = true;
+        }
+
+        public override void WaitForCompletion()
+        {
+            try
+            {
+                deliveryLock.Wait(DeliveryWaitTimeout);
+            }
+            finally
+            {
+                deliveryLock.Release();
+            }
+        }
+    }
+}
diff --git a/src/AddUp.FakeRabbitMQ/FakeModel.cs b/src/AddUp.FakeRabbitMQ/FakeModel.cs
@@ -3,8 +3,6 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
-using System.Threading.Channels;
-using System.Threading.Tasks;
 using RabbitMQ.Client;
 using RabbitMQ.Client.Events;
 using RabbitMQ.Client.Exceptions;
@@ -15,20 +13,18 @@ internal sealed class FakeModel : IModel
     {
         private readonly ConcurrentDictionary<ulong, RabbitMessage> workingMessages = new ConcurrentDictionary<ulong, RabbitMessage>();
         private readonly ConcurrentDictionary<string, ConsumerData> consumers = new ConcurrentDictionary<string, ConsumerData>();
-        private readonly Channel<Action> deliveries = Channel.CreateUnbounded<Action>(new UnboundedChannelOptions
-        {
-            SingleReader = true,
-            SingleWriter = false,
-        });
+
+        private readonly ConsumerDeliveryQueue deliveryQueue;
         private readonly RabbitServer server;
-        private readonly Task deliveriesTask;
-        private readonly AsyncLocal<bool> isDeliveriesTask = new AsyncLocal<bool>();
         private long lastDeliveryTag;
 
         public FakeModel(RabbitServer rabbitServer)
         {
             server = rabbitServer;
-            deliveriesTask = Task.Run(HandleDeliveries);
+            deliveryQueue = ConsumerDeliveryQueue.Create(
+                this,
+                rabbitServer.BlockingConsumerDelivery,
+                onDeliveryException: args => CallbackException(this, args));
         }
 
 #pragma warning disable 67
@@ -127,7 +123,8 @@ void notifyConsumerOfMessage(RabbitMessage message)
             if (queueInstance != null)
             {
                 void publishedAction(object sender, RabbitMessage message) =>
-                    deliveries.Writer.TryWrite(() => notifyConsumerOfMessage(message));
+                    deliveryQueue.Deliver(() => notifyConsumerOfMessage(message));
+
                 var consumerData = new ConsumerData(consumer, queueInstance, publishedAction);
 
                 // https://www.rabbitmq.com/amqp-0-9-1-reference.html#basic.consume.consumer-tag
@@ -138,7 +135,8 @@ ConsumerData updateFunction(string s, ConsumerData _) =>
                 _ = consumers.AddOrUpdate(consumerTag, consumerData, updateFunction);
 
                 foreach (var message in queueInstance.Messages)
-                    _ = deliveries.Writer.TryWrite(() => notifyConsumerOfMessage(message));
+                    consumerData.QueueMessagePublished(this, message);
+
                 queueInstance.MessagePublished += consumerData.QueueMessagePublished;
 
                 if (consumer is IAsyncBasicConsumer asyncBasicConsumer)
@@ -279,7 +277,7 @@ private void Close(ushort replyCode, string replyText, bool abort)
                     foreach (var consumerTag in consumerTags)
                         BasicCancel(consumerTag);
 
-                    _ = deliveries.Writer.TryComplete();
+                    deliveryQueue.Complete();
                     ModelShutdown?.Invoke(this, reason);
                 }
                 catch
@@ -288,10 +286,7 @@ private void Close(ushort replyCode, string replyText, bool abort)
                 }
             }
 
-            // It's possible that we can end up calling Close on a model from within the delivery handler.
-            // If this is the case, we must not wait on it to complete as this will deadlock!
-            if (!isDeliveriesTask.Value)
-                deliveriesTask.Wait();
+            deliveryQueue.WaitForCompletion();
         }
 
         public void ConfirmSelect()
@@ -482,47 +477,5 @@ public bool WaitForConfirms(TimeSpan timeout, out bool timedOut)
 
         public void WaitForConfirmsOrDie() => WaitForConfirmsOrDie(Timeout.InfiniteTimeSpan);
         public void WaitForConfirmsOrDie(TimeSpan timeout) => _ = WaitForConfirms(timeout);
-
-        /// <summary>
-        /// Rabbit docs states that each connection is backed by a single background thread:
-        /// 
-        /// https://www.rabbitmq.com/dotnet-api-guide.html#concurrency-thread-usage
-        /// 
-        /// However, this is not actually true, it's backed by a Task:
-        /// 
-        /// https://github.com/rabbitmq/rabbitmq-dotnet-client/blob/65dd5f92dda130ec35b4ad6fe7bc54dbcb1637fd/projects/RabbitMQ.Client/client/impl/ConsumerWorkService.cs#L81
-        /// 
-        /// FakeModels aren't aware of their connection, so in order to emulate this, just
-        /// run a task that handles deliveries per task. It's necessary to match RabbitMQ
-        /// semantics as running delivery callbacks synchronously can cause deadlocks in
-        /// code under test.
-        /// </summary>
-        private async Task HandleDeliveries()
-        {
-            try
-            {
-                isDeliveriesTask.Value = true;
-                while (await deliveries.Reader.WaitToReadAsync().ConfigureAwait(false))
-                {
-                    while (deliveries.Reader.TryRead(out var delivery))
-                    {
-                        try
-                        {
-                            if (!IsOpen) break;
-                            delivery();
-                        }
-                        catch (Exception ex)
-                        {
-                            var callbackArgs = CallbackExceptionEventArgs.Build(ex, "");
-                            CallbackException(this, callbackArgs);
-                        }
-                    }
-                }
-            }
-            catch
-            {
-                // Swallow exceptions so Close() doesn't have to deal with it.
-            }
-        }
     }
 }
diff --git a/src/AddUp.FakeRabbitMQ/RabbitServer.cs b/src/AddUp.FakeRabbitMQ/RabbitServer.cs
@@ -23,6 +23,12 @@ public RabbitServer()
         public ConcurrentDictionary<string, RabbitExchange> Exchanges { get; }
         public ConcurrentDictionary<string, RabbitQueue> Queues { get; }
 
+        /// <summary>
+        /// If true, deliveries to consumers will execute in a blocking manner, meaning the publish will not
+        /// finish until the message has reached all registered consumers.
+        /// </summary>
+        public bool BlockingConsumerDelivery { get; set; }
+
         public void Reset()
         {
             Exchanges.Clear();
