Add WaitForShutdown to WebHostExtensions

src/Microsoft.AspNetCore.Hosting/WebHostExtensions.cs

@@ -23,6 +23,32 @@ namespace Microsoft.AspNetCore.Hosting
             return host.StopAsync(new CancellationTokenSource(timeout).Token);
         }
 
+        /// <summary>
+        /// Block the calling thread until shutdown is triggered via Ctrl+C or SIGTERM.
+        /// </summary>
+        /// <param name="host">The running <see cref="IWebHost"/>.</param>
+        public static void WaitForShutdown(this IWebHost host)
+        {
+            host.WaitForShutdownAsync().GetAwaiter().GetResult();
+        }
+
+        /// <summary>
+        /// Returns a Task that completes when shutdown is triggered via the given token, Ctrl+C or SIGTERM.
+        /// </summary>
+        /// <param name="host">The running <see cref="IWebHost"/>.</param>
+        /// <param name="token">The token to trigger shutdown.</param>
+        public static async Task WaitForShutdownAsync(this IWebHost host, CancellationToken token = default(CancellationToken))
+        {
+            var done = new ManualResetEventSlim(false);
+            using (var cts = CancellationTokenSource.CreateLinkedTokenSource(token))
+            {
+                AttachCtrlcSigtermShutdown(cts, done, shutdownMessage: string.Empty);
+
+                await host.WaitForTokenShutdownAsync(cts.Token);
+                done.Set();
+            }
+        }
+
         /// <summary>
         /// Runs a web application and block the calling thread until host shutdown.
         /// </summary>
@@ -33,54 +59,30 @@ namespace Microsoft.AspNetCore.Hosting
         }
 
         /// <summary>
-        /// Runs a web application and returns a Task that only completes on host shutdown.
+        /// Runs a web application and returns a Task that only completes when the token is triggered or shutdown is triggered.
         /// </summary>
         /// <param name="host">The <see cref="IWebHost"/> to run.</param>
-        public static async Task RunAsync(this IWebHost host)
+        /// <param name="token">The token to trigger shutdown.</param>
+        public static async Task RunAsync(this IWebHost host, CancellationToken token = default(CancellationToken))
         {
+            // Wait for token shutdown if it can be canceled
+            if (token.CanBeCanceled)
+            {
+                await host.RunAsync(token, shutdownMessage: null);
+                return;
+            }
+
+            // If token cannot be canceled, attach Ctrl+C and SIGTERM shutdown
             var done = new ManualResetEventSlim(false);
             using (var cts = new CancellationTokenSource())
             {
-                Action shutdown = () =>
+                AttachCtrlcSigtermShutdown(cts, done, shutdownMessage: "Application is shutting down...");
-                {
-                    if (!cts.IsCancellationRequested)
-                    {
-                        Console.WriteLine("Application is shutting down...");
-                        try
-                        {
-                            cts.Cancel();
-                        }
-                        catch (ObjectDisposedException)
-                        {
-                        }
-                    }
-
-                    done.Wait();
-                };
-
-                AppDomain.CurrentDomain.ProcessExit += (sender, eventArgs) => shutdown();
-                Console.CancelKeyPress += (sender, eventArgs) =>
-                {
-                    shutdown();
-                    // Don't terminate the process immediately, wait for the Main thread to exit gracefully.
-                    eventArgs.Cancel = true;
-                };
 
                 await host.RunAsync(cts.Token, "Application started. Press Ctrl+C to shut down.");
                 done.Set();
             }
         }
 
-        /// <summary>
-        /// Runs a web application and and returns a Task that only completes when the token is triggered or shutdown is triggered.
-        /// </summary>
-        /// <param name="host">The <see cref="IWebHost"/> to run.</param>
-        /// <param name="token">The token to trigger shutdown.</param>
-        public static Task RunAsync(this IWebHost host, CancellationToken token)
-        {
-            return host.RunAsync(token, shutdownMessage: null);
-        }
-
         private static async Task RunAsync(this IWebHost host, CancellationToken token, string shutdownMessage)
         {
             using (host)
@@ -107,24 +109,61 @@ namespace Microsoft.AspNetCore.Hosting
                     Console.WriteLine(shutdownMessage);
                 }
 
