Improve async lock (#1520)

Source/MQTTnet/Implementations/PlatformAbstractionLayer.cs

@@ -11,11 +11,11 @@ namespace MQTTnet.Implementations
     public static class PlatformAbstractionLayer
     {
 #if NET452
-        public static Task CompletedTask => Task.FromResult(0);
+        public static Task CompletedTask { get; } = Task.FromResult(0);
 
         public static byte[] EmptyByteArray { get; } = new byte[0];
 #else
-        public static Task CompletedTask => Task.CompletedTask;
+        public static Task CompletedTask { get; } = Task.CompletedTask;
 
         public static byte[] EmptyByteArray { get; } = Array.Empty<byte>();
 #endif

Source/MQTTnet/Internal/AsyncLock.cs

@@ -3,6 +3,8 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -10,73 +12,173 @@ namespace MQTTnet.Internal
 {
     public sealed class AsyncLock : IDisposable
     {
-        readonly Task<IDisposable> _releaser;
+        /*
+         * This async supporting lock does not support reentrancy!
+         */
+
+        readonly List<Releaser> _queuedTasks = new List<Releaser>(64);
+        readonly Task<IDisposable> _releaserTaskWithDirectApproval;
+
+        readonly Releaser _releaserWithDirectApproval;
+
         readonly object _syncRoot = new object();
 
-        SemaphoreSlim _semaphore = new SemaphoreSlim(1, 1);
-
+        bool _isDisposed;
+        
         public AsyncLock()
         {
-            _releaser = Task.FromResult((IDisposable)new Releaser(this));
+            _releaserWithDirectApproval = new Releaser(this, null, CancellationToken.None);
+            _releaserTaskWithDirectApproval = Task.FromResult((IDisposable)_releaserWithDirectApproval);
         }
 
         public void Dispose()
         {
             lock (_syncRoot)
             {
-                _semaphore?.Dispose();
-                _semaphore = null;
+                foreach (var waitingTask in _queuedTasks)
+                {
+                    waitingTask.Fail(new ObjectDisposedException(nameof(AsyncLock)));
+                }
+
+                _queuedTasks.Clear();
+                
+                _isDisposed = true;
             }
         }
 
         public Task<IDisposable> WaitAsync(CancellationToken cancellationToken)
         {
-            Task task;
+            var hasDirectApproval = false;
+            Releaser releaser;
 
-            // This lock is required to avoid ObjectDisposedExceptions.
-            // These are fired when this lock gets disposed (and thus the semaphore)
-            // and a worker thread tries to call this method at the same time.
-            // Another way would be catching all ObjectDisposedExceptions but this situation happens
-            // quite often when clients are disconnecting.
             lock (_syncRoot)
             {
-                task = _semaphore?.WaitAsync(cancellationToken);
+                if (_isDisposed)
+                {
+                    throw new ObjectDisposedException(nameof(AsyncLock));
+                }
+                
+                if (_queuedTasks.Count == 0)
+                {
+                    // There is no other waiting task apart from the current one.
+                    // So we can approve the current task directly.
+                    releaser = _releaserWithDirectApproval;
+                    hasDirectApproval = true;
+                    Debug.WriteLine("AsyncLock: Task -1 directly approved.");
+                }
+                else
+                {
+                    releaser = new Releaser(this, new TaskCompletionSource<IDisposable>(), cancellationToken);
+                }
+
+                _queuedTasks.Add(releaser);
             }
 
-            if (task == null)
+            if (!hasDirectApproval)
             {
-                throw new ObjectDisposedException("The AsyncLock is disposed.");
+                return releaser.Task;
             }
 
-            if (task.Status == TaskStatus.RanToCompletion)
-            {
-                return _releaser;
-            }
-
-            // Wait for the _WaitAsync_ method and return the releaser afterwards.
-            return task.ContinueWith((_, state) => (IDisposable)state, _releaser.Result, cancellationToken, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+            return _releaserTaskWithDirectApproval;
         }
 
-        void Release()
+        void Release(Releaser releaser)
         {
             lock (_syncRoot)
             {
-                _semaphore?.Release();
+                if (_isDisposed)
+                {
+                    // There is no much left to do!
+                    return;
+                }
+
+                var activeTask = _queuedTasks[0];
+                if (!ReferenceEquals(activeTask, releaser))
+                {
+                    throw new InvalidOperationException("The active task must be the current releaser.");
+                }
+
+                _queuedTasks.RemoveAt(0);
+                
+                while (_queuedTasks.Count > 0)
+                {
+                    var nextTask = _queuedTasks[0];
+                    if (!nextTask.IsPending)
+                    {
+                        // Dequeue all canceled or failed tasks.
+                        _queuedTasks.RemoveAt(0);
+                        continue;
+                    }
+
+                    nextTask.Approve();
+                    return;
+                }
+
+                Debug.WriteLine("AsyncLock: No Task pending.");
             }
         }
 
         sealed class Releaser : IDisposable
         {
-            readonly AsyncLock _lock;
+            readonly AsyncLock _asyncLock;
+            readonly CancellationToken _cancellationToken;
+            readonly int _id;
+            readonly TaskCompletionSource<IDisposable> _promise;
 
-            internal Releaser(AsyncLock @lock)
+            // ReSharper disable once FieldCanBeMadeReadOnly.Local
+            CancellationTokenRegistration _cancellationTokenRegistration;
+
+            internal Releaser(AsyncLock asyncLock, TaskCompletionSource<IDisposable> promise, CancellationToken cancellationToken)
             {
-                _lock = @lock;
+                _asyncLock = asyncLock ?? throw new ArgumentNullException(nameof(asyncLock));
+                _promise = promise;
+                _cancellationToken = cancellationToken;
+
+                if (cancellationToken.CanBeCanceled)
+                {
+                    _cancellationTokenRegistration = cancellationToken.Register(Cancel);
+                }
+
+                _id = promise?.Task.Id ?? -1;
+
+                Debug.WriteLine($"AsyncLock: Task {_id} queued.");
+            }
+
+            public bool IsPending => _promise != null && !_promise.Task.IsCanceled && !_promise.Task.IsFaulted && !_promise.Task.IsCompleted;
+
+            public Task<IDisposable> Task => _promise?.Task;
+
+            public void Approve()
+            {
+                _promise?.TrySetResult(this);
+
+                Debug.WriteLine($"AsyncLock: Task {_id} approved.");
             }
 
             public void Dispose()
             {
-                _lock.Release();
+                if (_cancellationToken.CanBeCanceled)
+                {
+                    _cancellationTokenRegistration.Dispose();
+                }
+
+                Debug.WriteLine($"AsyncLock: Task {_id} completed.");
+
+                _asyncLock.Release(this);
+            }
+
+            public void Fail(Exception exception)
+            {
+                _promise?.TrySetException(exception);
+
+                Debug.WriteLine($"AsyncLock: Task {_id} failed ({exception.GetType().Name}).");
+            }
+
+            void Cancel()
+            {
+                _promise?.TrySetCanceled();
+
+                Debug.WriteLine($"AsyncLock: Task {_id} canceled.");
             }
         }
     }