[S.N.Quic] Observe exceptions in ResettableValueTaskSource (#114226)

* Observe exceptions in ResettableValueTaskSource

* Update src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs

Co-authored-by: Marie Píchová <[email protected]>

---------

Co-authored-by: Marie Píchová <[email protected]>

Author: antonfirsov

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

@@ -52,7 +52,7 @@ public ResettableValueTaskSource()
     public Action<object?> CancellationAction { init { _cancellationAction = value; } }
 
     /// <summary>
-    /// Returns <c>true</c> is this task source has entered its final state, i.e. <see cref="TrySetResult(bool)"/> or <see cref="TrySetException(Exception, bool)"/>
+    /// Returns <c>true</c> is this task source has entered its final state, i.e. <see cref="TrySetResult(bool)"/> or <see cref="TrySetException(Exception)"/>
     /// was called with <c>final</c> set to <c>true</c> and the result was propagated.
     /// </summary>
     public bool IsCompleted => (State)Volatile.Read(ref Unsafe.As<State, byte>(ref _state)) == State.Completed;
@@ -173,6 +173,7 @@ private bool TryComplete(Exception? exception, bool final)
                 // Unblock the current task source and in case of a final also the final task source.
                 if (exception is not null)
                 {
+                    Debug.Assert(final);
                     // Set up the exception stack trace for the caller.
                     exception = exception.StackTrace is null ? ExceptionDispatchInfo.SetCurrentStackTrace(exception) : exception;
                     if (state is State.None or State.Awaiting)
@@ -192,7 +193,7 @@ private bool TryComplete(Exception? exception, bool final)
                     if (_finalTaskSource.TryComplete(exception))
                     {
                         // Signal the final task only if we don't have another result in the value task source.
-                        // In that case, the final task will be signalled after the value task result is retrieved.
+                        // In that case, the final task will be signaled after the value task result is retrieved.
                         if (state != State.Ready)
                         {
                             _finalTaskSource.TrySignal(out _);
@@ -226,15 +227,14 @@ public bool TrySetResult(bool final = false)
     }
 
     /// <summary>
-    /// Tries to transition from <see cref="State.Awaiting"/> to either <see cref="State.Ready"/> or <see cref="State.Completed"/>, depending on the value of <paramref name="final"/>.
-    /// Only the first call is able to do that with the exception of <c>TrySetResult()</c> followed by <c>TrySetResult(true)</c>, which will both return <c>true</c>.
+    /// Tries to transition from <see cref="State.Awaiting"/> to <see cref="State.Completed"/>, setting an exception.
+    /// Only the first call is able to do that.
     /// </summary>
-    /// <param name="final">Whether this is the final transition to <see cref="State.Completed" /> or just a transition into <see cref="State.Ready"/> from which the task source can be reset back to <see cref="State.None"/>.</param>
     /// <param name="exception">The exception to set as a result of the value task.</param>
     /// <returns><c>true</c> if this is the first call that set the result; otherwise, <c>false</c>.</returns>
-    public bool TrySetException(Exception exception, bool final = false)
+    public bool TrySetException(Exception exception)
     {
-        return TryComplete(exception, final);
+        return TryComplete(exception, final: true);
     }
 
     ValueTaskSourceStatus IValueTaskSource.GetStatus(short token)
@@ -297,6 +297,7 @@ private struct FinalTaskSource
         private bool _isCompleted;
         private bool _isSignaled;
         private Exception? _exception;
+        private Task? _signaledTask;
 
         public FinalTaskSource()
         {
@@ -312,9 +313,9 @@ public Task GetTask(object? keepAlive)
             {
                 if (_isSignaled)
                 {
-                    return _exception is null
-                        ? Task.CompletedTask
-                        : Task.FromException(_exception);
+                    _signaledTask ??= _exception is null ? Task.CompletedTask : Task.FromException(_exception);
+                    _ = _signaledTask.Exception; // Observe the exception.
+                    return _signaledTask;
                 }
 
                 _finalTaskSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
@@ -350,13 +351,17 @@ public bool TrySignal(out Exception? exception)
                 return false;
             }
 
-            if (_exception is not null)
+            if (_finalTaskSource is not null)
             {
-                _finalTaskSource?.SetException(_exception);
-            }
-            else
-            {
-                _finalTaskSource?.SetResult();
+                if (_exception is not null)
+                {
+                    _finalTaskSource.SetException(_exception);
+                    _ = _finalTaskSource.Task.Exception; // Observe the exception.
+                }
+                else
+                {
+                    _finalTaskSource.SetResult();
+                }
             }
 
             exception = _exception;

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

@@ -430,7 +430,7 @@ public ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, bool completeWrites, Ca
                     exception = Volatile.Read(ref _sendException);
                     if (exception is not null)
                     {
-                        _sendTcs.TrySetException(exception, final: true);
+                        _sendTcs.TrySetException(exception);
                     }
                 }
                 // SEND_COMPLETE expected, buffer and lock will be released then.
@@ -487,15 +487,15 @@ public void Abort(QuicAbortDirection abortDirection, long errorCode)
 
         if (abortDirection.HasFlag(QuicAbortDirection.Read))
         {
-            _receiveTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_reading_aborted), final: true);
+            _receiveTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_reading_aborted));
         }
         if (abortDirection.HasFlag(QuicAbortDirection.Write))
         {
             var exception = ThrowHelper.GetOperationAbortedException(SR.net_quic_writing_aborted);
             Interlocked.CompareExchange(ref _sendException, exception, null);
             if (Interlocked.CompareExchange(ref _sendLocked, 1, 0) == 0)
             {
-                _sendTcs.TrySetException(_sendException, final: true);
+                _sendTcs.TrySetException(_sendException);
                 Volatile.Write(ref _sendLocked, 0);
             }
         }
