Ensure lockfree public async methods enforce asynchronous continuations.

StephenCleary · Jan 1, 2024 · 8f642df · 8f642df
1 parent a27d7a7
commit 8f642df
Show file tree

Hide file tree

Showing 7 changed files with 70 additions and 24 deletions.
diff --git a/src/Nito.AsyncEx.Coordination/AsyncAutoResetEvent.cs b/src/Nito.AsyncEx.Coordination/AsyncAutoResetEvent.cs
@@ -47,17 +47,23 @@ public AsyncAutoResetEvent()
         /// Asynchronously waits for this event to be set. If the event is set, this method will auto-reset it and return immediately, even if the cancellation token is already signalled. If the wait is canceled, then it will not auto-reset this event.
         /// </summary>
         /// <param name="cancellationToken">The cancellation token used to cancel this wait.</param>
-        public Task WaitAsync(CancellationToken cancellationToken)
+        internal Task InternalWaitAsync(CancellationToken cancellationToken)
         {
 	        Task<object>? result = null;
 	        InterlockedState.Transform(ref _state, s => s switch
 	        {
 		        { IsSet: true } => new State(false, s.Queue),
-                _ => new State(false, s.Queue.Enqueue(ApplyCancel, cancellationToken, out result)),
+		        _ => new State(false, s.Queue.Enqueue(ApplyCancel, cancellationToken, out result)),
 	        });
-            return result ?? Task.CompletedTask;
+	        return result ?? Task.CompletedTask;
         }
 
+		/// <summary>
+		/// Asynchronously waits for this event to be set. If the event is set, this method will auto-reset it and return immediately, even if the cancellation token is already signalled. If the wait is canceled, then it will not auto-reset this event.
+		/// </summary>
+		/// <param name="cancellationToken">The cancellation token used to cancel this wait.</param>
+		public Task WaitAsync(CancellationToken cancellationToken) => AsyncUtility.ForceAsync(InternalWaitAsync(cancellationToken));
+
         /// <summary>
         /// Asynchronously waits for this event to be set. If the event is set, this method will auto-reset it and return immediately.
         /// </summary>
@@ -72,7 +78,7 @@ public Task WaitAsync()
         /// <param name="cancellationToken">The cancellation token used to cancel this wait.</param>
         public void Wait(CancellationToken cancellationToken)
         {
-            WaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
+	        InternalWaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
         }
 
         /// <summary>
@@ -140,11 +146,11 @@ public DebugView(AsyncAutoResetEvent are)
                 _are = are;
             }
 
-            public int Id { get { return _are.Id; } }
+            public int Id => _are.Id;
 
-            public bool IsSet { get { return _are._state.IsSet; } }
+            public bool IsSet => _are._state.IsSet;
 
-            public IAsyncWaitQueue<object> WaitQueue { get { return _are._state.Queue; } }
+            public IAsyncWaitQueue<object> WaitQueue => _are._state.Queue;
         }
         // ReSharper restore UnusedMember.Local
     }

diff --git a/src/Nito.AsyncEx.Coordination/AsyncConditionVariable.cs b/src/Nito.AsyncEx.Coordination/AsyncConditionVariable.cs
@@ -57,7 +57,7 @@ public void NotifyAll()
         /// Asynchronously waits for a signal on this condition variable. The associated lock MUST be held when calling this method, and it will still be held when this method returns, even if the method is cancelled.
         /// </summary>
         /// <param name="cancellationToken">The cancellation signal used to cancel this wait.</param>
-        public Task WaitAsync(CancellationToken cancellationToken)
+        internal Task InternalWaitAsync(CancellationToken cancellationToken)
         {
 	        Task<object> task = null!;
 
@@ -88,6 +88,12 @@ static async Task WaitAndRetakeLockAsync(Task task, AsyncLock asyncLock)
             }
         }
 
+        /// <summary>
+        /// Asynchronously waits for a signal on this condition variable. The associated lock MUST be held when calling this method, and it will still be held when this method returns, even if the method is cancelled.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation signal used to cancel this wait.</param>
+        public Task WaitAsync(CancellationToken cancellationToken) => AsyncUtility.ForceAsync(InternalWaitAsync(cancellationToken));
+
 		/// <summary>
 		/// Asynchronously waits for a signal on this condition variable. The associated lock MUST be held when calling this method, and it will still be held when the returned task completes.
 		/// </summary>
@@ -102,7 +108,7 @@ public Task WaitAsync()
         /// <param name="cancellationToken">The cancellation signal used to cancel this wait.</param>
         public void Wait(CancellationToken cancellationToken)
         {
-            WaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
+	        InternalWaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
         }
 
         /// <summary>

