Revert "Use ordinary TCS and revert changes in ValueTaskSource"

This reverts commit f393313f886bf2073ddaea9b02f9bb30b80f047b.

Author: ManickaP

src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ValueTaskSource.cs

@@ -25,13 +25,15 @@ private enum State : byte
     private State _state;
     private ManualResetValueTaskSourceCore<bool> _valueTaskSource;
     private CancellationTokenRegistration _cancellationRegistration;
+    private Exception? _exception;
     private GCHandle _keepAlive;
 
     public ValueTaskSource()
     {
         _state = State.None;
         _valueTaskSource = new ManualResetValueTaskSourceCore<bool>() { RunContinuationsAsynchronously = true };
         _cancellationRegistration = default;
+        _exception = default;
         _keepAlive = default;
     }
 
@@ -75,7 +77,7 @@ public bool TryInitialize(out ValueTask valueTask, object? keepAlive = null, Can
             State state = _state;
 
             // If we're the first here, we will return true.
-            if (state == State.None)
+            if (state == State.None && _valueTaskSource.GetStatus(_valueTaskSource.Version) == ValueTaskSourceStatus.Pending)
             {
                 // Keep alive the caller object until the result is read from the task.
                 // Used for keeping caller alive during async interop calls.
@@ -104,30 +106,42 @@ private bool TryComplete(Exception? exception)
                 {
                     State state = _state;
 
-                    if (state != State.Completed)
+                    // Completed: nothing to do.
+                    if (state == State.Completed)
                     {
-                        _state = State.Completed;
+                        return false;
+                    }
+
+                    // With cancellation, keep the state as-is so it can be restored after the OCE is consumed.
+                    _state = exception is OperationCanceledException ? state : State.Completed;
 
-                        // Swap the cancellation registration so the one that's been registered gets eventually Disposed.
-                        // Ideally, we would dispose it here, but if the callbacks kicks in, it tries to take the lock held by this thread leading to deadlock.
-                        cancellationRegistration = _cancellationRegistration;
-                        _cancellationRegistration = default;
+                    // Swap the cancellation registration so the one that's been registered gets eventually Disposed.
+                    // Ideally, we would dispose it here, but if the callbacks kicks in, it tries to take the lock held by this thread leading to deadlock.
+                    cancellationRegistration = _cancellationRegistration;
+                    _cancellationRegistration = default;
 
-                        if (exception is not null)
+                    if (exception is not null)
+                    {
+                        // Set up the exception stack trace for the caller.
+                        exception = exception.StackTrace is null ? ExceptionDispatchInfo.SetCurrentStackTrace(exception) : exception;
+                        if (_valueTaskSource.GetStatus(_valueTaskSource.Version) == ValueTaskSourceStatus.Pending)
                         {
-                            // Set up the exception stack trace for the caller.
-                            exception = exception.StackTrace is null ? ExceptionDispatchInfo.SetCurrentStackTrace(exception) : exception;
                             _valueTaskSource.SetException(exception);
                         }
-                        else
+                        else if (exception is not OperationCanceledException)
+                        {
+                            _exception = exception;
+                        }
+                    }
+                    else
+                    {
+                        if (_valueTaskSource.GetStatus(_valueTaskSource.Version) == ValueTaskSourceStatus.Pending)
                         {
                             _valueTaskSource.SetResult(true);
                         }
-
-                        return true;
                     }
 
-                    return false;
+                    return true;
                 }
                 finally
                 {
@@ -173,5 +187,34 @@ void IValueTaskSource.OnCompleted(Action<object?> continuation, object? state, s
         => _valueTaskSource.OnCompleted(continuation, state, token, flags);
 
     void IValueTaskSource.GetResult(short token)
-        => _valueTaskSource.GetResult(token);
+    {
+        try
+        {
+            _valueTaskSource.GetResult(token);
+        }
+        finally
+        {
+            lock (this)
+            {
+                State state = _state;
+
+                // In case of a cancellation, reset the task and set the stored results if necessary.
+                if (_valueTaskSource.GetStatus(_valueTaskSource.Version) == ValueTaskSourceStatus.Canceled)
+                {
+                    _valueTaskSource.Reset();
+                    if (state == State.Completed)
+                    {
+                        if (_exception is not null)
+                        {
+                            _valueTaskSource.SetException(_exception);
+                        }
+                        else
+                        {
+                            _valueTaskSource.SetResult(true);
+                        }
+                    }
+                }
+            }
+        }
+    }
 }

src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs

@@ -109,7 +109,7 @@ static async ValueTask<QuicConnection> StartConnectAsync(QuicClientConnectionOpt
     private int _disposed;
 
     private readonly ValueTaskSource _connectedTcs = new ValueTaskSource();
-    private TaskCompletionSource? _shutdownTcs;
+    private readonly ValueTaskSource _shutdownTcs = new ValueTaskSource();
 
     private readonly CancellationTokenSource _shutdownTokenSource = new CancellationTokenSource();
 
@@ -467,7 +467,7 @@ public ValueTask CloseAsync(long errorCode, CancellationToken cancellationToken
     {
         ObjectDisposedException.ThrowIf(_disposed == 1, this);
 
-        if (Interlocked.CompareExchange(ref _shutdownTcs, new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), null) == null)
+        if (_shutdownTcs.TryInitialize(out ValueTask valueTask, this, cancellationToken))
         {
             unsafe
             {
@@ -478,7 +478,7 @@ public ValueTask CloseAsync(long errorCode, CancellationToken cancellationToken
             }
         }
 
-        return new ValueTask(_shutdownTcs.Task.WaitAsync(cancellationToken));
+        return valueTask;
     }
 
     private unsafe int HandleEventConnected(ref CONNECTED_DATA data)
@@ -520,7 +520,6 @@ private unsafe int HandleEventShutdownComplete()
         _acceptQueue.Writer.TryComplete(exception);
         _connectedTcs.TrySetException(exception);
         _shutdownTokenSource.Cancel();
-        Interlocked.CompareExchange(ref _shutdownTcs, new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), null);
         _shutdownTcs.TrySetResult();
         return QUIC_STATUS_SUCCESS;
     }
@@ -627,7 +626,7 @@ public async ValueTask DisposeAsync()
         }
 
         // Check if the connection has been shut down and if not, shut it down.
-        if (Interlocked.CompareExchange(ref _shutdownTcs, new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), null) == null)
+        if (_shutdownTcs.TryInitialize(out ValueTask valueTask, this))
         {
             unsafe
             {
@@ -637,7 +636,7 @@ public async ValueTask DisposeAsync()
                     (ulong)_defaultCloseErrorCode);
             }
         }
