Enable a way to Unregister Message Handler and Session Handler

sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/IReceiverClient.cs

@@ -63,6 +63,13 @@ public interface IReceiverClient : IClientEntity
         /// <remarks>Enable prefetch to speed up the receive rate.</remarks>
         void RegisterMessageHandler(Func<Message, CancellationToken, Task> handler, MessageHandlerOptions messageHandlerOptions);
 
+        /// <summary>
+        /// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
+        /// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously 
+        /// registered. 
+        /// </summary>
+        Task UnregisterMessageHandlerAsync();
+
         /// <summary>
         /// Completes a <see cref="Message"/> using its lock token. This will delete the message from the queue.
         /// </summary>
@@ -115,4 +122,4 @@ public interface IReceiverClient : IClientEntity
         /// </remarks>
         Task DeadLetterAsync(string lockToken, string deadLetterReason, string deadLetterErrorDescription = null);
     }
-}
+}

sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/MessageReceiver.cs

@@ -54,7 +54,11 @@ public class MessageReceiver : ClientEntity, IMessageReceiver
         int prefetchCount;
         long lastPeekedSequenceNumber;
         MessageReceivePump receivePump;
+        // Cancellation token to cancel the message pump. Once this is fired, all future message handling operations registered by application will be
+        // cancelled. 
         CancellationTokenSource receivePumpCancellationTokenSource;
+        // Cancellation token to cancel the inflight message handling operations registered by application in the message pump. 
+        CancellationTokenSource runningTaskCancellationTokenSource;
 
         /// <summary>
         /// Creates a new MessageReceiver from a <see cref="ServiceBusConnectionStringBuilder"/>.
@@ -899,6 +903,41 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
             this.OnMessageHandler(messageHandlerOptions, handler);
         }
 
+        /// <summary>
+        /// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
+        /// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously 
+        /// registered. 
+        /// </summary>
+        public async Task UnregisterMessageHandlerAsync()
+        {
+            this.ThrowIfClosed();
+
+            MessagingEventSource.Log.UnregisterMessageHandlerStart(this.ClientId);
+            lock (this.messageReceivePumpSyncLock)
+            {
+                if (this.receivePump == null || this.receivePumpCancellationTokenSource.IsCancellationRequested)
+                {
+                    // Silently return if handler has already been unregistered. 
+                    return;
+                }
+
+                this.receivePumpCancellationTokenSource.Cancel();
+                this.receivePumpCancellationTokenSource.Dispose();
+            }
+
+            while (this.receivePump != null
+                && this.receivePump.maxConcurrentCallsSemaphoreSlim.CurrentCount < this.receivePump.registerHandlerOptions.MaxConcurrentCalls)
+            {
+                await Task.Delay(10).ConfigureAwait(false);
+            }
+
+            lock (this.messageReceivePumpSyncLock)
+            {
+                this.receivePump = null;
+            }
+            MessagingEventSource.Log.UnregisterMessageHandlerStop(this.ClientId);
+        }
+
         /// <summary>
         /// Registers a <see cref="ServiceBusPlugin"/> to be used with this receiver.
         /// </summary>
@@ -1003,6 +1042,9 @@ protected override async Task OnClosingAsync()
                 {
                     this.receivePumpCancellationTokenSource.Cancel();
                     this.receivePumpCancellationTokenSource.Dispose();
+                    // For back-compatibility 
+                    this.runningTaskCancellationTokenSource.Cancel();
+                    this.runningTaskCancellationTokenSource.Dispose();
                     this.receivePump = null;
                 }
             }
@@ -1279,7 +1321,14 @@ protected virtual void OnMessageHandler(
                 }
 
                 this.receivePumpCancellationTokenSource = new CancellationTokenSource();
-                this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint, this.receivePumpCancellationTokenSource.Token);
+
+                // Running task cancellation token source can be reused if previously UnregisterMessageHandlerAsync was called
+                if (this.runningTaskCancellationTokenSource == null)
+                {
+                    this.runningTaskCancellationTokenSource = new CancellationTokenSource();
+                }
+
+                this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint, this.receivePumpCancellationTokenSource.Token, this.runningTaskCancellationTokenSource.Token);
             }
 
             try
@@ -1295,6 +1344,8 @@ protected virtual void OnMessageHandler(
                     {
                         this.receivePumpCancellationTokenSource.Cancel();
                         this.receivePumpCancellationTokenSource.Dispose();
+                        //this.runningTaskCancellationTokenSource.Cancel();
+                        //this.runningTaskCancellationTokenSource.Dispose();
                         this.receivePump = null;
                     }
                 }

