Chain hosting timeout to StopAsync cancellation token (#1078)

build/dependencies.props

@@ -8,6 +8,7 @@
     <TestSdkVersion>15.3.0-*</TestSdkVersion>
     <WindowsApiSetsVersion>1.0.1</WindowsApiSetsVersion>
     <XunitVersion>2.3.0-beta2-*</XunitVersion>
+    <MoqVersion>4.7.1</MoqVersion>
     <SerilogExtensionsLoggingVersion>1.4.0</SerilogExtensionsLoggingVersion>
     <SerilogFileSinkVersion>3.2.0</SerilogFileSinkVersion>
   </PropertyGroup>

src/Microsoft.AspNetCore.Hosting/Internal/WebHost.cs

@@ -277,9 +277,14 @@ namespace Microsoft.AspNetCore.Hosting.Internal
 
             _logger?.Shutdown();
 
+            var timeoutToken = new CancellationTokenSource(Options.ShutdownTimeout).Token;
             if (!cancellationToken.CanBeCanceled)
             {
-                cancellationToken = new CancellationTokenSource(Options.ShutdownTimeout).Token;
+                cancellationToken = timeoutToken;
+            }
+            else
+            {
+                cancellationToken = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutToken).Token;
             }
 
             // Fire IApplicationLifetime.Stopping

test/Microsoft.AspNetCore.Hosting.Tests/Microsoft.AspNetCore.Hosting.Tests.csproj

@@ -22,6 +22,7 @@
     <PackageReference Include="Microsoft.Extensions.Logging.Testing" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.Extensions.Options" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="$(TestSdkVersion)" />
+    <PackageReference Include="Moq" Version="$(MoqVersion)" />
     <PackageReference Include="xunit.runner.visualstudio" Version="$(XunitVersion)" />
     <PackageReference Include="xunit" Version="$(XunitVersion)" />
   </ItemGroup>

test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs

@@ -20,6 +20,7 @@ using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Microsoft.Extensions.Primitives;
+using Moq;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Hosting
@@ -187,6 +188,114 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        [Fact]
+        public async Task WebHostStopAsyncUsesDefaultTimeoutIfGivenTokenDoesNotFire()
+        {
+            var data = new Dictionary<string, string>
+            {
+                { WebHostDefaults.ShutdownTimeoutKey, "1" }
+            };
+
+            var config = new ConfigurationBuilder().AddInMemoryCollection(data).Build();
+
+            var server = new Mock<IServer>();
+            server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask)
+                .Callback<CancellationToken>(token =>
+                {
+                    token.WaitHandle.WaitOne();
+                });
+
+            using (var host = CreateBuilder(config)
+                .ConfigureServices(services =>
+                {
+                    services.AddSingleton(server.Object);
+                })
+                .UseStartup("Microsoft.AspNetCore.Hosting.Tests")
+                .Build())
+            {
+                await host.StartAsync();
+
+                var cts = new CancellationTokenSource();
+
+                // Purposefully don't trigger cts
+                var task = host.StopAsync(cts.Token);
+
+                Assert.Equal(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(10))));
+            }
+        }
+
+        [Fact]
+        public async Task WebHostStopAsyncUsesDefaultTimeoutIfNoTokenProvided()
+        {
+            var data = new Dictionary<string, string>
+            {
+                { WebHostDefaults.ShutdownTimeoutKey, "1" }
+            };
+
+            var config = new ConfigurationBuilder().AddInMemoryCollection(data).Build();
+
+            var server = new Mock<IServer>();
+            server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask)
+                .Callback<CancellationToken>(token =>
+                {
+                    token.WaitHandle.WaitOne();
+                });
+
+            using (var host = CreateBuilder(config)
+                .ConfigureServices(services =>
+                {
+                    services.AddSingleton(server.Object);
+                })
+                .UseStartup("Microsoft.AspNetCore.Hosting.Tests")
+                .Build())
+            {
+                await host.StartAsync();
+
+                var task = host.StopAsync();
+
+                Assert.Equal(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(10))));
+            }
+        }
+
+        [Fact]
+        public async Task WebHostStopAsyncCanBeCancelledEarly()
+        {
+            var data = new Dictionary<string, string>
+            {
+                { WebHostDefaults.ShutdownTimeoutKey, "10" }
+            };
+
+            var config = new ConfigurationBuilder().AddInMemoryCollection(data).Build();
+
+            var server = new Mock<IServer>();
+            server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask)
+                .Callback<CancellationToken>(token =>
+                {
+                    token.WaitHandle.WaitOne();
+                });
+
+            using (var host = CreateBuilder(config)
+                .ConfigureServices(services =>
+                {
+                    services.AddSingleton(server.Object);
+                })
+                .UseStartup("Microsoft.AspNetCore.Hosting.Tests")
+                .Build())
+            {
+                await host.StartAsync();
+
+                var cts = new CancellationTokenSource();
+
+                var task = host.StopAsync(cts.Token);
+                cts.Cancel();
+
+                Assert.Equal(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(8))));
+            }
+        }
+
         [Fact]
         public void WebHostApplicationLifetimeEventsOrderedCorrectlyDuringShutdown()
         {