-        else if (!_shutdownTcs.Task.IsCompletedSuccessfully)
+        else if (!valueTask.IsCompletedSuccessfully)
         {
             unsafe
             {
@@ -649,7 +648,7 @@ public async ValueTask DisposeAsync()
         }
 
         // Wait for SHUTDOWN_COMPLETE, the last event, so that all resources can be safely released.
-        await _shutdownTcs.Task.ConfigureAwait(false);
+        await valueTask.ConfigureAwait(false);
         Debug.Assert(_connectedTcs.IsCompleted);
         _handle.Dispose();
         _shutdownTokenSource.Dispose();

src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs

@@ -134,27 +134,7 @@ await RunClientServer(
                 {
                     var cts = new CancellationTokenSource();
                     cts.Cancel();
-                    await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await clientConnection.CloseAsync(ExpectedErrorCode, cts.Token));
-                    await clientConnection.DisposeAsync();
-                    sync.Release();
-                },
-                async serverConnection =>
-                {
-                    await sync.WaitAsync();
-                    await serverConnection.DisposeAsync();
-                });
-        }
-
-        [Fact]
-        public async Task DisposeAfterCloseTaskStored()
-        {
-            using var sync = new SemaphoreSlim(0);
-
-            await RunClientServer(
-                async clientConnection =>
-                {
-                    var cts = new CancellationTokenSource();
-                    var task = clientConnection.CloseAsync(0).AsTask();
+                    await Assert.ThrowsAsync<OperationCanceledException>(async () => await clientConnection.CloseAsync(ExpectedErrorCode, cts.Token));
                     await clientConnection.DisposeAsync();
                     sync.Release();
                 },