sdk/servicebus/Microsoft.Azure.ServiceBus/src/IQueueClient.cs

@@ -77,5 +77,12 @@ public interface IQueueClient : IReceiverClient, ISenderClient
         /// <param name="sessionHandlerOptions">Options used to configure the settings of the session pump.</param>
         /// <remarks>Enable prefetch to speed up the receive rate. </remarks>
         void RegisterSessionHandler(Func<IMessageSession, Message, CancellationToken, Task> handler, SessionHandlerOptions sessionHandlerOptions);
+
+        /// <summary>
+        /// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
+        /// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously 
+        /// registered. 
+        /// </summary>
+        Task UnregisterSessionHandlerAsync();
     }
 }

sdk/servicebus/Microsoft.Azure.ServiceBus/src/ISubscriptionClient.cs

@@ -114,5 +114,12 @@ public interface ISubscriptionClient : IReceiverClient
         /// <param name="sessionHandlerOptions">Options used to configure the settings of the session pump.</param>
         /// <remarks>Enable prefetch to speed up the receive rate. </remarks>
         void RegisterSessionHandler(Func<IMessageSession, Message, CancellationToken, Task> handler, SessionHandlerOptions sessionHandlerOptions);
+
+        /// <summary>
+        /// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
+        /// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously 
+        /// registered. 
+        /// </summary>
+        Task UnregisterSessionHandlerAsync();
     }
 }

sdk/servicebus/Microsoft.Azure.ServiceBus/src/MessageReceivePump.cs

@@ -12,25 +12,28 @@ namespace Microsoft.Azure.ServiceBus
 
     sealed class MessageReceivePump
     {
+        public readonly SemaphoreSlim maxConcurrentCallsSemaphoreSlim;
+        public readonly MessageHandlerOptions registerHandlerOptions;
         readonly Func<Message, CancellationToken, Task> onMessageCallback;
         readonly string endpoint;
-        readonly MessageHandlerOptions registerHandlerOptions;
         readonly IMessageReceiver messageReceiver;
         readonly CancellationToken pumpCancellationToken;
-        readonly SemaphoreSlim maxConcurrentCallsSemaphoreSlim;
+        readonly CancellationToken runningTaskCancellationToken;
         readonly ServiceBusDiagnosticSource diagnosticSource;
 
         public MessageReceivePump(IMessageReceiver messageReceiver,
             MessageHandlerOptions registerHandlerOptions,
             Func<Message, CancellationToken, Task> callback,
             Uri endpoint,
-            CancellationToken pumpCancellationToken)
+            CancellationToken pumpCancellationToken,
+            CancellationToken runningTaskCancellationToken)
         {
             this.messageReceiver = messageReceiver ?? throw new ArgumentNullException(nameof(messageReceiver));
             this.registerHandlerOptions = registerHandlerOptions;
             this.onMessageCallback = callback;
             this.endpoint = endpoint.Authority;
             this.pumpCancellationToken = pumpCancellationToken;
+            this.runningTaskCancellationToken = runningTaskCancellationToken;
             this.maxConcurrentCallsSemaphoreSlim = new SemaphoreSlim(this.registerHandlerOptions.MaxConcurrentCalls);
             this.diagnosticSource = new ServiceBusDiagnosticSource(messageReceiver.Path, endpoint);
         }
@@ -163,7 +166,7 @@ async Task MessageDispatchTask(Message message)
             try
             {
                 MessagingEventSource.Log.MessageReceiverPumpUserCallbackStart(this.messageReceiver.ClientId, message);
-                await this.onMessageCallback(message, this.pumpCancellationToken).ConfigureAwait(false);
+                await this.onMessageCallback(message, this.runningTaskCancellationToken).ConfigureAwait(false);
 
                 MessagingEventSource.Log.MessageReceiverPumpUserCallbackStop(this.messageReceiver.ClientId, message);
             }

sdk/servicebus/Microsoft.Azure.ServiceBus/src/MessagingEventSource.cs

@@ -1380,6 +1380,42 @@ public void ManagementSerializationException(string objectName, string details =
                 this.WriteEvent(117, objectName, details);
             }
         }
