Fix potential crash

Author: rzikm

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

@@ -525,12 +525,12 @@ private unsafe int HandleEventShutdownInitiatedByTransport(ref SHUTDOWN_INITIATE
     {
         Exception exception = ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetExceptionForMsQuicStatus(data.Status, (long)data.ErrorCode));
         _connectedTcs.TrySetException(exception);
-        CompleteAndDrainAcceptQueue(exception);
+        CompleteAcceptQueue(exception, false);
         return QUIC_STATUS_SUCCESS;
     }
     private unsafe int HandleEventShutdownInitiatedByPeer(ref SHUTDOWN_INITIATED_BY_PEER_DATA data)
     {
-        CompleteAndDrainAcceptQueue(ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetConnectionAbortedException((long)data.ErrorCode)));
+        CompleteAcceptQueue(ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetConnectionAbortedException((long)data.ErrorCode)), false);
         return QUIC_STATUS_SUCCESS;
     }
     private unsafe int HandleEventShutdownComplete()
@@ -539,7 +539,7 @@ private unsafe int HandleEventShutdownComplete()
         _tlsSecret?.WriteSecret();
 
         Exception exception = ExceptionDispatchInfo.SetCurrentStackTrace(_disposed == 1 ? new ObjectDisposedException(GetType().FullName) : ThrowHelper.GetOperationAbortedException());
-        CompleteAndDrainAcceptQueue(exception);
+        CompleteAcceptQueue(exception, true);
         _connectedTcs.TrySetException(exception);
         _shutdownTokenSource.Cancel();
         _shutdownTcs.TrySetResult(final: true);
@@ -643,22 +643,26 @@ private static unsafe int NativeCallback(QUIC_HANDLE* connection, void* context,
         }
     }
 
-    private void CompleteAndDrainAcceptQueue(Exception? ex)
+    private void CompleteAcceptQueue(Exception? ex, bool drain)
     {
-        if (_acceptQueue.Writer.TryComplete(ex))
+        _acceptQueue.Writer.TryComplete(ex);
+
+        if (drain)
         {
-            // also drain the queue. Because stream shutdown events are indicated before connection shutdown,
-            // the QuicStream instances have already been signaled and closed internally. We only need to dispose them,
-            // which in this situation should complete synchronously.
+            // This should be only called after connection SHUTDOWN_COMPLETE has been indicated.
+            // At that point, all streams have been already shut down internally and we need
+            // only to close the handle via dispose, so DisposeAsync below should complete
+            // synchronously (which is necessary for this method to be callable from MsQuic
+            // event callback).
             while (_acceptQueue.Reader.TryRead(out QuicStream? stream))
             {
                 ValueTask task = stream.DisposeAsync();
                 Debug.Assert(task.IsCompletedSuccessfully);
                 task.GetAwaiter().GetResult();
             }
-        }
 
-        Debug.Assert(_acceptQueue.Reader.Completion.IsCompleted);
+            Debug.Assert(_acceptQueue.Reader.Completion.IsCompleted);
+        }
     }
 
     /// <summary>
@@ -709,6 +713,6 @@ public async ValueTask DisposeAsync()
         }
 
         // Flush the queue and dispose all remaining streams.
-        CompleteAndDrainAcceptQueue(ExceptionDispatchInfo.SetCurrentStackTrace(new ObjectDisposedException(GetType().FullName)));
+        CompleteAcceptQueue(ExceptionDispatchInfo.SetCurrentStackTrace(new ObjectDisposedException(GetType().FullName)), true);
     }
 }