76831-SystemNetQuicTestsQuicStreamTestsWriteCanceled_NextWriteThrows

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

@@ -29,6 +29,7 @@ private enum State
     private ManualResetValueTaskSourceCore<bool> _valueTaskSource;
     private CancellationTokenRegistration _cancellationRegistration;
     private Action<object?>? _cancellationAction;
+    private object? _cancellationActionLock;
     private GCHandle _keepAlive;
 
     private readonly TaskCompletionSource _finalTaskSource;
@@ -50,6 +51,11 @@ public ResettableValueTaskSource(bool runContinuationsAsynchronously = true)
     /// </summary>
     public Action<object?> CancellationAction { init { _cancellationAction = value; } }
 
+    /// <summary>
+    /// Optional lock object to be held during the <see cref="CancellationAction"/> callback. If no lock is provided, the callback will be invoked without holding any lock.
+    /// </summary>
+    public object? CancellationActionLock { init { _cancellationActionLock = 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)"/>
     /// was called with <c>final</c> set to <c>true</c> and the result was propagated.
@@ -91,10 +97,27 @@ public bool TryGetValueTask(out ValueTask valueTask, object? keepAlive = null, C
                     _cancellationRegistration = cancellationToken.UnsafeRegister(static (obj, cancellationToken) =>
                     {
                         (ResettableValueTaskSource thisRef, object? target) = ((ResettableValueTaskSource, object?))obj!;
-                        // This will transition the state to Ready.
-                        if (thisRef.TrySetException(new OperationCanceledException(cancellationToken)))
+
+                        object? lockObject = thisRef._cancellationActionLock;
+                        try
+                        {
+                            if (lockObject != null)
+                            {
+                                Monitor.Enter(lockObject);
+                            }
+
+                            // This will transition the state to Ready.
+                            if (thisRef.TrySetException(new OperationCanceledException(cancellationToken)))
+                            {
+                                thisRef._cancellationAction?.Invoke(target);
+                            }
+                        }
+                        finally
                         {
-                            thisRef._cancellationAction?.Invoke(target);
+                            if (lockObject != null)
+                            {
+                                Monitor.Exit(lockObject);
+                            }
                         }
                     }, (this, keepAlive));
                 }

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

@@ -89,7 +89,9 @@ public sealed partial class QuicStream
     private ReceiveBuffers _receiveBuffers = new ReceiveBuffers();
     private int _receivedNeedsEnable;
 
-    private readonly ResettableValueTaskSource _sendTcs = new ResettableValueTaskSource()
+    private readonly ResettableValueTaskSource _sendTcs;
+    // helper function since we need to pass the _shutdownTcs to the _sendTcs
+    private ResettableValueTaskSource CreateSendTaskCompletionSource() => new ResettableValueTaskSource()
     {
         CancellationAction = target =>
         {
@@ -106,7 +108,12 @@ public sealed partial class QuicStream
                 // when using CancellationTokenSource.CancelAfter.
                 // Ignore the exception
             }
-        }
+        },
+
+        // since we're using _sendTcs (Abortive) and _shutdownTcs (Graceful) for controlling write direction shutdown,
+        // CancellationAction above can race with disposing the stream. Holding _shutdownTcs will prevent Dispose from
+        // performing graceful shutdown while we are Aborting it.
+        CancellationActionLock = _shutdownTcs
     };
     private MsQuicBuffers _sendBuffers = new MsQuicBuffers();
     private readonly object _sendBuffersLock = new object();
@@ -185,6 +192,7 @@ internal unsafe QuicStream(MsQuicContextSafeHandle connectionHandle, QuicStreamT
             _receiveTcs.TrySetResult(final: true);
         }
         _type = type;
+        _sendTcs = CreateSendTaskCompletionSource();
     }
 
     /// <summary>
@@ -211,6 +219,7 @@ internal unsafe QuicStream(MsQuicContextSafeHandle connectionHandle, QUIC_HANDLE
             context.Free();
             throw;
         }
+        _sendTcs = CreateSendTaskCompletionSource();
 
         _defaultErrorCode = defaultErrorCode;
 
@@ -355,15 +364,18 @@ public ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, bool completeWrites, Ca
             return ValueTask.FromCanceled(cancellationToken);
         }
 
-        // Concurrent call, this one lost the race.
+        // attempt to reserve the task. If the token was already canceled, this will also invoke Abort(Write, DefaultErrorCode) and transition to final state.
         if (!_sendTcs.TryGetValueTask(out ValueTask valueTask, this, cancellationToken))
         {
+            // Concurrent call, this one lost the race.
             throw new InvalidOperationException(SR.Format(SR.net_io_invalidnestedcall, "write"));
         }
 
-        // No need to call anything since we already have a result, most likely an exception.
+        // Task.Delay(TimeSpan.FromSeconds(1), CancellationToken.None).Wait(CancellationToken.None);
+
         if (valueTask.IsCompleted)
         {
+            // No need to continue since we already have a result, most likely an exception.
             return valueTask;
         }