optimize async path for immediate acquisition

Author: mgravell

src/StackExchange.Redis/AwaitableMutex.netfx.cs

@@ -97,18 +97,44 @@ public bool TryTakeSync()
         }
 
         public ValueTask<bool> TryTakeAsync(CancellationToken cancellationToken)
+        {
+            bool lockTaken = false;
+            try
+            {
+                // try to acquire uncontested lock - that way we can avoid allocating the pending caller
+                Monitor.TryEnter(_queue, 0, ref lockTaken);
+                if (lockTaken)
+                {
+                    if (_isDisposed) return DisposedAsync();
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        return CanceledAsync(cancellationToken);
+                    }
+
+                    if (TryTakeInsideLockCore()) return new ValueTask<bool>(true);
+                }
+            }
+            finally
+            {
+                if (lockTaken) Monitor.Exit(_queue);
+            }
+
+            return TryTakeAsyncSlow(cancellationToken);
+        }
+
+        private ValueTask<bool> TryTakeAsyncSlow(CancellationToken cancellationToken)
         {
             lock (_queue)
             {
+                if (_isDisposed) return DisposedAsync();
                 if (cancellationToken.IsCancellationRequested)
                 {
-                    ThrowIfDisposed();
-                    return AsyncPendingCaller.Canceled();
+                    return CanceledAsync(cancellationToken);
                 }
 
-                if (TryTakeInsideLock()) return new ValueTask<bool>(true);
+                if (TryTakeInsideLockCore()) return new ValueTask<bool>(true);
                 if (TimeoutMilliseconds == 0) return new ValueTask<bool>(false);
-                if (cancellationToken.IsCancellationRequested) return AsyncPendingCaller.Canceled();
+                if (cancellationToken.IsCancellationRequested) return CanceledAsync(cancellationToken);
 
                 var pending = new AsyncPendingCaller(TimeoutMilliseconds, cancellationToken);
                 _queue.Enqueue(pending);
@@ -137,6 +163,11 @@ public void Release()
         private bool TryTakeInsideLock()
         {
             ThrowIfDisposed();
+            return TryTakeInsideLockCore();
+        }
+
+        private bool TryTakeInsideLockCore()
+        {
             if (_isHeld || _queue.Count != 0) return false;
             _isHeld = true;
             return true;
@@ -191,6 +222,12 @@ private void ThrowIfDisposed()
 
         private static void ThrowDisposed() => throw new ObjectDisposedException(nameof(AwaitableMutex));
 
+        private static ValueTask<bool> DisposedAsync()
+            => new(Task.FromException<bool>(new ObjectDisposedException(nameof(AwaitableMutex))));
+
+        private static ValueTask<bool> CanceledAsync(CancellationToken cancellationToken)
+            => new(Task.FromCanceled<bool>(cancellationToken));
+
         private static uint GetTime() => (uint)Environment.TickCount;
 
         private static int GetRemainingTimeout(uint startTime, int originalTimeoutMilliseconds)
@@ -280,15 +317,16 @@ private sealed class AsyncPendingCaller : TaskCompletionSource<bool>, IPendingCa
         {
             private static readonly TimerCallback s_onTimeout = state => ((AsyncPendingCaller)state!).TryComplete(CompletionState.TimedOut);
             private static readonly Action<object?> s_onCanceled = state => ((AsyncPendingCaller)state!).TryComplete(CompletionState.Canceled);
-            private static readonly Task<bool> s_canceledTask = CreateCanceledTask();
 
+            private readonly CancellationToken _cancellationToken;
             private readonly CancellationTokenRegistration _cancellation;
             private readonly Timer? _timeout;
             private int _completionState;
 
             public AsyncPendingCaller(int timeoutMilliseconds, CancellationToken cancellationToken)
                 : base(TaskCreationOptions.RunContinuationsAsynchronously)
             {
+                _cancellationToken = cancellationToken;
                 if (timeoutMilliseconds != Timeout.Infinite)
                 {
                     _timeout = new Timer(s_onTimeout, this, timeoutMilliseconds, Timeout.Infinite);
@@ -300,8 +338,6 @@ public AsyncPendingCaller(int timeoutMilliseconds, CancellationToken cancellatio
                 }
             }
 
-            public static ValueTask<bool> Canceled() => new(s_canceledTask);
-
             public bool TryGrant() => TryComplete(CompletionState.Granted);
 
             public void Abort() => TryComplete(CompletionState.Disposed);
@@ -331,21 +367,14 @@ private void Complete(CompletionState completionState)
                         TrySetResult(false);
                         break;
                     case CompletionState.Canceled:
-                        TrySetCanceled();
+                        TrySetCanceled(_cancellationToken);
                         break;
                     case CompletionState.Disposed:
                         TrySetException(new ObjectDisposedException(nameof(AwaitableMutex)));
                         break;
                 }
             }
 
-            private static Task<bool> CreateCanceledTask()
-            {
-                var source = new TaskCompletionSource<bool>();
-                source.SetCanceled();
-                return source.Task;
-            }
-
             private enum CompletionState
             {
                 Pending = 0,

tests/StackExchange.Redis.Tests/AwaitableMutexTests.cs

@@ -55,7 +55,7 @@ public async Task AsyncCallerTimesOutWhileHeld()
     }
 
     [Fact]
-    public void DisposalPreventsNewAcquisitions()
+    public async Task DisposalPreventsNewAcquisitions()
     {
         var mutex = AwaitableMutex.Create(timeoutMilliseconds: 100);
         Assert.True(mutex.TryTakeInstant());
@@ -65,10 +65,7 @@ public void DisposalPreventsNewAcquisitions()
         Assert.False(mutex.IsAvailable);
         Assert.Throws<ObjectDisposedException>(() => mutex.TryTakeInstant());
         Assert.Throws<ObjectDisposedException>(() => mutex.TryTakeSync());
-        Assert.Throws<ObjectDisposedException>(() =>
-        {
-            _ = mutex.TryTakeAsync();
-        });
+        await Assert.ThrowsAsync<ObjectDisposedException>(async () => await mutex.TryTakeAsync().AsTask());
         Assert.Throws<ObjectDisposedException>(() => mutex.Release());
     }