Reduce overhead of SemaphoreSlim.WaitAsync (#24687)

This affects the case where SemaphoreSlim.WaitAsync is invoked, the operation completes asynchronously because there's no count available in the semaphore, and where a timeout and/or a cancellation token is passed to WaitAsync.  In that case, we need to wait for the underlying wait task to complete, for cancellation to be requested, or for the timeout to elapse.

There's currently one code path that handles this.  This change improves that slightly by avoiding a defensive array copy.  However, we can do much better when we get a cancelable token but no timeout, in which case we can avoid operations like Task.Delay, can register with the original cancellation token rather than a new one (which means we avoid forcing a new CancellationTokenSource's lazily-allocated partitions into existence), etc.

diff --git a/src/System.Private.CoreLib/shared/System/Threading/SemaphoreSlim.cs b/src/System.Private.CoreLib/shared/System/Threading/SemaphoreSlim.cs
index 6437966..0214456 100644 (file)
--- a/src/System.Private.CoreLib/shared/System/Threading/SemaphoreSlim.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/SemaphoreSlim.cs
@@ -712,19 +712,31 @@ namespace System.Threading
             Debug.Assert(asyncWaiter != null, "Waiter should have been constructed");
             Debug.Assert(Monitor.IsEntered(m_lockObj!), "Requires the lock be held");
 
-            // Wait until either the task is completed, timeout occurs, or cancellation is requested.
-            // We need to ensure that the Task.Delay task is appropriately cleaned up if the await
-            // completes due to the asyncWaiter completing, so we use our own token that we can explicitly
-            // cancel, and we chain the caller's supplied token into it.
-            using (var cts = cancellationToken.CanBeCanceled ?
-                CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, default) :
-                new CancellationTokenSource())
+            if (millisecondsTimeout != Timeout.Infinite)
             {
-                var waitCompleted = Task.WhenAny(asyncWaiter, Task.Delay(millisecondsTimeout, cts.Token));
-                if (asyncWaiter == await waitCompleted.ConfigureAwait(false))
+                // Wait until either the task is completed, cancellation is requested, or the timeout occurs.
+                // We need to ensure that the Task.Delay task is appropriately cleaned up if the await
+                // completes due to the asyncWaiter completing, so we use our own token that we can explicitly
+                // cancel, and we chain the caller's supplied token into it.
+                using (var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, default))
                 {
-                    cts.Cancel(); // ensure that the Task.Delay task is cleaned up
-                    return true; // successfully acquired
+                    if (asyncWaiter == await TaskFactory.CommonCWAnyLogic(new Task[] { asyncWaiter, Task.Delay(millisecondsTimeout, cts.Token) }).ConfigureAwait(false))
+                    {
+                        cts.Cancel(); // ensure that the Task.Delay task is cleaned up
+                        return true; // successfully acquired
+                    }
+                }
+            }
+            else // millisecondsTimeout == Timeout.Infinite
+            {
+                // Wait until either the task is completed or cancellation is requested.
+                var cancellationTask = new Task();
+                using (cancellationToken.UnsafeRegister(s => ((Task)s!).TrySetResult(), cancellationTask))
+                {
+                    if (asyncWaiter == await TaskFactory.CommonCWAnyLogic(new Task[] { asyncWaiter, cancellationTask }).ConfigureAwait(false))
+                    {
+                        return true; // successfully acquired
+                    }
                 }
             }