Fix flaky graceful shutdown tests

test/Kestrel.Transport.FunctionalTests/Http2/ShutdownTests.cs

@@ -3,18 +3,22 @@
 
 #if NETCOREAPP2_2
 
+using System;
 using System.Collections.Generic;
 using System.Net;
 using System.Net.Http;
 using System.Runtime.InteropServices;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.AspNetCore.Testing.xunit;
+using Moq;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
@@ -46,13 +50,22 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
         {
             var requestStarted = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var requestUnblocked = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var requestStopping = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var mockKestrelTrace = new Mock<KestrelTrace>(TestApplicationErrorLogger)
+            {
+                CallBase = true
+            };
+            mockKestrelTrace
+                .Setup(m => m.Http2ConnectionClosing(It.IsAny<string>()))
+                .Callback(() => requestStopping.SetResult(null));
+
             using (var server = new TestServer(async context =>
             {
                 requestStarted.SetResult(null);
                 await requestUnblocked.Task.DefaultTimeout();
                 await context.Response.WriteAsync("hello world " + context.Request.Protocol);
             },
-            new TestServiceContext(LoggerFactory),
+            new TestServiceContext(LoggerFactory, mockKestrelTrace.Object),
             kestrelOptions =>
             {
                 kestrelOptions.Listen(IPAddress.Loopback, 0, listenOptions =>
@@ -69,6 +82,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
 
                 var stopTask = server.StopAsync();
 
+                await requestStopping.Task.DefaultTimeout();
+
                 // Unblock the request
                 requestUnblocked.SetResult(null);
 
@@ -116,7 +131,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
                 Assert.False(requestTask.IsCompleted);
                 await requestStarted.Task.DefaultTimeout();
 
-                await server.StopAsync().DefaultTimeout();
+                await server.StopAsync(new CancellationToken(true)).DefaultTimeout();
             }
 
             Assert.Contains(TestApplicationErrorLogger.Messages, m => m.Message.Contains("is closing."));

test/shared/TransportTestHelpers/TestServer.cs

@@ -6,6 +6,7 @@ using System.Linq;
 using System.Net;
 using System.Net.Sockets;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
@@ -81,7 +82,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                     configureServices(services);
                 })
                 .UseSetting(WebHostDefaults.ApplicationKey, typeof(TestServer).GetTypeInfo().Assembly.FullName)
-                .UseSetting(WebHostDefaults.ShutdownTimeoutKey, "1")
                 .Build();
 
             _host.Start();
@@ -110,9 +110,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             return new TestConnection(Port, AddressFamily);
         }
 
-        public Task StopAsync()
+        public Task StopAsync(CancellationToken token = default)
         {
-            return _host.StopAsync();
+            return _host.StopAsync(token);
         }
 
         public void Dispose()