-                token.Register(state =>
+                await host.WaitForTokenShutdownAsync(token);
-                {
-                    ((IApplicationLifetime)state).StopApplication();
-                },
-                applicationLifetime);
-
-                var waitForStop = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
-                applicationLifetime.ApplicationStopping.Register(obj =>
-                {
-                    var tcs = (TaskCompletionSource<object>)obj;
-                    tcs.TrySetResult(null);
-                }, waitForStop);
-
-                await waitForStop.Task;
-
-                // WebHost will use its default ShutdownTimeout if none is specified.
-                await host.StopAsync();
             }
         }
+
+        private static void AttachCtrlcSigtermShutdown(CancellationTokenSource cts, ManualResetEventSlim resetEvent, string shutdownMessage)
+        {
+            void Shutdown()
+            {
+                if (!cts.IsCancellationRequested)
+                {
+                    if (!string.IsNullOrEmpty(shutdownMessage))
+                    {
+                        Console.WriteLine(shutdownMessage);
+                    }
+                    try
+                    {
+                        cts.Cancel();
+                    }
+                    catch (ObjectDisposedException) { }
+                }
+
+                // Wait on the given reset event
+                resetEvent.Wait();
+            };
+
+            AppDomain.CurrentDomain.ProcessExit += (sender, eventArgs) => Shutdown();
+            Console.CancelKeyPress += (sender, eventArgs) =>
+            {
+                Shutdown();
+                // Don't terminate the process immediately, wait for the Main thread to exit gracefully.
+                eventArgs.Cancel = true;
+            };
+        }
+
+        private static async Task WaitForTokenShutdownAsync(this IWebHost host, CancellationToken token)
+        {
+            var applicationLifetime = host.Services.GetService<IApplicationLifetime>();
+
+            token.Register(state =>
+            {
+                ((IApplicationLifetime)state).StopApplication();
+            },
+            applicationLifetime);
+
+            var waitForStop = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            applicationLifetime.ApplicationStopping.Register(obj =>
+            {
+                var tcs = (TaskCompletionSource<object>)obj;
+                tcs.TrySetResult(null);
+            }, waitForStop);
+
+            await waitForStop.Task;
+
+            // WebHost will use its default ShutdownTimeout if none is specified.
+            await host.StopAsync();
+        }
     }
 }

test/Microsoft.AspNetCore.Hosting.FunctionalTests/ShutdownTests.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Threading.Tasks;
@@ -22,11 +23,11 @@ namespace Microsoft.AspNetCore.Hosting.FunctionalTests
         [ConditionalFact]
         [OSSkipCondition(OperatingSystems.Windows)]
         [OSSkipCondition(OperatingSystems.MacOSX)]
-        public async Task ShutdownTest()
+        public async Task ShutdownTestRun()
         {
             using (StartLog(out var loggerFactory))
             {
-                var logger = loggerFactory.CreateLogger(nameof(ShutdownTest));
+                var logger = loggerFactory.CreateLogger(nameof(ShutdownTestRun));
 
                 var applicationPath = Path.Combine(TestPathUtilities.GetSolutionRootDirectory("Hosting"), "test",
                     "Microsoft.AspNetCore.Hosting.TestSites");
@@ -43,13 +44,12 @@ namespace Microsoft.AspNetCore.Hosting.FunctionalTests
                     PublishApplicationBeforeDeployment = true
                 };
 
+                deploymentParameters.EnvironmentVariables.Add(new KeyValuePair<string, string>("ASPNETCORE_STARTMECHANIC", "Run"));
+
                 using (var deployer = new SelfHostDeployer(deploymentParameters, loggerFactory))
                 {
                     await deployer.DeployAsync();
 
-                    // Wait for application to start
-                    await Task.Delay(1000);
-
                     string output = string.Empty;
                     deployer.HostProcess.OutputDataReceived += (sender, args) => output += args.Data + '\n';
 
@@ -68,6 +68,53 @@ namespace Microsoft.AspNetCore.Hosting.FunctionalTests
             }
         }
 
