[QUIC] Add ShutdownCompleted method and fix receive and shutdown behavior in tests

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

@@ -104,6 +104,7 @@ public override void Flush() { }
         public override void SetLength(long value) { }
         public void Shutdown() { }
         public System.Threading.Tasks.ValueTask ShutdownWriteCompleted(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
+        public System.Threading.Tasks.ValueTask ShutdownCompleted(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public override void Write(byte[] buffer, int offset, int count) { }
         public override void Write(System.ReadOnlySpan<byte> buffer) { }
         public System.Threading.Tasks.ValueTask WriteAsync(System.Buffers.ReadOnlySequence<byte> buffers, bool endStream, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/Mock/MockConnection.cs

@@ -210,6 +210,12 @@ internal override ValueTask CloseAsync(long errorCode, CancellationToken cancell
             ConnectionState? state = _state;
             if (state is not null)
             {
+                if (state._closed)
+                {
+                    return default;
+                }
+                state._closed = true;
+
                 if (_isClient)
                 {
                     state._clientErrorCode = errorCode;
@@ -272,6 +278,7 @@ internal sealed class ConnectionState
             public Channel<MockStream.StreamState> _serverInitiatedStreamChannel;
             public long _clientErrorCode;
             public long _serverErrorCode;
+            public bool _closed;
 
             public ConnectionState(SslApplicationProtocol applicationProtocol)
             {

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/Mock/MockStream.cs

@@ -178,6 +178,14 @@ internal override ValueTask ShutdownWriteCompleted(CancellationToken cancellatio
             return default;
         }
 
+
+        internal override ValueTask ShutdownCompleted(CancellationToken cancellationToken = default)
+        {
+            CheckDisposed();
+
+            return default;
+        }
+
         internal override void Shutdown()
         {
             CheckDisposed();

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicConnection.cs

@@ -361,8 +361,6 @@ private ValueTask ShutdownAsync(
             QUIC_CONNECTION_SHUTDOWN_FLAGS Flags,
             long ErrorCode)
         {
-            Debug.Assert(!_state.ShutdownTcs.Task.IsCompleted);
-
             // Store the connection into the GCHandle'd state to prevent GC if user calls ShutdownAsync and gets rid of all references to the MsQuicConnection.
             Debug.Assert(_state.Connection == null);
             _state.Connection = this;

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicStream.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Net.Quic.Implementations.MsQuic.Internal;
+using System.Runtime.ExceptionServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -62,6 +63,12 @@ private sealed class State
 
             // Set once writes have been shutdown.
             public readonly TaskCompletionSource ShutdownWriteCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+
+            public ShutdownState ShutdownState;
+
+            // Set once stream have been shutdown.
+            public readonly TaskCompletionSource ShutdownCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
         }
 
         // inbound.
@@ -199,23 +206,25 @@ private async ValueTask<CancellationTokenRegistration> HandleWriteStartState(Can
                 }
             }
 
-            CancellationTokenRegistration registration = cancellationToken.Register(() =>
+            CancellationTokenRegistration registration = cancellationToken.UnsafeRegister(static (s, token) =>
             {
+                var stream = (MsQuicStream)s!;
                 bool shouldComplete = false;
-                lock (_state)
+                lock (stream._state)
                 {
-                    if (_state.SendState == SendState.None)
+                    if (stream._state.SendState == SendState.None)
                     {
-                        _state.SendState = SendState.Aborted;
+                        stream._state.SendState = SendState.Aborted;
                         shouldComplete = true;
                     }
                 }
 
                 if (shouldComplete)
                 {
-                    _state.SendResettableCompletionSource.CompleteException(new OperationCanceledException("Write was canceled", cancellationToken));
+                    stream._state.SendResettableCompletionSource.CompleteException(
+                        ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceledException("Write was canceled", token)));
                 }
-            });
+            }, this);
 
             // Make sure start has completed
             if (!_started)
@@ -268,24 +277,26 @@ internal override async ValueTask<int> ReadAsync(Memory<byte> destination, Cance
                 }
             }
 
-            using CancellationTokenRegistration registration = cancellationToken.Register(() =>
+            using CancellationTokenRegistration registration = cancellationToken.UnsafeRegister(static (s, token) =>
             {
+                var stream = (MsQuicStream)s!;
                 bool shouldComplete = false;
-                lock (_state)
+                lock (stream._state)
                 {
-                    if (_state.ReadState == ReadState.None)
+                    if (stream._state.ReadState == ReadState.None)
                     {
                         shouldComplete = true;
                     }
 
-                    _state.ReadState = ReadState.Aborted;
+                    stream._state.ReadState = ReadState.Aborted;
                 }
 
                 if (shouldComplete)
                 {
-                    _state.ReceiveResettableCompletionSource.CompleteException(new OperationCanceledException("Read was canceled", cancellationToken));
+                    stream._state.ReceiveResettableCompletionSource.CompleteException(
+                        ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceledException("Read was canceled", token)));
                 }
-            });
+            }, this);
 
             // TODO there could potentially be a perf gain by storing the buffer from the inital read
             // This reduces the amount of async calls, however it makes it so MsQuic holds onto the buffers
@@ -358,7 +369,8 @@ internal override void AbortWrite(long errorCode)
 
             if (shouldComplete)
             {
-                _state.ShutdownWriteCompletionSource.SetException(new QuicStreamAbortedException("Shutdown was aborted.", errorCode));
+                _state.ShutdownWriteCompletionSource.SetException(
+                    ExceptionDispatchInfo.SetCurrentStackTrace(new QuicStreamAbortedException("Shutdown was aborted.", errorCode)));
             }
 
             StartShutdown(QUIC_STREAM_SHUTDOWN_FLAGS.ABORT_SEND, errorCode);
@@ -375,30 +387,61 @@ internal override async ValueTask ShutdownWriteCompleted(CancellationToken cance
             ThrowIfDisposed();
 
             // TODO do anything to stop writes?
-            using CancellationTokenRegistration registration = cancellationToken.Register(() =>
+            using CancellationTokenRegistration registration = cancellationToken.UnsafeRegister(static (s, token) =>
             {
+                var stream = (MsQuicStream)s!;
                 bool shouldComplete = false;
-                lock (_state)
+                lock (stream._state)
                 {
-                    if (_state.ShutdownWriteState == ShutdownWriteState.None)
+                    if (stream._state.ShutdownWriteState == ShutdownWriteState.None)
                     {
-                        _state.ShutdownWriteState = ShutdownWriteState.Canceled;
+                        stream._state.ShutdownWriteState = ShutdownWriteState.Canceled; // TODO: should we separate states for cancelling here vs calling Abort?
                         shouldComplete = true;
                     }
                 }
 
                 if (shouldComplete)
                 {
-                    _state.ShutdownWriteCompletionSource.SetException(new OperationCanceledException("Shutdown was canceled", cancellationToken));
+                    stream._state.ShutdownWriteCompletionSource.SetException(
+                        ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceledException("Wait for shutdown write was canceled", token)));
                 }
-            });
+            }, this);
 
             await _state.ShutdownWriteCompletionSource.Task.ConfigureAwait(false);
         }
 
