Update the unregister method to be async and await for inflight opera…

…tions to finish

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

@@ -66,7 +66,7 @@ public interface IReceiverClient : IClientEntity
         /// <summary>
         /// Unregister message handler from the receiver if there is an active message handler registered. 
         /// </summary>
-        void UnregisterMessageHandler();
+        Task UnregisterMessageHandler();
 
         /// <summary>
         /// Completes a <see cref="Message"/> using its lock token. This will delete the message from the queue.

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

@@ -902,23 +902,31 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
         /// <summary>
         /// Unregister message hander from the receiver if there is an active message handler registered. 
         /// </summary>
-        public void UnregisterMessageHandler()
+        public async Task UnregisterMessageHandler()
         {
             this.ThrowIfClosed();
 
             MessagingEventSource.Log.UnregisterMessageHandlerStart(this.ClientId);
             lock (this.messageReceivePumpSyncLock)
             {
-                if (this.receivePump != null)
-                {
-                    this.receivePumpCancellationTokenSource.Cancel();
-                    this.receivePumpCancellationTokenSource.Dispose();
-                    this.receivePump = null;
-                }
-                else
+                if (this.receivePump == null || this.receivePumpCancellationTokenSource.IsCancellationRequested)
                 {
-                    throw new InvalidOperationException(Resources.MessageHandlerNotRegisteredYet);
+                    return;
                 }
+
+                this.receivePumpCancellationTokenSource.Cancel();
+                this.receivePumpCancellationTokenSource.Dispose();
+            }
+
+            while (this.receivePump != null
+                && this.receivePump.maxConcurrentCallsSemaphoreSlim.CurrentCount < this.receivePump.registerHandlerOptions.MaxConcurrentCalls)
+            {
+                await Task.Delay(10);
+            }
+
+            lock (this.messageReceivePumpSyncLock)
+            {
+                this.receivePump = null;
             }
             MessagingEventSource.Log.UnregisterMessageHandlerStop(this.ClientId);
         }

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

@@ -12,14 +12,15 @@ 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 ServiceBusDiagnosticSource diagnosticSource;
 
+
         public MessageReceivePump(IMessageReceiver messageReceiver,
             MessageHandlerOptions registerHandlerOptions,
             Func<Message, CancellationToken, Task> callback,

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

@@ -448,10 +448,10 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
         /// <summary>
         /// Unregister messgae hander from the receiver if there is active message handler registered. 
         /// </summary>
-        public void UnregisterMessageHandler()
+        public async Task UnregisterMessageHandler()
         {
             this.ThrowIfClosed();
-            this.InnerReceiver.UnregisterMessageHandler();
+            await this.InnerReceiver.UnregisterMessageHandler();
         }
 
         /// <summary>

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

@@ -422,10 +422,10 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
         /// <summary>
         /// Unregister messgae hander from the receiver if there is active message handler registered. 
         /// </summary>
-        public void UnregisterMessageHandler()
+        public async Task UnregisterMessageHandler()
         {
             this.ThrowIfClosed();
-            this.InnerSubscriptionClient.InnerReceiver.UnregisterMessageHandler();
+            await this.InnerSubscriptionClient.InnerReceiver.UnregisterMessageHandler();
         }
 
         /// <summary>