PR feedback

Author: JamesNK

src/Servers/Kestrel/Core/src/Internal/Http3/Http3Connection.cs

@@ -46,6 +46,7 @@ internal class Http3Connection : IHttp3StreamLifetimeHandler, IRequestProcessor
         private int _stoppedAcceptingStreams;
         private bool _gracefulCloseStarted;
         private int _activeRequestCount;
+        private CancellationTokenSource _acceptStreamsCts = new CancellationTokenSource();
         private readonly Http3PeerSettings _serverSettings = new Http3PeerSettings();
         private readonly Http3PeerSettings _clientSettings = new Http3PeerSettings();
         private readonly StreamCloseAwaitable _streamCompletionAwaitable = new StreamCloseAwaitable();
@@ -107,7 +108,8 @@ public void StopProcessingNextRequest(bool serverInitiated)
 
                 if (Interlocked.CompareExchange(ref _gracefulCloseInitiator, initiator, GracefulCloseInitiator.None) == GracefulCloseInitiator.None)
                 {
-                    UpdateConnectionState();
+                    // Break out of AcceptStreams so connection state can be updated.
+                    _acceptStreamsCts.Cancel();
                 }
             }
         }
@@ -240,15 +242,20 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                 // Don't delay on waiting to send outbound control stream settings.
                 outboundControlStreamTask = ProcessOutboundControlStreamAsync(outboundControlStream);
 