+        internal override async ValueTask ShutdownCompleted(CancellationToken cancellationToken = default)
+        {
+            ThrowIfDisposed();
+
+            // TODO do anything to stop writes?
+            using CancellationTokenRegistration registration = cancellationToken.UnsafeRegister(static (s, token) =>
+            {
+                var stream = (MsQuicStream)s!;
+                bool shouldComplete = false;
+                lock (stream._state)
+                {
+                    if (stream._state.ShutdownState == ShutdownState.None)
+                    {
+                        stream._state.ShutdownState = ShutdownState.Canceled;
+                        shouldComplete = true;
+                    }
+                }
+
+                if (shouldComplete)
+                {
+                    stream._state.ShutdownWriteCompletionSource.SetException(
+                        ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceledException("Wait for shutdown was canceled", token)));
+                }
+            }, this);
+
+            await _state.ShutdownCompletionSource.Task.ConfigureAwait(false);
+        }
+
         internal override void Shutdown()
         {
             ThrowIfDisposed();
+            // it is ok to send shutdown several times, MsQuic will ignore it
             StartShutdown(QUIC_STREAM_SHUTDOWN_FLAGS.GRACEFUL, errorCode: 0);
         }
 
@@ -481,6 +524,11 @@ private static uint NativeCallbackHandler(
 
         private static uint HandleEvent(State state, ref StreamEvent evt)
         {
+            if (NetEventSource.Log.IsEnabled())
+            {
+                NetEventSource.Info(state, $"[{state.GetHashCode()}] received event {evt.Type}");
+            }
+
             try
             {
                 switch ((QUIC_STREAM_EVENT_TYPE)evt.Type)
@@ -564,7 +612,8 @@ private static uint HandleEventPeerRecvAborted(State state, ref StreamEvent evt)
 
             if (shouldComplete)
             {
-                state.SendResettableCompletionSource.CompleteException(new QuicStreamAbortedException(state.SendErrorCode));
+                state.SendResettableCompletionSource.CompleteException(
+                    ExceptionDispatchInfo.SetCurrentStackTrace(new QuicStreamAbortedException(state.SendErrorCode)));
             }
 
             return MsQuicStatusCodes.Success;
@@ -604,7 +653,7 @@ private static uint HandleEventSendShutdownComplete(State state, ref StreamEvent
 
             if (shouldComplete)
             {
-                state.ShutdownWriteCompletionSource.TrySetResult();
+                state.ShutdownWriteCompletionSource.SetResult();
             }
 
             return MsQuicStatusCodes.Success;
@@ -614,6 +663,7 @@ private static uint HandleEventShutdownComplete(State state)
         {
             bool shouldReadComplete = false;
             bool shouldShutdownWriteComplete = false;
+            bool shouldShutdownComplete = false;
 
             lock (state)
             {
@@ -632,6 +682,12 @@ private static uint HandleEventShutdownComplete(State state)
                     state.ShutdownWriteState = ShutdownWriteState.Finished;
                     shouldShutdownWriteComplete = true;
                 }
+
+                if (state.ShutdownState == ShutdownState.None)
+                {
+                    state.ShutdownState = ShutdownState.Finished;
+                    shouldShutdownComplete = true;
+                }
             }
 
             if (shouldReadComplete)
@@ -641,7 +697,12 @@ private static uint HandleEventShutdownComplete(State state)
 
             if (shouldShutdownWriteComplete)
             {
-                state.ShutdownWriteCompletionSource.TrySetResult();
+                state.ShutdownWriteCompletionSource.SetResult();
+            }
+
+            if (shouldShutdownComplete)
+            {
+                state.ShutdownCompletionSource.SetResult();
             }
 
             return MsQuicStatusCodes.Success;
@@ -662,7 +723,8 @@ private static uint HandleEventPeerSendAborted(State state, ref StreamEvent evt)
 
             if (shouldComplete)
             {
-                state.ReceiveResettableCompletionSource.CompleteException(new QuicStreamAbortedException(state.ReadErrorCode));
+                state.ReceiveResettableCompletionSource.CompleteException(
+                    ExceptionDispatchInfo.SetCurrentStackTrace(new QuicStreamAbortedException(state.ReadErrorCode)));
             }
 
             return MsQuicStatusCodes.Success;
@@ -942,6 +1004,13 @@ private enum ShutdownWriteState
             Finished
         }
 
+        private enum ShutdownState
+        {
+            None,
+            Canceled,
+            Finished
+        }
+
         private enum SendState
         {
             None,

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/QuicStreamProvider.cs

@@ -39,6 +39,8 @@ internal abstract class QuicStreamProvider : IDisposable, IAsyncDisposable
 
         internal abstract ValueTask ShutdownWriteCompleted(CancellationToken cancellationToken = default);
 
+        internal abstract ValueTask ShutdownCompleted(CancellationToken cancellationToken = default);
+
         internal abstract void Shutdown();
 
         internal abstract void Flush();

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

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net.Quic.Implementations;
-using System.Net.Quic.Implementations.MsQuic.Internal;
 using System.Net.Security;
 using System.Threading;
 using System.Threading.Tasks;

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

@@ -101,6 +101,8 @@ public override Task WriteAsync(byte[] buffer, int offset, int count, Cancellati
 
         public ValueTask ShutdownWriteCompleted(CancellationToken cancellationToken = default) => _provider.ShutdownWriteCompleted(cancellationToken);
 
+        public ValueTask ShutdownCompleted(CancellationToken cancellationToken = default) => _provider.ShutdownCompleted(cancellationToken);
+
         public void Shutdown() => _provider.Shutdown();
 
         protected override void Dispose(bool disposing)

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

@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using System.Net.Security;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Net.Quic.Tests
@@ -50,18 +51,27 @@ internal async Task RunClientServer(Func<QuicConnection, Task> clientFunction, F
         {
             using QuicListener listener = CreateQuicListener();
 
+            var serverFinished = new ManualResetEventSlim();
+            var clientFinished = new ManualResetEventSlim();
+
             await new[]
             {
                 Task.Run(async () =>
                 {
                     using QuicConnection serverConnection = await listener.AcceptConnectionAsync();
                     await serverFunction(serverConnection);
+                    serverFinished.Set();
+                    clientFinished.Wait();
+                    await serverConnection.CloseAsync(0);
                 }),
                 Task.Run(async () =>
                 {
                     using QuicConnection clientConnection = CreateQuicConnection(listener.ListenEndPoint);
                     await clientConnection.ConnectAsync();
                     await clientFunction(clientConnection);
+                    clientFinished.Set();
+                    serverFinished.Wait();
+                    await clientConnection.CloseAsync(0);
                 })
             }.WhenAllOrAnyFailed(millisecondsTimeout);
         }