diff --git a/src/Nito.AsyncEx.Coordination/AsyncLock.cs b/src/Nito.AsyncEx.Coordination/AsyncLock.cs
@@ -65,7 +65,7 @@ public AsyncLock()
         /// </summary>
         /// <param name="cancellationToken">The cancellation token used to cancel the lock. If this is already set, then this method will attempt to take the lock immediately (succeeding if the lock is currently available).</param>
         /// <returns>A disposable that releases the lock when disposed.</returns>
-        private Task<IDisposable> RequestLockAsync(CancellationToken cancellationToken)
+        internal Task<IDisposable> InternalLockAsync(CancellationToken cancellationToken)
         {
             Task<IDisposable>? result = null;
 	        InterlockedState.Transform(ref _state, s => s switch
@@ -85,7 +85,7 @@ private Task<IDisposable> RequestLockAsync(CancellationToken cancellationToken)
 		/// <returns>A disposable that releases the lock when disposed.</returns>
 		public AwaitableDisposable<IDisposable> LockAsync(CancellationToken cancellationToken)
         {
-            return new AwaitableDisposable<IDisposable>(RequestLockAsync(cancellationToken));
+            return new AwaitableDisposable<IDisposable>(AsyncUtility.ForceAsync(InternalLockAsync(cancellationToken)));
         }
 
         /// <summary>
@@ -103,7 +103,7 @@ public AwaitableDisposable<IDisposable> LockAsync()
         /// <param name="cancellationToken">The cancellation token used to cancel the lock. If this is already set, then this method will attempt to take the lock immediately (succeeding if the lock is currently available).</param>
         public IDisposable Lock(CancellationToken cancellationToken)
         {
-            return RequestLockAsync(cancellationToken).WaitAndUnwrapException();
+            return InternalLockAsync(cancellationToken).WaitAndUnwrapException();
         }
 
         /// <summary>

diff --git a/src/Nito.AsyncEx.Coordination/AsyncReaderWriterLock.cs b/src/Nito.AsyncEx.Coordination/AsyncReaderWriterLock.cs
@@ -98,7 +98,7 @@ private void ReleaseWaitersWhenCanceled(Task task)
         /// </summary>
         /// <param name="cancellationToken">The cancellation token used to cancel the lock. If this is already set, then this method will attempt to take the lock immediately (succeeding if the lock is currently available).</param>
         /// <returns>A disposable that releases the lock when disposed.</returns>
-        private Task<IDisposable> RequestReaderLockAsync(CancellationToken cancellationToken)
+        internal Task<IDisposable> InternalReaderLockAsync(CancellationToken cancellationToken)
         {
 	        Task<IDisposable>? task = null;
             InterlockedState.Transform(ref _state, s => s switch
@@ -119,7 +119,7 @@ private Task<IDisposable> RequestReaderLockAsync(CancellationToken cancellationT
 		/// <returns>A disposable that releases the lock when disposed.</returns>
 		public AwaitableDisposable<IDisposable> ReaderLockAsync(CancellationToken cancellationToken)
         {
-            return new AwaitableDisposable<IDisposable>(RequestReaderLockAsync(cancellationToken));
+            return new AwaitableDisposable<IDisposable>(AsyncUtility.ForceAsync(InternalReaderLockAsync(cancellationToken)));
         }
 
         /// <summary>
@@ -138,7 +138,7 @@ public AwaitableDisposable<IDisposable> ReaderLockAsync()
         /// <returns>A disposable that releases the lock when disposed.</returns>
         public IDisposable ReaderLock(CancellationToken cancellationToken)
         {
-            return RequestReaderLockAsync(cancellationToken).WaitAndUnwrapException();
+            return InternalReaderLockAsync(cancellationToken).WaitAndUnwrapException();
         }
 
         /// <summary>
@@ -155,7 +155,7 @@ public IDisposable ReaderLock()
         /// </summary>
         /// <param name="cancellationToken">The cancellation token used to cancel the lock. If this is already set, then this method will attempt to take the lock immediately (succeeding if the lock is currently available).</param>
         /// <returns>A disposable that releases the lock when disposed.</returns>
-        private Task<IDisposable> RequestWriterLockAsync(CancellationToken cancellationToken)
+        internal Task<IDisposable> InternalWriterLockAsync(CancellationToken cancellationToken)
         {
 	        Task<IDisposable>? task = null;
             InterlockedState.Transform(ref _state, s => s switch
@@ -178,7 +178,7 @@ private Task<IDisposable> RequestWriterLockAsync(CancellationToken cancellationT
         /// <returns>A disposable that releases the lock when disposed.</returns>
         public AwaitableDisposable<IDisposable> WriterLockAsync(CancellationToken cancellationToken)
         {
-            return new AwaitableDisposable<IDisposable>(RequestWriterLockAsync(cancellationToken));
+            return new AwaitableDisposable<IDisposable>(AsyncUtility.ForceAsync(InternalWriterLockAsync(cancellationToken)));
         }
 
         /// <summary>
@@ -197,7 +197,7 @@ public AwaitableDisposable<IDisposable> WriterLockAsync()
         /// <returns>A disposable that releases the lock when disposed.</returns>
         public IDisposable WriterLock(CancellationToken cancellationToken)
         {
-            return RequestWriterLockAsync(cancellationToken).WaitAndUnwrapException();
+            return InternalWriterLockAsync(cancellationToken).WaitAndUnwrapException();
         }
 
         /// <summary>

diff --git a/src/Nito.AsyncEx.Coordination/AsyncSemaphore.cs b/src/Nito.AsyncEx.Coordination/AsyncSemaphore.cs
@@ -36,7 +36,7 @@ public sealed class AsyncSemaphore
         /// Asynchronously waits for a slot in the semaphore to be available.
         /// </summary>
         /// <param name="cancellationToken">The cancellation token used to cancel the wait. If this is already set, then this method will attempt to take the slot immediately (succeeding if a slot is currently available).</param>
-        public Task WaitAsync(CancellationToken cancellationToken)
+        internal Task InternalWaitAsync(CancellationToken cancellationToken)
         {
 	        Task<object>? result = null;
 	        InterlockedState.Transform(ref _state, s => s switch
@@ -47,6 +47,12 @@ public Task WaitAsync(CancellationToken cancellationToken)
 	        return result ?? TaskConstants.Completed;
         }
 
+        /// <summary>
+        /// Asynchronously waits for a slot in the semaphore to be available.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation token used to cancel the wait. If this is already set, then this method will attempt to take the slot immediately (succeeding if a slot is currently available).</param>
+        public Task WaitAsync(CancellationToken cancellationToken) => AsyncUtility.ForceAsync(InternalWaitAsync(cancellationToken));
+
 		/// <summary>
 		/// Asynchronously waits for a slot in the semaphore to be available.
 		/// </summary>
@@ -61,7 +67,7 @@ public Task WaitAsync()
         /// <param name="cancellationToken">The cancellation token used to cancel the wait. If this is already set, then this method will attempt to take the slot immediately (succeeding if a slot is currently available).</param>
         public void Wait(CancellationToken cancellationToken)
         {
-            WaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
+            InternalWaitAsync(cancellationToken).WaitAndUnwrapException(CancellationToken.None);
         }
 
         /// <summary>

diff --git a/src/Nito.AsyncEx.Coordination/Internals/AsyncUtility.cs b/src/Nito.AsyncEx.Coordination/Internals/AsyncUtility.cs
@@ -0,0 +1,32 @@
+using System;
+using System.Threading.Tasks;
+
+namespace Nito.AsyncEx.Internals;
+
+internal static class AsyncUtility
+{
+	public static Task ForceAsync(Task task)
+	{
+		_ = task ?? throw new ArgumentNullException(nameof(task));
+		return task.IsCompleted ? task : ForceAsyncCore(task);
+
+		static async Task ForceAsyncCore(Task task)
+		{
+			await task.ConfigureAwait(false);
+			await Task.Yield();
+		}
+	}
+
+	public static Task<TResult> ForceAsync<TResult>(Task<TResult> task)
+	{
+		_ = task ?? throw new ArgumentNullException(nameof(task));
+		return task.IsCompleted ? task : ForceAsyncCore(task);
+
+		static async Task<TResult> ForceAsyncCore(Task<TResult> task)
+		{
+			var result = await task.ConfigureAwait(false);
+			await Task.Yield();
+			return result;
+		}
+	}
+}
diff --git a/src/Nito.AsyncEx.Coordination/Nito.AsyncEx.Coordination.csproj b/src/Nito.AsyncEx.Coordination/Nito.AsyncEx.Coordination.csproj
@@ -6,10 +6,6 @@
     <PackageTags>$(PackageTags);asynclock;asynclazy</PackageTags>
   </PropertyGroup>
 
-  <ItemGroup>
-    <PackageReference Include="Nito.Collections.Deque" Version="1.1.1" />
-  </ItemGroup>
-
   <ItemGroup>
     <ProjectReference Include="..\Nito.AsyncEx.Tasks\Nito.AsyncEx.Tasks.csproj" />
   </ItemGroup>