Graceful stop of worker services

projects/client/RabbitMQ.Client/src/client/impl/AsyncConsumerDispatcher.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 
 namespace RabbitMQ.Client.Impl
 {
@@ -19,9 +20,9 @@ public void Quiesce()
             IsShutdown = true;
         }
 
-        public void Shutdown(IModel model)
+        public Task Shutdown(IModel model)
         {
-            _workService.Stop(model);
+            return _workService.Stop(model);
         }
 
         public bool IsShutdown

projects/client/RabbitMQ.Client/src/client/impl/AsyncConsumerWorkService.cs

@@ -20,20 +20,14 @@ private WorkPool StartNewWorkPool(IModel model)
             return newWorkPool;
         }
 
-        public void Stop(IModel model)
+        public Task Stop(IModel model)
         {
             if (_workPools.TryRemove(model, out WorkPool workPool))
             {
-                workPool.Stop();
+                return workPool.Stop();
             }
-        }
 
-        public void Stop()
-        {
-            foreach (IModel model in _workPools.Keys)
-            {
-                Stop(model);
-            }
+            return Task.CompletedTask;
         }
 
         class WorkPool
@@ -78,17 +72,18 @@ async Task Loop()
                         // Swallowing the task cancellation in case we are stopping work.
                     }
 
-                    while (_tokenSource.IsCancellationRequested == false && _workQueue.TryDequeue(out Work work))
+                    while (_workQueue.TryDequeue(out Work work))
                     {
                         await work.Execute(_model).ConfigureAwait(false);
                     }
                 }
             }
 
-            public void Stop()
+            public Task Stop()
             {
                 _tokenSource.Cancel();
                 _tokenRegistration.Dispose();
+                return _worker;
             }
         }
     }

projects/client/RabbitMQ.Client/src/client/impl/AutorecoveringModel.cs

@@ -279,7 +279,7 @@ public void Close(ShutdownEventArgs reason, bool abort)
         {
             try
             {
-                _delegate.Close(reason, abort);
+                _delegate.Close(reason, abort).GetAwaiter().GetResult();;
             }
             finally
             {

projects/client/RabbitMQ.Client/src/client/impl/ConcurrentConsumerDispatcher.cs

@@ -1,6 +1,6 @@
 using System;
 using System.Collections.Generic;
-
+using System.Threading.Tasks;
 using RabbitMQ.Client.Events;
 
 namespace RabbitMQ.Client.Impl
@@ -22,9 +22,9 @@ public void Quiesce()
             IsShutdown = true;
         }
 
-        public void Shutdown(IModel model)
+        public Task Shutdown(IModel model)
         {
-            _workService.StopWork(model);
+            return _workService.StopWorkAsync(model);
         }
 
         public bool IsShutdown

projects/client/RabbitMQ.Client/src/client/impl/ConsumerWorkService.cs

@@ -21,20 +21,27 @@ private WorkPool StartNewWorkPool(IModel model)
             return newWorkPool;
         }
 
-        public void StopWork(IModel model)
+        public void StopWork()
         {
-            if (_workPools.TryRemove(model, out WorkPool workPool))
+            foreach (IModel model in _workPools.Keys)
             {
-                workPool.Stop();
+                StopWork(model);
             }
         }
 
-        public void StopWork()
+        public void StopWork(IModel model)
         {
-            foreach (IModel model in _workPools.Keys)
+            StopWorkAsync(model).GetAwaiter().GetResult();
+        }
+
+        internal Task StopWorkAsync(IModel model)
+        {
+            if (_workPools.TryRemove(model, out WorkPool workPool))
             {
-                StopWork(model);
+                return workPool.Stop();
             }
+
+            return Task.CompletedTask;
         }
 
         class WorkPool
@@ -77,7 +84,7 @@ async Task Loop()
                         // Swallowing the task cancellation exception for the semaphore in case we are stopping.
                     }
 