-                while (true)
+                while (_stoppedAcceptingStreams == 0)
                 {
-                    var streamContext = await _multiplexedContext.AcceptAsync();
+                    var streamContext = await _multiplexedContext.AcceptAsync(_acceptStreamsCts.Token);
 
                     try
                     {
                         if (streamContext == null)
                         {
-                            break;
+                            if (_acceptStreamsCts.Token.IsCancellationRequested)
+                            {
+                                _acceptStreamsCts = new CancellationTokenSource();
+                            }
+
+                            continue;
                         }
 
                         var streamDirectionFeature = streamContext.Features.Get<IStreamDirectionFeature>();
@@ -267,9 +274,10 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                         }
                         else
                         {
-                            // TODO race condition between checking this and updating highest stream ID
+                            // Request stream
+
                             // https://quicwg.org/base-drafts/draft-ietf-quic-http.html#section-5.2-2
-                            if (_stoppedAcceptingStreams == 1)
+                            if (_gracefulCloseStarted)
                             {
                                 // https://quicwg.org/base-drafts/draft-ietf-quic-http.html#section-4.1.2-3
                                 streamContext.Features.Get<IProtocolErrorCodeFeature>()!.Error = (long)Http3ErrorCode.RequestRejected;
@@ -280,7 +288,6 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                             // Request stream IDs are tracked.
                             UpdateHighestOpenedRequestStreamId(streamIdFeature.StreamId);
 
-                            // Request stream
                             var persistentStateFeature = streamContext.Features.Get<IPersistentStateFeature>();
                             Debug.Assert(persistentStateFeature != null, $"Required {nameof(IPersistentStateFeature)} not on stream context.");
 
@@ -353,7 +360,10 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                     // Only send goaway if the connection close was initiated on the server.
                     if (!clientAbort)
                     {
-                        await SendGoAwayAsync(GetCurrentGoAwayStreamId());
+                        if (TryStopAcceptingStreams() || _gracefulCloseStarted)
+                        {
+                            await SendGoAwayAsync(GetCurrentGoAwayStreamId());
+                        }
                     }
 
                     // Abort active request streams.
@@ -430,37 +440,31 @@ private void UpdateConnectionState()
                 return;
             }
 
-            int activeRequestCount;
-            lock (_streams)
+            if (_gracefulCloseInitiator != GracefulCloseInitiator.None)
             {
-                activeRequestCount = _activeRequestCount;
-            }
+                int activeRequestCount;
+                lock (_streams)
+                {
+                    activeRequestCount = _activeRequestCount;
+                }
 
-            if (_gracefulCloseInitiator != GracefulCloseInitiator.None && !_gracefulCloseStarted)
-            {
-                _gracefulCloseStarted = true;
+                if (!_gracefulCloseStarted)
+                {
+                    _gracefulCloseStarted = true;
 
-                _errorCodeFeature.Error = (long)Http3ErrorCode.NoError;
-                Log.Http3ConnectionClosing(_context.ConnectionId);
+                    _errorCodeFeature.Error = (long)Http3ErrorCode.NoError;
+                    Log.Http3ConnectionClosing(_context.ConnectionId);
 
-                if (_gracefulCloseInitiator == GracefulCloseInitiator.Server && activeRequestCount > 0)
-                {
-                    if (TryStopAcceptingStreams())
+                    if (_gracefulCloseInitiator == GracefulCloseInitiator.Server && activeRequestCount > 0)
                     {
                         // Go away with largest streamid to initiate graceful shutdown.
                         SendGoAwayAsync(VariableLengthIntegerHelper.EightByteLimit).Preserve();
                     }
                 }
-            }
 
-            if (_activeRequestCount == 0)
-            {
-                if (_gracefulCloseStarted)
+                if (activeRequestCount == 0)
                 {
-                    if (TryStopAcceptingStreams())
-                    {
-                        SendGoAwayAsync(GetCurrentGoAwayStreamId()).Preserve();
-                    }
+                    TryStopAcceptingStreams();
                 }
             }
         }

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.cs

@@ -327,6 +327,17 @@ private async Task DoSend()
 
                 _clientAbort = true;
             }
+            catch (QuicConnectionAbortedException ex)
+            {
+                // Abort from peer.
+                Error = ex.ErrorCode;
+                _log.StreamAborted(this, ex.ErrorCode, ex);
+
+                // This could be ignored if _shutdownReason is already set.
+                shutdownReason = new ConnectionResetException(ex.Message, ex);
+
+                _clientAbort = true;
+            }
             catch (QuicOperationAbortedException ex)
             {
                 // AbortWrite has been called for the stream.

src/Servers/Kestrel/Transport.Quic/test/QuicConnectionContextTests.cs

@@ -24,6 +24,48 @@ public class QuicConnectionContextTests : TestApplicationErrorLoggerLoggedTest
     {
         private static readonly byte[] TestData = Encoding.UTF8.GetBytes("Hello world");
 
+        [ConditionalFact]
+        [MsQuicSupported]
+        public async Task AcceptAsync_CancellationThenAccept_AcceptStreamAfterCancellation()
+        {
+            // Arrange
+            var connectionClosedTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            await using var connectionListener = await QuicTestHelpers.CreateConnectionListenerFactory(LoggerFactory);
+
+            // Act
+            var acceptTask = connectionListener.AcceptAndAddFeatureAsync().DefaultTimeout();
+
+            var options = QuicTestHelpers.CreateClientConnectionOptions(connectionListener.EndPoint);
+
+            using var clientConnection = new QuicConnection(QuicImplementationProviders.MsQuic, options);
+            await clientConnection.ConnectAsync().DefaultTimeout();
+
+            await using var serverConnection = await acceptTask.DefaultTimeout();
+
+            // Wait for stream and then cancel
+            var cts = new CancellationTokenSource();
+            var acceptStreamTask = serverConnection.AcceptAsync(cts.Token);
+            cts.Cancel();
+
+            var serverStream = await acceptStreamTask.DefaultTimeout();
+            Assert.Null(serverStream);
+
+            // Wait for stream after cancellation
+            acceptStreamTask = serverConnection.AcceptAsync();
+
+            await using var clientStream = clientConnection.OpenBidirectionalStream();
+            await clientStream.WriteAsync(TestData);
+
+            // Assert
+            serverStream = await acceptStreamTask.DefaultTimeout();
+            Assert.NotNull(serverStream);
+
+            var read = await serverStream.Transport.Input.ReadAtLeastAsync(TestData.Length).DefaultTimeout();
+            Assert.Equal(TestData, read.Buffer.ToArray());
+            serverStream.Transport.Input.AdvanceTo(read.Buffer.End);
+        }
+
         [ConditionalFact]
         [MsQuicSupported]
         public async Task AcceptAsync_ClientClosesConnection_ServerNotified()
@@ -495,10 +537,6 @@ public async Task StreamPool_ManyConcurrentStreams_StreamPoolFull()
             const int StreamsSent = 101;
             for (var i = 0; i < StreamsSent; i++)
             {
-                // TODO: Race condition in QUIC library.
-                // Delay between sending streams to avoid
-                // https://github.com/dotnet/runtime/issues/55249
-                await Task.Delay(100);
                 streamTasks.Add(SendStream(requestState));
             }
 

src/Servers/Kestrel/test/Interop.FunctionalTests/Http3/Http3RequestTests.cs

@@ -219,6 +219,48 @@ public async Task GET_ServerStreaming_ClientReadsPartialResponse(HttpProtocols p
             }
         }
 
+        [ConditionalTheory]
+        [MsQuicSupported]
+        [InlineData(HttpProtocols.Http3, Skip = "https://github.com/dotnet/runtime/issues/56969")]
+        [InlineData(HttpProtocols.Http2)]
+        public async Task POST_ClientSendsOnlyHeaders_RequestReceivedOnServer(HttpProtocols protocol)
+        {
+            // Arrange
+            var builder = CreateHostBuilder(context =>
+            {
+                return Task.CompletedTask;
+            }, protocol: protocol);
+
+            using (var host = builder.Build())
+            using (var client = CreateClient())
+            {
+                await host.StartAsync();
+
+                var requestContent = new StreamingHttpContext();
+
+                var request = new HttpRequestMessage(HttpMethod.Post, $"https://127.0.0.1:{host.GetPort()}/");
+                request.Content = requestContent;
+                request.Version = GetProtocol(protocol);
+                request.VersionPolicy = HttpVersionPolicy.RequestVersionExact;
+
+                // Act
+                var responseTask = client.SendAsync(request, CancellationToken.None).DefaultTimeout();
+
+                var requestStream = await requestContent.GetStreamAsync().DefaultTimeout();
+
+                // Send headers
+                await requestStream.FlushAsync().DefaultTimeout();
+
+                var response = await responseTask.DefaultTimeout();
+
+                // Assert
+                response.EnsureSuccessStatusCode();
+                Assert.Equal(GetProtocol(protocol), response.Version);
+
+                await host.StopAsync();
+            }
+        }
+
         [ConditionalFact]
         [MsQuicSupported]
         public async Task POST_ServerCompletesWithoutReadingRequestBody_ClientGetsResponse()
@@ -937,7 +979,8 @@ private async Task WaitForLogAsync(Func<IEnumerable<WriteContext>, bool> testLog
         {
             Logger.LogInformation($"Started waiting for logs: {message}");
 
-            for (int i = 0; i < 5; i++)
+            var retryCount = !Debugger.IsAttached ? 5 : int.MaxValue;
+            for (var i = 0; i < retryCount; i++)
             {
                 if (testLogs(TestSink.Writes))
                 {