Fix bug in graceful shutdown

src/Kestrel.Core/Internal/Http2/Http2Connection.cs

@@ -910,8 +910,12 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
                         Input.CancelPendingRead();
                     }
 
-                    // Complete the task waiting on all streams to finish
+
-                    _streamsCompleted.TrySetResult(null);
+                    if (_state != Http2ConnectionState.Open)
+                    {
+                        // Complete the task waiting on all streams to finish
+                        _streamsCompleted.TrySetResult(null);
+                    }
                 }
             }
         }

test/Kestrel.Core.Tests/Http2ConnectionTests.cs

@@ -3324,6 +3324,54 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             VerifyGoAway(await ReceiveFrameAsync(), 0, Http2ErrorCode.NO_ERROR);
         }
 
+        [Fact]
+        public async Task StopProcessingNextRequestSendsGracefulGOAWAYAndWaitsForStreamsToComplete()
+        {
+            var task = Task.CompletedTask;
+            await InitializeConnectionAsync(context => task);
+
+            // Send and receive an unblocked request
+            await StartStreamAsync(1, _browserRequestHeaders, endStream: true);
+
+            await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 55,
+                withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
+                withStreamId: 1);
+            await ExpectAsync(Http2FrameType.DATA,
+                withLength: 0,
+                withFlags: (byte)Http2DataFrameFlags.END_STREAM,
+                withStreamId: 1);
+
+            // Send a blocked request
+            var tcs = new TaskCompletionSource<object>(TaskContinuationOptions.RunContinuationsAsynchronously);
+            task = tcs.Task;
+            await StartStreamAsync(3, _browserRequestHeaders, endStream: false);
+
+            // Close pipe
+            _pair.Application.Output.Complete();
+
+            // Assert connection closed
+            await _closedStateReached.Task.DefaultTimeout();
+            VerifyGoAway(await ReceiveFrameAsync(), 3, Http2ErrorCode.NO_ERROR);
+
+            // Assert connection shutdown is still blocked
+            // ProcessRequestsAsync completes the connection's Input pipe
+            var readTask = _pair.Application.Input.ReadAsync();
+            _pair.Application.Input.CancelPendingRead();
+            var result = await readTask;
+            Assert.False(result.IsCompleted);
+
+            // Unblock the request and ProcessRequestsAsync
+            tcs.TrySetResult(null);
+            await _connectionTask;
+
+            // Assert connection's Input pipe is completed
+            readTask = _pair.Application.Input.ReadAsync();
+            _pair.Application.Input.CancelPendingRead();
+            result = await readTask;
+            Assert.True(result.IsCompleted);
+        }
+
         [Fact]
         public async Task StopProcessingNextRequestSendsGracefulGOAWAYThenFinalGOAWAYWhenAllStreamsComplete()
         {

test/Kestrel.Core.Tests/Http2StreamTests.cs

@@ -964,12 +964,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         [Fact]
         public async Task ContentLength_Received_SingleDataFrameUnderSize_Reset()
         {
-            // I hate doing this, but it avoids exceptions from MemoryPool.Dipose() in debug mode. The problem is since
-            // the stream's ProcessRequestsAsync loop is never awaited by the connection, it's not really possible to
-            // observe when all the blocks are returned. This can be removed after we implement graceful shutdown.
-            Dispose();
-            InitializeConnectionFields(new DiagnosticMemoryPool(KestrelMemoryPool.CreateSlabMemoryPool(), allowLateReturn: true));
-
             var headers = new[]
             {
                 new KeyValuePair<string, string>(HeaderNames.Method, "POST"),
@@ -997,12 +991,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         [Fact]
         public async Task ContentLength_Received_MultipleDataFramesOverSize_Reset()
         {
-            // I hate doing this, but it avoids exceptions from MemoryPool.Dipose() in debug mode. The problem is since
-            // the stream's ProcessRequestsAsync loop is never awaited by the connection, it's not really possible to
-            // observe when all the blocks are returned. This can be removed after we implement graceful shutdown.
-            Dispose();
-            InitializeConnectionFields(new DiagnosticMemoryPool(KestrelMemoryPool.CreateSlabMemoryPool(), allowLateReturn: true));
-
             var headers = new[]
             {
                 new KeyValuePair<string, string>(HeaderNames.Method, "POST"),
@@ -1033,12 +1021,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         [Fact]
         public async Task ContentLength_Received_MultipleDataFramesUnderSize_Reset()
         {
-            // I hate doing this, but it avoids exceptions from MemoryPool.Dipose() in debug mode. The problem is since
-            // the stream's ProcessRequestsAsync loop is never awaited by the connection, it's not really possible to
-            // observe when all the blocks are returned. This can be removed after we implement graceful shutdown.
-            Dispose();
-            InitializeConnectionFields(new DiagnosticMemoryPool(KestrelMemoryPool.CreateSlabMemoryPool(), allowLateReturn: true));
-
             var headers = new[]
             {
                 new KeyValuePair<string, string>(HeaderNames.Method, "POST"),