1552 mqttnetservermqttclientsessionsmanager subscribersessions is not proper removed on special case (#1562)

Source/MQTTnet/Internal/AsyncQueue.cs

@@ -11,64 +11,39 @@ namespace MQTTnet.Internal
 {
     public sealed class AsyncQueue<TItem> : IDisposable
     {
+        readonly AsyncSignal _signal = new AsyncSignal();
         readonly object _syncRoot = new object();
-        SemaphoreSlim _semaphore = new SemaphoreSlim(0);
+
         ConcurrentQueue<TItem> _queue = new ConcurrentQueue<TItem>();
 
+        bool _isDisposed;
+        
         public int Count => _queue.Count;
 
+        public void Clear()
+        {
+            Interlocked.Exchange(ref _queue, new ConcurrentQueue<TItem>());
+        }
+
+        public void Dispose()
+        {
+            lock (_syncRoot)
+            {
+                _signal?.Dispose();
+
+                _isDisposed = true;
+            }
+        }
+
         public void Enqueue(TItem item)
         {
             lock (_syncRoot)
             {
                 _queue.Enqueue(item);
-                _semaphore?.Release();
+                _signal.Set();
             }
         }
 
-        public async Task<AsyncQueueDequeueResult<TItem>> TryDequeueAsync(CancellationToken cancellationToken)
-        {
-            while (!cancellationToken.IsCancellationRequested)
-            {
-                try
-                {
-                    Task task;
-                    lock (_syncRoot)
-                    {
-                        if (_semaphore == null)
-                        {
-                            return new AsyncQueueDequeueResult<TItem>(false, default);
-                        }
-
-                        task = _semaphore.WaitAsync(cancellationToken);
-                    }
-                    
-                    await task.ConfigureAwait(false);
-
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        return new AsyncQueueDequeueResult<TItem>(false, default);
-                    }
-
-                    if (_queue.TryDequeue(out var item))
-                    {
-                        return new AsyncQueueDequeueResult<TItem>(true, item);
-                    }
-                }
-                catch (ArgumentNullException)
-                {
-                    // The semaphore throws this internally sometimes.
-                    return new AsyncQueueDequeueResult<TItem>(false, default);
-                }
-                catch (OperationCanceledException)
-                {
-                    return new AsyncQueueDequeueResult<TItem>(false, default);
-                }
-            }
-
-            return new AsyncQueueDequeueResult<TItem>(false, default);
-        }
-
         public AsyncQueueDequeueResult<TItem> TryDequeue()
         {
             if (_queue.TryDequeue(out var item))
@@ -79,18 +54,42 @@ namespace MQTTnet.Internal
             return new AsyncQueueDequeueResult<TItem>(false, default);
         }
 
-        public void Clear()
+        public async Task<AsyncQueueDequeueResult<TItem>> TryDequeueAsync(CancellationToken cancellationToken)
         {
-            Interlocked.Exchange(ref _queue, new ConcurrentQueue<TItem>());
-        }
-
-        public void Dispose()
-        {
-            lock (_syncRoot)
+            while (!cancellationToken.IsCancellationRequested)
             {
-                _semaphore?.Dispose();
-                _semaphore = null;
+                try
+                {
+                    Task task;
+                    lock (_syncRoot)
+                    {
+                        if (_isDisposed)
+                        {
+                            return new AsyncQueueDequeueResult<TItem>(false, default);
+                        }
+
+                        task = _signal.WaitAsync(cancellationToken);
+                    }
+
+                    await task.ConfigureAwait(false);
+
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        return new AsyncQueueDequeueResult<TItem>(false, default);
+                    }
+
+                    if (_queue.TryDequeue(out var item))
+                    {
+                        return new AsyncQueueDequeueResult<TItem>(true, item);
+                    }
+                }
+                catch (OperationCanceledException)
+                {
+                    return new AsyncQueueDequeueResult<TItem>(false, default);
+                }
             }
+
+            return new AsyncQueueDequeueResult<TItem>(false, default);
         }
     }
-}
+}

Source/MQTTnet/Internal/AsyncSignal.cs

