add Http3LoopbackConnection.ShutdownAsync and use in AcceptConnection…

…Async

src/libraries/Common/tests/System/Net/Http/Http3LoopbackConnection.cs

@@ -36,7 +36,11 @@ internal sealed class Http3LoopbackConnection : GenericLoopbackConnection
 
         // This is specifically request streams, not control streams
         private readonly Dictionary<int, Http3LoopbackStream> _openStreams = new Dictionary<int, Http3LoopbackStream>();
+
         private Http3LoopbackStream _currentStream;
+        // We can't retrieve the stream ID after the stream is disposed, so store it separately
+        // Initialize it to -4 so that the firstInvalidStreamId calculation will work even if we never process a request
+        private long _currentStreamId = -4;
 
         private Http3LoopbackStream _inboundControlStream;      // Inbound control stream from client
         private Http3LoopbackStream _outboundControlStream;     // Our outbound control stream
@@ -111,6 +115,7 @@ public async Task<Http3LoopbackStream> AcceptStreamAsync()
             {
                 _openStreams.Add(checked((int)quicStream.StreamId), stream);
                 _currentStream = stream;
+                _currentStreamId = quicStream.StreamId;
             }
 
             return stream;
@@ -223,6 +228,8 @@ public override async Task<HttpRequestData> HandleRequestAsync(HttpStatusCode st
 
             // We are about to close the connection, after we send the response.
             // So, send a GOAWAY frame now so the client won't inadvertantly try to reuse the connection.
+            // Note that in HTTP3 (unlike HTTP2) there is no strict ordering between the GOAWAY and the response below;
+            // so the client may race in processing them and we need to handle this.
             await _outboundControlStream.SendGoAwayFrameAsync(stream.StreamId + 4);
 
             await stream.SendResponseAsync(statusCode, headers, content).ConfigureAwait(false);
@@ -232,6 +239,34 @@ public override async Task<HttpRequestData> HandleRequestAsync(HttpStatusCode st
             return request;
         }
 
+        public async Task ShutdownAsync()
+        {
+            try
+            {
+                long firstInvalidStreamId = _currentStreamId + 4;
+                await _outboundControlStream.SendGoAwayFrameAsync(firstInvalidStreamId);
+            }
+            catch (QuicConnectionAbortedException abortException) when (abortException.ErrorCode == H3_NO_ERROR)
+            {
+                // Client must have closed the connection already because the HttpClientHandler instance was disposed.
+                // So nothing to do.
+                return;
+            }
+            catch (OperationCanceledException)
+            {
+                // If the client is closing the connection at the same time we are trying to send the GOAWAY above,
+                // this can result in OperationCanceledException from QuicStream.WriteAsync.
+                // See https://github.com/dotnet/runtime/issues/58078
+                // I saw this consistently with GetAsync_EmptyResponseHeader_Success.
+                // To work around this, just eat the exception for now.
+                // Also, be aware of this issue as it will do weird things with OperationCanceledException and can
+                // make debugging this very confusing: https://github.com/dotnet/runtime/issues/58081
+                return;
+            }
+
+            await WaitForClientDisconnectAsync();
+        }
+
         // Wait for the client to close the connection, e.g. after we send a GOAWAY, or after the HttpClient is disposed.
         public async Task WaitForClientDisconnectAsync(bool refuseNewRequests = true)
         {

src/libraries/Common/tests/System/Net/Http/Http3LoopbackServer.cs

@@ -52,7 +52,7 @@ public override void Dispose()
             _cert.Dispose();
         }
 
-        public override async Task<GenericLoopbackConnection> EstablishGenericConnectionAsync()
+        private async Task<Http3LoopbackConnection> EstablishHttp3ConnectionAsync()
         {
             QuicConnection con = await _listener.AcceptConnectionAsync().ConfigureAwait(false);
             Http3LoopbackConnection connection = new Http3LoopbackConnection(con);
@@ -61,10 +61,16 @@ public override async Task<GenericLoopbackConnection> EstablishGenericConnection
             return connection;
         }
 
+        public override async Task<GenericLoopbackConnection> EstablishGenericConnectionAsync()
+        {
+            return await EstablishHttp3ConnectionAsync();
+        }
+
         public override async Task AcceptConnectionAsync(Func<GenericLoopbackConnection, Task> funcAsync)
         {
-            using GenericLoopbackConnection con = await EstablishGenericConnectionAsync().ConfigureAwait(false);
+            using Http3LoopbackConnection con = await EstablishHttp3ConnectionAsync().ConfigureAwait(false);
             await funcAsync(con).ConfigureAwait(false);
+            await con.ShutdownAsync();
         }
 
         public override async Task<HttpRequestData> HandleRequestAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null, string content = "")

src/libraries/Common/tests/System/Net/Http/HttpClientHandlerTest.Cancellation.cs

@@ -210,12 +210,9 @@ public async Task GetAsync_CancelDuringResponseBodyReceived_Unbuffered_TaskCance
             using (HttpClient client = CreateHttpClient())
             {
                 client.Timeout = Timeout.InfiniteTimeSpan;
-                var cts = new CancellationTokenSource();
 
                 await LoopbackServerFactory.CreateServerAsync(async (server, url) =>
                 {
-                    var clientFinished = new TaskCompletionSource<bool>();
-
                     Task serverTask = server.AcceptConnectionAsync(async connection =>
                     {
                         var headers = new List<HttpHeaderData>();
@@ -227,28 +224,31 @@ await LoopbackServerFactory.CreateServerAsync(async (server, url) =>
 
                         await connection.ReadRequestDataAsync();
                         await connection.SendResponseAsync(HttpStatusCode.OK, headers: headers, isFinal: false);
-                        await clientFinished.Task;
+
+                        await connection.WaitForCancellationAsync();
                     });
 
                     var req = new HttpRequestMessage(HttpMethod.Get, url) { Version = UseVersion };
                     req.Headers.ConnectionClose = connectionClose;
-                    Task<HttpResponseMessage> getResponse = client.SendAsync(TestAsync, req, HttpCompletionOption.ResponseHeadersRead, cts.Token);
+                    Task<HttpResponseMessage> getResponse = client.SendAsync(TestAsync, req, HttpCompletionOption.ResponseHeadersRead);
                     await ValidateClientCancellationAsync(async () =>
                     {
-                        HttpResponseMessage resp = await getResponse;
+                        // This 'using' shouldn't be necessary in general. However, HTTP3 does not remove the request stream from the
+                        // active stream table until the user disposes the response (or it gets finalized).
+                        // This means the connection will fail to shut down promptly.
+                        // See https://github.com/dotnet/runtime/issues/58072
+                        using HttpResponseMessage resp = await getResponse;
+
                         Stream respStream = await resp.Content.ReadAsStreamAsync(TestAsync);
+                        using var cts = new CancellationTokenSource();
                         Task readTask = readOrCopyToAsync ?
                             respStream.ReadAsync(new byte[1], 0, 1, cts.Token) :
                             respStream.CopyToAsync(Stream.Null, 10, cts.Token);
                         cts.Cancel();
                         await readTask;
                     });
 
-                    try
-                    {
-                        clientFinished.SetResult(true);
-                        await serverTask;
-                    } catch { }
+                    await serverTask;
                 });
             }
         }