+
+        [Event(118, Level = EventLevel.Informational, Message = "{0}: Unregister MessageHandler start.")]
+        public void UnregisterMessageHandlerStart(string clientId)
+        {
+            if (this.IsEnabled())
+            {
+                this.WriteEvent(118, clientId);
+            }
+        }
+
+        [Event(119, Level = EventLevel.Informational, Message = "{0}: Unregister MessageHandler done.")]
+        public void UnregisterMessageHandlerStop(string clientId)
+        {
+            if (this.IsEnabled())
+            {
+                this.WriteEvent(119, clientId);
+            }
+        }
+
+        [Event(120, Level = EventLevel.Informational, Message = "{0}: Unregister SessionHandler start.")]
+        public void UnregisterSessionHandlerStart(string clientId)
+        {
+            if (this.IsEnabled())
+            {
+                this.WriteEvent(120, clientId);
+            }
+        }
+
+        [Event(121, Level = EventLevel.Informational, Message = "{0}: Unregister SessionHandler done.")]
+        public void UnregisterSessionHandlerStop(string clientId)
+        {
+            if (this.IsEnabled())
+            {
+                this.WriteEvent(121, clientId);
+            }
+        }
     }
 
     internal static class TraceHelper

sdk/servicebus/Microsoft.Azure.ServiceBus/src/QueueClient.cs

@@ -446,6 +446,17 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
             this.InnerReceiver.RegisterMessageHandler(handler, messageHandlerOptions);
         }
 
+        /// <summary>
+        /// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
+        /// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously 
+        /// registered. 
+        /// </summary>
+        public async Task UnregisterMessageHandlerAsync()
+        {
+            this.ThrowIfClosed();
+            await this.InnerReceiver.UnregisterMessageHandlerAsync().ConfigureAwait(false);
+        }
+
         /// <summary>
         /// Receive session messages continuously from the queue. Registers a message handler and begins a new thread to receive session-messages.
         /// This handler(<see cref="Func{IMessageSession, Message, CancellationToken, Task}"/>) is awaited on every time a new message is received by the queue client.
@@ -476,6 +487,17 @@ public void RegisterSessionHandler(Func<IMessageSession, Message, CancellationTo
             this.SessionPumpHost.OnSessionHandler(handler, sessionHandlerOptions);
         }
 
+        /// <summary>
+        /// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
+        /// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously 
+        /// registered. 
+        /// </summary>
+        public async Task UnregisterSessionHandlerAsync()
+        {
+            this.ThrowIfClosed();
+            await this.SessionPumpHost.UnregisterSessionHandlerAsync().ConfigureAwait(false);
+        }
+
         /// <summary>
         /// Schedules a message to appear on Service Bus at a later time.
         /// </summary>

sdk/servicebus/Microsoft.Azure.ServiceBus/src/SessionPumpHost.cs

@@ -12,6 +12,7 @@ internal sealed class SessionPumpHost
         readonly object syncLock;
         SessionReceivePump sessionReceivePump;
         CancellationTokenSource sessionPumpCancellationTokenSource;
+        CancellationTokenSource runningTaskCancellationTokenSource;
         readonly Uri endpoint;
 
         public SessionPumpHost(string clientId, ReceiveMode receiveMode, ISessionClient sessionClient, Uri endpoint)
@@ -35,6 +36,9 @@ public void Close()
             {
                 this.sessionPumpCancellationTokenSource?.Cancel();
                 this.sessionPumpCancellationTokenSource?.Dispose();
+                // For back-compatibility 
+                this.runningTaskCancellationTokenSource?.Cancel();
+                this.runningTaskCancellationTokenSource?.Dispose();
                 this.sessionReceivePump = null;
             }
         }
@@ -53,14 +57,22 @@ public void OnSessionHandler(
                 }
 
                 this.sessionPumpCancellationTokenSource = new CancellationTokenSource();
+
+                // Running task cancellation token source can be reused if previously UnregisterSessionHandlerAsync was called
+                if (this.runningTaskCancellationTokenSource == null)
+                {
+                    this.runningTaskCancellationTokenSource = new CancellationTokenSource();
+                }
+
                 this.sessionReceivePump = new SessionReceivePump(
                     this.ClientId,
                     this.SessionClient,
                     this.ReceiveMode,
                     sessionHandlerOptions,
                     callback,
                     this.endpoint,
-                    this.sessionPumpCancellationTokenSource.Token);
+                    this.sessionPumpCancellationTokenSource.Token,
+                    this.runningTaskCancellationTokenSource.Token);
             }
 
             try
