Stephen Cleary's POC vs mutual exclusion

CSharp · Console

/Program.cs

using System;
using System.Threading.Tasks;
using System.Threading;
					
public class Program
{
	public static void Main()
	{
		Task.Run(async () =>
		{
			var asyncLock = new StephenCleary.RecursiveAsyncLock();
			for (var i = 0; i < 100000; ++i)
			{
				var raceCondition = 0;
				// You can acquire the lock asynchronously
				using (await asyncLock.LockAsync(CancellationToken.None))
				{
					await Task.WhenAll(
						Task.Run(async () =>
						{
							// The lock is reentrant
							using (await asyncLock.LockAsync(CancellationToken.None))
							{
								// The lock provides mutual exclusion
								raceCondition++;
							}
						}),
						Task.Run(async () =>
						{
							using (await asyncLock.LockAsync(CancellationToken.None))
							{
								raceCondition++;
							}
						})
					);
				}
				if (raceCondition != 2)
				{
					Console.WriteLine("Not okay");
					return;
				}
			}
			Console.WriteLine("Ok");
		}).Wait();
	}
}

/*
 * Copied with minimal changes from https://github.com/StephenCleary/AsyncEx/blob/v4/Source/Unit%20Tests/AdvancedExamples/RecursiveAsyncLockExample.cs
 */
namespace StephenCleary
{
    using System;
    using System.Collections.Immutable;
    using System.Runtime.CompilerServices;
    using System.Runtime.Remoting.Messaging;
    using System.Threading;
    using System.Threading.Tasks;
    using Nito.AsyncEx;

    // Use a logical call context to track which locks the current async logical stack frame "owns".
        // NOTE: This approach will *only* work on .NET 4.5!
        static class AsyncLockTracker
        {
            private static readonly string slotName = Guid.NewGuid().ToString("N");

            private static IImmutableDictionary<RecursiveAsyncLock, Tuple<int, Task<IDisposable>>> OwnedLocks
            {
                get
                {
                    var ret = CallContext.LogicalGetData(slotName) as ImmutableDictionary<RecursiveAsyncLock, Tuple<int, Task<IDisposable>>>;
                    if (ret == null)
                        return ImmutableDictionary.Create<RecursiveAsyncLock, Tuple<int, Task<IDisposable>>>();
                    return ret;
                }

                set
                {
                    CallContext.LogicalSetData(slotName, value);
                }
            }

            public static bool Contains(RecursiveAsyncLock mutex)
            {
                return OwnedLocks.ContainsKey(mutex);
            }

            public static Task<IDisposable> Lookup(RecursiveAsyncLock mutex)
            {
                Tuple<int, Task<IDisposable>> value;
                if (OwnedLocks.TryGetValue(mutex, out value))
                    return value.Item2;
                return null;
            }

            public static void Add(RecursiveAsyncLock mutex, Task<IDisposable> key)
            {
                Tuple<int, Task<IDisposable>> value;
                if (!OwnedLocks.TryGetValue(mutex, out value))
                    value = Tuple.Create(0, key);
                OwnedLocks = OwnedLocks.SetItem(mutex, Tuple.Create(value.Item1 + 1, value.Item2));
            }

            public static void Remove(RecursiveAsyncLock mutex)
            {
                var value = OwnedLocks[mutex];
                if (value.Item1 == 1)
                {
                    OwnedLocks = OwnedLocks.Remove(mutex);
                    value.Item2.Result.Dispose();
                }
                else
                {
                    OwnedLocks = OwnedLocks.SetItem(mutex, Tuple.Create(value.Item1 - 1, value.Item2));
                }
            }
        }

        public sealed class RecursiveAsyncLock
        {
            private readonly AsyncLock mutex;

            public RecursiveAsyncLock()
            {
                mutex = new AsyncLock();
            }

            public RecursiveAsyncLock(IAsyncWaitQueue<IDisposable> queue)
            {
                mutex = new AsyncLock(queue);
            }

            public int Id { get { return mutex.Id; } }

            public RecursiveLockAwaitable LockAsync(CancellationToken token)
            {
                var key = AsyncLockTracker.Lookup(this);
                if (key == null)
                    key = mutex.LockAsync(token).AsTask();
                return new RecursiveLockAwaitable(key, this);
            }

            public RecursiveLockAwaitable LockAsync()
            {
                return LockAsync(CancellationToken.None);
            }

            public sealed class RecursiveLockAwaitable : INotifyCompletion
            {
                private readonly Task<IDisposable> _key;
                private readonly TaskAwaiter<IDisposable> _awaiter;
                private readonly RecursiveAsyncLock _mutex;

                public RecursiveLockAwaitable(Task<IDisposable> key, RecursiveAsyncLock mutex)
                {
                    _key = key;
                    _awaiter = key.GetAwaiter();
                    _mutex = mutex;
                }

                public RecursiveLockAwaitable GetAwaiter()
                {
                    return this;
                }

                public bool IsCompleted
                {
                    get { return _awaiter.IsCompleted; }
                }

                public IDisposable GetResult()
                {
                    var ret = _awaiter.GetResult();
                    return new KeyDisposable(_key, _mutex);
                }

                public void OnCompleted(Action continuation)
                {
                    _awaiter.OnCompleted(continuation);
                }

                private sealed class KeyDisposable : IDisposable
                {
                    private RecursiveAsyncLock _mutex;

                    public KeyDisposable(Task<IDisposable> keyTask, RecursiveAsyncLock mutex)
                    {
                        _mutex = mutex;
                        AsyncLockTracker.Add(mutex, keyTask);
                    }

                    public void Dispose()
                    {
                        if (_mutex == null)
                            return;
                        AsyncLockTracker.Remove(_mutex);
                        _mutex = null;
                    }
                }
            }
        }
}