implement AwaitableMutex for netfx

Author: mgravell

src/StackExchange.Redis/AwaitableMutex.net.cs

@@ -1,12 +1,17 @@
 using System.Threading;
 using System.Threading.Tasks;
 
-// #if NET
+#if NET
 namespace StackExchange.Redis;
 
 internal partial struct AwaitableMutex
 {
     private readonly int _timeoutMilliseconds;
+
+    // note: this does not guarantee "fairness", but that's OK for our use-case - we mostly just want
+    // a sync+async awaitable mutex, which this does; the .NET Framework version has a hand-written
+    // implementation (see .netfx.cx for reasons), which *is* fair, but we'd rather not pay that overhead
+    // here. Good-enough-is.
     private readonly SemaphoreSlim _mutex;
 
     private partial AwaitableMutex(int timeoutMilliseconds)
@@ -28,4 +33,4 @@ public partial ValueTask<bool> TryTakeAsync(CancellationToken cancellationToken)
 
     public partial void Release() => _mutex.Release();
 }
-// #endif
+#endif

src/StackExchange.Redis/AwaitableMutex.netfx.cs

@@ -1,9 +1,321 @@
-#if !NET
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+#if !NET
 namespace StackExchange.Redis;
 
-// compensating for the fact that netfx SemaphoreSlim is kinda janky
-// (https://blog.marcgravell.com/2019/02/fun-with-spiral-of-death.html)
+/*
+Compensating for the fact that netfx SemaphoreSlim is kinda janky (https://blog.marcgravell.com/2019/02/fun-with-spiral-of-death.html).
+
+This uses a simple queue of sync/async callers, and assumes a reasonable caller (the original MutexSlim is more defensive, as
+a general purpose public API).
+*/
+
 internal partial struct AwaitableMutex
 {
+    private readonly State _state;
+
+    private partial AwaitableMutex(int timeoutMilliseconds)
+    {
+        _state = new(timeoutMilliseconds);
+    }
+
+    public partial void Dispose() => _state?.Dispose();
+
+    public partial bool IsAvailable => _state.IsAvailable;
+    public partial int TimeoutMilliseconds => _state.TimeoutMilliseconds;
+
+    public partial bool TryTakeInstant() => _state.TryTakeInstant();
+
+    public partial ValueTask<bool> TryTakeAsync(CancellationToken cancellationToken)
+        => _state.TryTakeAsync(cancellationToken);
+
+    public partial bool TryTakeSync() => _state.TryTakeSync();
+
+    public partial void Release() => _state.Release();
+
+    private sealed class State : IDisposable
+    {
+        private readonly Queue<IPendingCaller> _queue = new();
+        private bool _isHeld;
+        private bool _isDisposed;
+
+        public State(int timeoutMilliseconds)
+        {
+            if (timeoutMilliseconds < Timeout.Infinite) ThrowOutOfRangeException();
+            TimeoutMilliseconds = timeoutMilliseconds;
+
+            static void ThrowOutOfRangeException() => throw new ArgumentOutOfRangeException(nameof(timeoutMilliseconds));
+        }
+
+        public int TimeoutMilliseconds { get; }
+
+        public bool IsAvailable
+        {
+            get
+            {
+                lock (_queue)
+                {
+                    return !_isDisposed && !_isHeld && _queue.Count == 0;
+                }
+            }
+        }
+
+        public bool TryTakeInstant()
+        {
+            lock (_queue)
+            {
+                ThrowIfDisposed();
+                return TryTakeInsideLock();
+            }
+        }
+
+        public bool TryTakeSync()
+        {
+            var start = GetTime();
+            SyncPendingCaller pending;
+            lock (_queue)
+            {
+                ThrowIfDisposed();
+                if (TryTakeInsideLock()) return true;
+                if (TimeoutMilliseconds == 0) return false;
+
+                pending = new(start, TimeoutMilliseconds);
+                _queue.Enqueue(pending);
+            }
+
+            return pending.Wait();
+        }
+
+        public ValueTask<bool> TryTakeAsync(CancellationToken cancellationToken)
+        {
+            lock (_queue)
+            {
+                ThrowIfDisposed();
+                if (cancellationToken.IsCancellationRequested) return AsyncPendingCaller.Canceled();
+                if (TryTakeInsideLock()) return new ValueTask<bool>(true);
+                if (TimeoutMilliseconds == 0) return new ValueTask<bool>(false);
+                if (cancellationToken.IsCancellationRequested) return AsyncPendingCaller.Canceled();
+
+                var pending = new AsyncPendingCaller(TimeoutMilliseconds, cancellationToken);
+                _queue.Enqueue(pending);
+                return new ValueTask<bool>(pending.Task);
+            }
+        }
+
+        public void Release()
+        {
+            lock (_queue)
+            {
+                ThrowIfDisposed();
+                if (!_isHeld) ThrowNotHeld();
+
+                while (_queue.Count != 0)
+                {
+                    if (_queue.Dequeue().TryGrant()) return;
+                }
+
+                _isHeld = false;
+            }
+
+            static void ThrowNotHeld() => throw new SemaphoreFullException();
+        }
+
+        private bool TryTakeInsideLock()
+        {
+            if (_isHeld || _queue.Count != 0) return false;
+            _isHeld = true;
+            return true;
+        }
+
+        public void Dispose()
+        {
+            lock (_queue)
+            {
+                if (_isDisposed) return;
+
+                _isDisposed = true;
+                _isHeld = false;
+                while (_queue.Count != 0)
+                {
+                    _queue.Dequeue().Abort();
+                }
+            }
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (_isDisposed) ThrowDisposed();
+        }
+
+        private static void ThrowDisposed() => throw new ObjectDisposedException(nameof(AwaitableMutex));
+
+        private static uint GetTime() => (uint)Environment.TickCount;
+
+        private static int GetRemainingTimeout(uint startTime, int originalTimeoutMilliseconds)
+        {
+            if (originalTimeoutMilliseconds == Timeout.Infinite) return Timeout.Infinite;
+
+            var elapsedMilliseconds = GetTime() - startTime;
+            if (elapsedMilliseconds > int.MaxValue) return 0;
+
+            var remaining = originalTimeoutMilliseconds - (int)elapsedMilliseconds;
+            return remaining <= 0 ? 0 : remaining;
+        }
+
+        private interface IPendingCaller
+        {
+            bool TryGrant();
+            void Abort();
+        }
+
+        private sealed class SyncPendingCaller : IPendingCaller
+        {
+            private readonly uint _start;
+            private readonly int _timeoutMilliseconds;
+            private bool _isComplete;
+            private bool _wasGranted;
+            private bool _wasAborted;
+
+            public SyncPendingCaller(uint start, int timeoutMilliseconds)
+            {
+                _start = start;
+                _timeoutMilliseconds = timeoutMilliseconds;
+            }
+
+            public bool Wait()
+            {
+                lock (this)
+                {
+                    while (!_isComplete)
+                    {
+                        var remaining = GetRemainingTimeout(_start, _timeoutMilliseconds);
+                        if (remaining == 0)
+                        {
+                            _isComplete = true;
+                            return false;
+                        }
+
+                        if (remaining == Timeout.Infinite)
+                        {
+                            Monitor.Wait(this);
+                        }
+                        else
+                        {
+                            Monitor.Wait(this, remaining);
+                        }
+                    }
+
+                    if (_wasAborted) ThrowDisposed();
+                    return _wasGranted;
+                }
+            }
+
+            public bool TryGrant()
+            {
+                lock (this)
+                {
+                    if (_isComplete) return false;
+                    _wasGranted = true;
+                    _isComplete = true;
+                    Monitor.Pulse(this);
+                    return true;
+                }
+            }
+
+            public void Abort()
+            {
+                lock (this)
+                {
+                    if (_isComplete) return;
+                    _wasAborted = true;
+                    _isComplete = true;
+                    Monitor.Pulse(this);
+                }
+            }
+        }
+
+        private sealed class AsyncPendingCaller : TaskCompletionSource<bool>, IPendingCaller
+        {
+            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 CancellationTokenRegistration _cancellation;
+            private readonly Timer? _timeout;
+            private int _completionState;
+
+            public AsyncPendingCaller(int timeoutMilliseconds, CancellationToken cancellationToken)
+                : base(TaskCreationOptions.RunContinuationsAsynchronously)
+            {
+                if (timeoutMilliseconds != Timeout.Infinite)
+                {
+                    _timeout = new Timer(s_onTimeout, this, timeoutMilliseconds, Timeout.Infinite);
+                }
+
+                if (cancellationToken.CanBeCanceled)
+                {
+                    _cancellation = cancellationToken.Register(s_onCanceled, this);
+                }
+            }
+
+            public static ValueTask<bool> Canceled() => new(s_canceledTask);
+
+            public bool TryGrant() => TryComplete(CompletionState.Granted);
+
+            public void Abort() => TryComplete(CompletionState.Disposed);
+
+            private bool TryComplete(CompletionState completionState)
+            {
+                var newState = (int)completionState;
+                if (Interlocked.CompareExchange(ref _completionState, newState, (int)CompletionState.Pending) != (int)CompletionState.Pending)
+                {
+                    return false;
+                }
+
+                if (completionState != CompletionState.TimedOut) _timeout?.Dispose();
+                if (completionState != CompletionState.Canceled) _cancellation.Dispose();
+                Complete(completionState);
+                return true;
+            }
+
+            private void Complete(CompletionState completionState)
+            {
+                switch (completionState)
+                {
+                    case CompletionState.Granted:
+                        TrySetResult(true);
+                        break;
+                    case CompletionState.TimedOut:
+                        TrySetResult(false);
+                        break;
+                    case CompletionState.Canceled:
+                        TrySetCanceled();
+                        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,
+                Granted = 1,
+                TimedOut = 2,
+                Canceled = 3,
+                Disposed = 4,
+            }
+        }
+    }
 }
 #endif