Fix reusing TaskCompletionSource

Author: Meyn

Requests/Request.Static.cs

@@ -1,4 +1,3 @@
-
 using System.Runtime.CompilerServices;
 
 namespace Requests
@@ -14,7 +13,7 @@ public static class Request
         /// <summary>
         /// Gets the current request context.
         /// </summary>
-        internal static IRequest? Current => _current.Value;
+        public static IRequest? Current => _current.Value;
 
         /// <summary>
         /// Sets the current request context.
@@ -35,16 +34,26 @@ public static class Request
     /// <summary>
     /// Provides an awaitable type that enables yielding in request contexts.
     /// </summary>
-    public readonly struct YieldAwaitable : ICriticalNotifyCompletion
+    public readonly struct YieldAwaitable
     {
         private readonly IRequest? _context;
 
         internal YieldAwaitable(IRequest? context) => _context = context;
 
         /// <summary>
-        /// Gets an awaiter for this awaitable (returns itself).
+        /// Gets an awaiter for this awaitable.
         /// </summary>
-        public YieldAwaitable GetAwaiter() => this;
+        public YieldAwaiter GetAwaiter() => new(_context);
+    }
+
+    /// <summary>
+    /// Provides an awaiter that handles yielding in request contexts.
+    /// </summary>
+    public readonly struct YieldAwaiter : ICriticalNotifyCompletion
+    {
+        private readonly IRequest? _context;
+
+        internal YieldAwaiter(IRequest? context) => _context = context;
 
         /// <summary>
         /// Returns true if there's no context (completes synchronously as no-op).
@@ -64,7 +73,16 @@ public void GetResult() { }
         public void OnCompleted(Action continuation)
         {
             ArgumentNullException.ThrowIfNull(continuation);
-            _ = ScheduleContinuationAsync(_context!, continuation);
+
+            ExecutionContext? ctx = ExecutionContext.Capture();
+            if (ctx == null)
+            {
+                UnsafeOnCompleted(continuation);
+            }
+            else
+            {
+                UnsafeOnCompleted(() => ExecutionContext.Run(ctx, static s => ((Action)s!)(), continuation));
+            }
         }
 
         /// <summary>
@@ -74,24 +92,7 @@ public void OnCompleted(Action continuation)
         public void UnsafeOnCompleted(Action continuation)
         {
             ArgumentNullException.ThrowIfNull(continuation);
-            _ = ScheduleContinuationAsync(_context!, continuation);
-        }
-
-        /// <summary>
-        /// Schedules the continuation after yielding through the request context.
-        /// </summary>
-        private static async Task ScheduleContinuationAsync(IRequest context, Action continuation)
-        {
-            try
-            {
-                await context.YieldAsync().ConfigureAwait(false);
-                continuation();
-            }
-            catch
-            {
-                // Still invoke continuation to prevent deadlock
-                continuation();
-            }
+            _context!.YieldAsync().ConfigureAwait(false).GetAwaiter().UnsafeOnCompleted(continuation);
         }
     }
 }

Requests/Request.cs

@@ -20,7 +20,6 @@ public abstract partial class Request<TOptions, TCompleted, TFailed> : IRequest,
         private CancellationTokenRegistration _ctr;
 
         private TaskCompletionSource<bool>? _pauseTcs;
-        private volatile bool _hasPausedExecution;
 
         private readonly RequestStateMachine _stateMachine;
 
@@ -291,27 +290,21 @@ async Task IRequest.StartRequestAsync()
                 return;
 
             // Check if we're resuming a paused execution
-            if (_hasPausedExecution)
+            TaskCompletionSource<bool>? tcs = Interlocked.Exchange(ref _pauseTcs, null);
+            if (tcs != null)
             {
-                // Get and clear the pause TCS atomically
-                TaskCompletionSource<bool>? tcs = Interlocked.Exchange(ref _pauseTcs, null);
-                if (tcs != null)
-                {
-                    _hasPausedExecution = false;
-
-                    if (!_stateMachine.TryTransition(RequestState.Running))
-                        return;
+                if (!_stateMachine.TryTransition(RequestState.Running))
+                    return;
 
-                    _runningSourceVersion++;
-                    _runningSource.Reset();
+                _runningSourceVersion++;
+                _runningSource.Reset();
 
-                    // Resume the paused execution
-                    tcs.TrySetResult(true);
+                // Resume the paused execution
+                tcs.TrySetResult(true);
 
-                    // Wait for the resumed execution to stop running
-                    await new ValueTask(this, _runningSourceVersion).ConfigureAwait(false);
-                    return;
-                }
+                // Wait for the resumed execution to stop running
+                await new ValueTask(this, _runningSourceVersion).ConfigureAwait(false);
+                return;
             }
 
             if (!_stateMachine.TryTransition(RequestState.Running))
@@ -492,10 +485,8 @@ private async ValueTask YieldAsyncSlow()
 
             if (State == RequestState.Paused)
             {
-                // Lazy-create pause TCS only when actually pausing
-                TaskCompletionSource<bool> tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
-                _pauseTcs = tcs;
-                _hasPausedExecution = true;
+                TaskCompletionSource<bool> tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                tcs = Interlocked.CompareExchange(ref _pauseTcs, tcs, null) ?? tcs;
 
                 // Wait for resume
                 await tcs.Task.ConfigureAwait(false);
@@ -573,7 +564,7 @@ public virtual void Dispose()
             _requestCts?.Dispose();
             _ctr.Dispose();
             _completionSource.TrySetCanceled();
-            _pauseTcs?.TrySetCanceled();
+            Interlocked.Exchange(ref _pauseTcs, null)?.TrySetCanceled();
 
             GC.SuppressFinalize(this);
         }