@@ -82,5 +94,36 @@ public void OnSessionHandler(
 
             MessagingEventSource.Log.RegisterOnSessionHandlerStop(this.ClientId);
         }
+
+        public async Task UnregisterSessionHandlerAsync()
+        {
+            MessagingEventSource.Log.UnregisterSessionHandlerStart(this.ClientId);
+            lock (this.syncLock)
+            {
+                if (this.sessionReceivePump == null || this.sessionPumpCancellationTokenSource.IsCancellationRequested)
+                {
+                    // Silently return if handler has already been unregistered. 
+                    return;
+                }
+
+                this.sessionPumpCancellationTokenSource.Cancel();
+                this.sessionPumpCancellationTokenSource.Dispose();
+            }
+
+            while (this.sessionReceivePump != null
+                && (this.sessionReceivePump.maxConcurrentSessionsSemaphoreSlim.CurrentCount < 
+                    this.sessionReceivePump.sessionHandlerOptions.MaxConcurrentSessions
+                || this.sessionReceivePump.maxPendingAcceptSessionsSemaphoreSlim.CurrentCount <           
+                this.sessionReceivePump.sessionHandlerOptions.MaxConcurrentAcceptSessionCalls))
+            {
+                await Task.Delay(10).ConfigureAwait(false);
+            }
+
+            lock (this.sessionPumpCancellationTokenSource)
+            {
+                this.sessionReceivePump = null;
+            }
+            MessagingEventSource.Log.UnregisterSessionHandlerStop(this.ClientId);
+        }
     }
 }

sdk/servicebus/Microsoft.Azure.ServiceBus/src/SessionReceivePump.cs

@@ -11,15 +11,16 @@ namespace Microsoft.Azure.ServiceBus
 
     sealed class SessionReceivePump
     {
+        public readonly SemaphoreSlim maxConcurrentSessionsSemaphoreSlim;
+        public readonly SemaphoreSlim maxPendingAcceptSessionsSemaphoreSlim;
+        public readonly SessionHandlerOptions sessionHandlerOptions;
         readonly string clientId;
         readonly ISessionClient client;
         readonly Func<IMessageSession, Message, CancellationToken, Task> userOnSessionCallback;
-        readonly SessionHandlerOptions sessionHandlerOptions;
         readonly string endpoint;
         readonly string entityPath;
         readonly CancellationToken pumpCancellationToken;
-        readonly SemaphoreSlim maxConcurrentSessionsSemaphoreSlim;
-        readonly SemaphoreSlim maxPendingAcceptSessionsSemaphoreSlim;
+        readonly CancellationToken runningTaskCancellationToken;
         private readonly ServiceBusDiagnosticSource diagnosticSource;
 
         public SessionReceivePump(string clientId,
@@ -28,7 +29,8 @@ public SessionReceivePump(string clientId,
             SessionHandlerOptions sessionHandlerOptions,
             Func<IMessageSession, Message, CancellationToken, Task> callback,
             Uri endpoint,
-            CancellationToken token)
+            CancellationToken pumpToken,
+            CancellationToken runningTaskToken)
         {
             this.client = client ?? throw new ArgumentException(nameof(client));
             this.clientId = clientId;
@@ -37,7 +39,8 @@ public SessionReceivePump(string clientId,
             this.userOnSessionCallback = callback;
             this.endpoint = endpoint.Authority;
             this.entityPath = client.EntityPath;
-            this.pumpCancellationToken = token;
+            this.pumpCancellationToken = pumpToken;
+            this.runningTaskCancellationToken = runningTaskToken;
             this.maxConcurrentSessionsSemaphoreSlim = new SemaphoreSlim(this.sessionHandlerOptions.MaxConcurrentSessions);
             this.maxPendingAcceptSessionsSemaphoreSlim = new SemaphoreSlim(this.sessionHandlerOptions.MaxConcurrentAcceptSessionCalls);
             this.diagnosticSource = new ServiceBusDiagnosticSource(client.EntityPath, endpoint);
@@ -229,7 +232,7 @@ async Task MessagePumpTaskAsync(IMessageSession session)
                         var callbackExceptionOccurred = false;
                         try
                         {
-                            processTask = this.userOnSessionCallback(session, message, this.pumpCancellationToken);
+                            processTask = this.userOnSessionCallback(session, message, this.runningTaskCancellationToken);
                             await processTask.ConfigureAwait(false);
                         }
                         catch (Exception exception)