@@ -585,7 +585,7 @@ private unsafe int HandleEventSendComplete(ref SEND_COMPLETE_DATA data)
         Exception? exception = Volatile.Read(ref _sendException);
         if (exception is not null)
         {
-            _sendTcs.TrySetException(exception, final: true);
+            _sendTcs.TrySetException(exception);
         }
         if (data.Canceled == 0)
         {
@@ -604,12 +604,12 @@ private unsafe int HandleEventPeerSendShutdown()
     }
     private unsafe int HandleEventPeerSendAborted(ref PEER_SEND_ABORTED_DATA data)
     {
-        _receiveTcs.TrySetException(ThrowHelper.GetStreamAbortedException((long)data.ErrorCode), final: true);
+        _receiveTcs.TrySetException(ThrowHelper.GetStreamAbortedException((long)data.ErrorCode));
         return QUIC_STATUS_SUCCESS;
     }
     private unsafe int HandleEventPeerReceiveAborted(ref PEER_RECEIVE_ABORTED_DATA data)
     {
-        _sendTcs.TrySetException(ThrowHelper.GetStreamAbortedException((long)data.ErrorCode), final: true);
+        _sendTcs.TrySetException(ThrowHelper.GetStreamAbortedException((long)data.ErrorCode));
         return QUIC_STATUS_SUCCESS;
     }
     private unsafe int HandleEventSendShutdownComplete(ref SEND_SHUTDOWN_COMPLETE_DATA data)
@@ -639,8 +639,8 @@ private unsafe int HandleEventShutdownComplete(ref SHUTDOWN_COMPLETE_DATA data)
                 (shutdownByApp: false, closedRemotely: false) => ThrowHelper.GetExceptionForMsQuicStatus(data.ConnectionCloseStatus, (long)data.ConnectionErrorCode),
             };
             _startedTcs.TrySetException(exception);
-            _receiveTcs.TrySetException(exception, final: true);
-            _sendTcs.TrySetException(exception, final: true);
+            _receiveTcs.TrySetException(exception);
+            _sendTcs.TrySetException(exception);
         }
         _startedTcs.TrySetException(ThrowHelper.GetOperationAbortedException());
         _shutdownTcs.TrySetResult();
@@ -766,11 +766,11 @@ unsafe void StreamShutdown(QUIC_STREAM_SHUTDOWN_FLAGS flags, long errorCode)
             {
                 if (flags.HasFlag(QUIC_STREAM_SHUTDOWN_FLAGS.ABORT_RECEIVE) && !_receiveTcs.IsCompleted)
                 {
-                    _receiveTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_reading_aborted), final: true);
+                    _receiveTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_reading_aborted));
                 }
                 if (flags.HasFlag(QUIC_STREAM_SHUTDOWN_FLAGS.ABORT_SEND) && !_sendTcs.IsCompleted)
                 {
-                    _sendTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_writing_aborted), final: true);
+                    _sendTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_writing_aborted));
                 }
             }
         }