@@ -0,0 +1,113 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace MQTTnet.Internal
+{
+    public sealed class AsyncSignal : IDisposable
+    {
+        readonly object _syncRoot = new object();
+
+        bool _isDisposed;
+        bool _isSignaled;
+        AsyncSignalWaiter _waiter;
+
+        public void Dispose()
+        {
+            lock (_syncRoot)
+            {
+                _waiter?.Dispose();
+                _waiter = null;
+
+                _isDisposed = true;
+            }
+        }
+
+        public void Set()
+        {
+            lock (_syncRoot)
+            {
+                _isSignaled = true;
+
+                // If there is already a waiting task let it run.
+                _waiter?.Approve();
+                _waiter = null;
+            }
+        }
+
+        public Task WaitAsync(CancellationToken cancellationToken = default)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            lock (_syncRoot)
+            {
+                ThrowIfDisposed();
+
+                if (_isSignaled)
+                {
+                    return CompletedTask.Instance;
+                }
+
+                if (_waiter != null)
+                {
+                    if (!_waiter.Task.IsCompleted)
+                    {
+                        throw new InvalidOperationException("Only one waiting task is permitted per async signal.");
+                    }
+                    
+                    _waiter.Dispose();
+                }
+
+                _waiter = new AsyncSignalWaiter(cancellationToken);
+                return _waiter.Task;
+            }
+        }
+
+        void ThrowIfDisposed()
+        {
+            if (_isDisposed)
+            {
+                throw new ObjectDisposedException(nameof(AsyncSignal));
+            }
+        }
+
+        sealed class AsyncSignalWaiter : IDisposable
+        {
+            readonly AsyncTaskCompletionSource<bool> _promise = new AsyncTaskCompletionSource<bool>();
+
+            // ReSharper disable once FieldCanBeMadeReadOnly.Local
+            CancellationTokenRegistration _cancellationTokenRegistration;
+
+            public AsyncSignalWaiter(CancellationToken cancellationToken)
+            {
+                if (cancellationToken.CanBeCanceled)
+                {
+                    _cancellationTokenRegistration = cancellationToken.Register(Cancel);
+                }
+            }
+
+            public Task Task => _promise.Task;
+
+            public void Approve()
+            {
+                _promise.TrySetResult(true);
+            }
+
+            public void Dispose()
+            {
+                _cancellationTokenRegistration.Dispose();
+
+                _promise.TrySetException(new ObjectDisposedException(nameof(AsyncSignalWaiter)));
+            }
+
+            void Cancel()
+            {
+                _promise.TrySetCanceled();
+            }
+        }
+    }
+}

Source/MQTTnet/Internal/MqttPacketBus.cs

@@ -20,7 +20,7 @@ namespace MQTTnet.Internal
             new LinkedList<MqttPacketBusItem>()
         };
 
-        readonly SemaphoreSlim _semaphore = new SemaphoreSlim(0, 1);
+        readonly AsyncSignal _signal = new AsyncSignal();
         readonly object _syncRoot = new object();
 
         int _activePartition = (int)MqttPacketBusPartition.Health;
@@ -82,7 +82,16 @@ namespace MQTTnet.Internal
 
                 // No partition contains data so that we have to wait and put
                 // the worker back to the thread pool.
-                await _semaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+                try
+                {
+                    await _signal.WaitAsync(cancellationToken).ConfigureAwait(false);
+                }
+                catch (ObjectDisposedException)
+                {
+                    // The cancelled token should "hide" the disposal of the signal.
+                    cancellationToken.ThrowIfCancellationRequested();
+                    throw;
+                }
             }
 
             cancellationToken.ThrowIfCancellationRequested();
@@ -92,7 +101,7 @@ namespace MQTTnet.Internal
 
         public void Dispose()
         {
-            _semaphore?.Dispose();
+            _signal.Dispose();
         }
 
         public void DropFirstItem(MqttPacketBusPartition partition)
@@ -118,11 +127,7 @@ namespace MQTTnet.Internal
             lock (_syncRoot)
             {
                 _partitions[(int)partition].AddLast(item);
-
-                if (_semaphore.CurrentCount == 0)
-                {
-                    _semaphore.Release();
-                }
+                _signal.Set();
             }
         }