+        [ConditionalFact]
+        [OSSkipCondition(OperatingSystems.Windows)]
+        [OSSkipCondition(OperatingSystems.MacOSX)]
+        public async Task ShutdownTestWaitForShutdown()
+        {
+            using (StartLog(out var loggerFactory))
+            {
+                var logger = loggerFactory.CreateLogger(nameof(ShutdownTestWaitForShutdown));
+
+                var applicationPath = Path.Combine(TestPathUtilities.GetSolutionRootDirectory("Hosting"), "test",
+                    "Microsoft.AspNetCore.Hosting.TestSites");
+
+                var deploymentParameters = new DeploymentParameters(
+                    applicationPath,
+                    ServerType.Kestrel,
+                    RuntimeFlavor.CoreClr,
+                    RuntimeArchitecture.x64)
+                {
+                    EnvironmentName = "Shutdown",
+                    TargetFramework = "netcoreapp2.0",
+                    ApplicationType = ApplicationType.Portable,
+                    PublishApplicationBeforeDeployment = true
+                };
+
+                deploymentParameters.EnvironmentVariables.Add(new KeyValuePair<string, string>("ASPNETCORE_STARTMECHANIC", "WaitForShutdown"));
+
+                using (var deployer = new SelfHostDeployer(deploymentParameters, loggerFactory))
+                {
+                    await deployer.DeployAsync();
+
+                    string output = string.Empty;
+                    deployer.HostProcess.OutputDataReceived += (sender, args) => output += args.Data + '\n';
+
+                    SendSIGINT(deployer.HostProcess.Id);
+
+                    WaitForExitOrKill(deployer.HostProcess);
+
+                    output = output.Trim('\n');
+
+                    Assert.Equal(output, "Stopping firing\n" +
+                                         "Stopping end\n" +
+                                         "Stopped firing\n" +
+                                         "Stopped end");
+                }
+            }
+        }
+
 
         private static void SendSIGINT(int processId)
         {

test/Microsoft.AspNetCore.Hosting.TestSites/Microsoft.AspNetCore.Hosting.TestSites.csproj

@@ -15,6 +15,7 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Configuration" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.Extensions.Configuration.CommandLine" Version="$(AspNetCoreVersion)" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.Extensions.Logging.Console" Version="$(AspNetCoreVersion)" />
   </ItemGroup>
 

test/Microsoft.AspNetCore.Hosting.TestSites/Program.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Hosting;
@@ -19,6 +20,7 @@ namespace ServerComparison.TestSites
         {
             var config = new ConfigurationBuilder()
                 .AddCommandLine(args)
+                .AddEnvironmentVariables(prefix: "ASPNETCORE_")
                 .Build();
 
             var builder = new WebHostBuilder()
@@ -31,9 +33,29 @@ namespace ServerComparison.TestSites
                 })
                 .UseStartup("Microsoft.AspNetCore.Hosting.TestSites");
 
-            var host = builder.Build();
+            if (config["STARTMECHANIC"] == "Run")
+            {
+                var host = builder.Build();
 
-            host.Run();
+                host.Run();
+            }
+            else if (config["STARTMECHANIC"] == "WaitForShutdown")
+            {
+                using (var host = builder.Build())
+                {
+                    host.Start();
+
+                    // Mimic application startup messages so application deployer knows that the application has started
+                    Console.WriteLine("Application started. Press Ctrl+C to shut down.");
+                    Console.WriteLine("Now listening on: http://localhost:5000");
+
+                    host.WaitForShutdown();
+                }
+            }
+            else
+            {
+                throw new InvalidOperationException("Starting mechanic not specified");
+            }
         }
     }