-                    while (_tokenSource.IsCancellationRequested == false && _actions.TryDequeue(out Action action))
+                    while (_actions.TryDequeue(out Action action))
                     {
                         try
                         {
@@ -91,10 +98,11 @@ async Task Loop()
                 }
             }
 
-            public void Stop()
+            public Task Stop()
             {
                 _tokenSource.Cancel();
                 _tokenRegistration.Dispose();
+                return _worker;
             }
         }
     }

projects/client/RabbitMQ.Client/src/client/impl/IConsumerDispatcher.cs

@@ -39,6 +39,7 @@
 //---------------------------------------------------------------------------
 
 using System;
+using System.Threading.Tasks;
 
 namespace RabbitMQ.Client.Impl
 {
@@ -69,6 +70,6 @@ void HandleModelShutdown(IBasicConsumer consumer,
 
         void Quiesce();
 
-        void Shutdown(IModel model);
+        Task Shutdown(IModel model);
     }
 }

projects/client/RabbitMQ.Client/src/client/impl/ModelBase.cs

@@ -45,7 +45,7 @@
 using System.Runtime.CompilerServices;
 using System.Text;
 using System.Threading;
-
+using System.Threading.Tasks;
 using RabbitMQ.Client.Events;
 using RabbitMQ.Client.Exceptions;
 using RabbitMQ.Client.Framing;
@@ -184,14 +184,14 @@ public bool IsOpen
 
         public ISession Session { get; private set; }
 
-        public void Close(ushort replyCode, string replyText, bool abort)
+        public Task Close(ushort replyCode, string replyText, bool abort)
         {
-            Close(new ShutdownEventArgs(ShutdownInitiator.Application,
+            return Close(new ShutdownEventArgs(ShutdownInitiator.Application,
                 replyCode, replyText),
                 abort);
         }
 
-        public void Close(ShutdownEventArgs reason, bool abort)
+        public async Task Close(ShutdownEventArgs reason, bool abort)
         {
             var k = new ShutdownContinuation();
             ModelShutdown += k.OnConnectionShutdown;
@@ -204,7 +204,7 @@ public void Close(ShutdownEventArgs reason, bool abort)
                     _Private_ChannelClose(reason.ReplyCode, reason.ReplyText, 0, 0);
                 }
                 k.Wait(TimeSpan.FromMilliseconds(10000));
-                ConsumerDispatcher.Shutdown(this);
+                await ConsumerDispatcher.Shutdown(this).ConfigureAwait(false);
             }
             catch (AlreadyClosedException)
             {
@@ -508,7 +508,7 @@ public void OnSessionShutdown(object sender, ShutdownEventArgs reason)
             SetCloseReason(reason);
             OnModelShutdown(reason);
             BroadcastShutdownToConsumers(m_consumers, reason);
-            ConsumerDispatcher.Shutdown(this);
+            ConsumerDispatcher.Shutdown(this).GetAwaiter().GetResult();;
         }
 
         protected void BroadcastShutdownToConsumers(IDictionary<string, IBasicConsumer> cs, ShutdownEventArgs reason)
@@ -1394,7 +1394,7 @@ public void WaitForConfirmsOrDie(TimeSpan timeout)
                 Close(new ShutdownEventArgs(ShutdownInitiator.Application,
                     Constants.ReplySuccess,
                     "Nacks Received", new IOException("nack received")),
-                    false);
+                    false).GetAwaiter().GetResult();
                 throw new IOException("Nacks Received");
             }
             if (timedOut)
@@ -1403,7 +1403,7 @@ public void WaitForConfirmsOrDie(TimeSpan timeout)
                     Constants.ReplySuccess,
                     "Timed out waiting for acks",
                     new IOException("timed out waiting for acks")),
-                    false);
+                    false).GetAwaiter().GetResult();
                 throw new IOException("Timed out waiting for acks");
             }
         }