From 0673acedc4de48fc86192ffbbb59d2a6b86bc3c3 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Fri, 15 Jan 2016 11:20:57 +0000
Subject: [PATCH] Added overload of Run that triggers shutdown on a
 cancellation token - Added test - Changed NotifyStopped() after all work has
 been done

---
 .../Internal/WebApplication.cs                |  2 +-
 .../WebApplicationBuilderExtensions.cs        | 42 +++++++++++++++----
 .../WebApplicationTests.cs                    | 32 +++++++++++++-
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/src/Microsoft.AspNet.Hosting/Internal/WebApplication.cs b/src/Microsoft.AspNet.Hosting/Internal/WebApplication.cs
index 74e95c4bc5..066ec19253 100644
--- a/src/Microsoft.AspNet.Hosting/Internal/WebApplication.cs
+++ b/src/Microsoft.AspNet.Hosting/Internal/WebApplication.cs
@@ -233,8 +233,8 @@ namespace Microsoft.AspNet.Hosting.Internal
             _logger?.Shutdown();
             _applicationLifetime.StopApplication();
             Server?.Dispose();
-            _applicationLifetime.NotifyStopped();
             (_applicationServices as IDisposable)?.Dispose();
+            _applicationLifetime.NotifyStopped();
         }
 
         private class Disposable : IDisposable
diff --git a/src/Microsoft.AspNet.Hosting/WebApplicationBuilderExtensions.cs b/src/Microsoft.AspNet.Hosting/WebApplicationBuilderExtensions.cs
index e4fd487507..abc9da77a3 100644
--- a/src/Microsoft.AspNet.Hosting/WebApplicationBuilderExtensions.cs
+++ b/src/Microsoft.AspNet.Hosting/WebApplicationBuilderExtensions.cs
@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Threading;
 using Microsoft.AspNet.Hosting.Server;
 using Microsoft.AspNet.Server.Features;
 using Microsoft.Extensions.Configuration;
@@ -105,6 +106,32 @@ namespace Microsoft.AspNet.Hosting
         /// </summary>
         /// <param name="application"></param>
         public static void Run(this IWebApplication application)
+        {
+            using (var cts = new CancellationTokenSource())
+            {
+                Console.CancelKeyPress += (sender, eventArgs) =>
+                {
+                    cts.Cancel();
+
+                    // Don't terminate the process immediately, wait for the Main thread to exit gracefully.
+                    eventArgs.Cancel = true;
+                };
+
+                application.Run(cts.Token, "Application started. Press Ctrl+C to shut down.");
+            }
+        }
+
+        /// <summary>
+        /// Runs a web application and block the calling thread until token is triggered or shutdown is triggered
+        /// </summary>
+        /// <param name="application"></param>
+        /// <param name="token">The token to trigger shutdown</param>
+        public static void Run(this IWebApplication application, CancellationToken token)
+        {
+            application.Run(token, shutdownMessage: null);
+        }
+
+        private static void Run(this IWebApplication application, CancellationToken token, string shutdownMessage)
         {
             using (application)
             {
@@ -124,15 +151,16 @@ namespace Microsoft.AspNet.Hosting
                     }
                 }
 
-                Console.WriteLine("Application started. Press Ctrl+C to shut down.");
-
-                Console.CancelKeyPress += (sender, eventArgs) =>
+                if (!string.IsNullOrEmpty(shutdownMessage))
                 {
-                    applicationLifetime.StopApplication();
+                    Console.WriteLine(shutdownMessage);
+                }
 
-                    // Don't terminate the process immediately, wait for the Main thread to exit gracefully.
-                    eventArgs.Cancel = true;
-                };
+                token.Register(state =>
+                {
+                    ((IApplicationLifetime)state).StopApplication();
+                },
+                applicationLifetime);
 
                 applicationLifetime.ApplicationStopping.WaitHandle.WaitOne();
             }
diff --git a/test/Microsoft.AspNet.Hosting.Tests/WebApplicationTests.cs b/test/Microsoft.AspNet.Hosting.Tests/WebApplicationTests.cs
index a83818a9cf..a759535009 100644
--- a/test/Microsoft.AspNet.Hosting.Tests/WebApplicationTests.cs
+++ b/test/Microsoft.AspNet.Hosting.Tests/WebApplicationTests.cs
@@ -6,6 +6,7 @@ using System.Collections;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNet.Builder;
 using Microsoft.AspNet.Hosting.Fakes;
@@ -30,7 +31,8 @@ namespace Microsoft.AspNet.Hosting
         private IFeatureCollection _featuresSupportedByThisHost = NewFeatureCollection();
         private IFeatureCollection _instanceFeaturesSupportedByThisHost;
 
-        public IFeatureCollection Features {
+        public IFeatureCollection Features
+        {
             get
             {
                 var features = new FeatureCollection();
@@ -158,6 +160,34 @@ namespace Microsoft.AspNet.Hosting
             Assert.Equal(1, _startInstances[0].DisposeCalls);
         }
 
+        [Fact]
+        public void WebApplicationShutsDownWhenTokenTriggers()
+        {
+            var app = CreateBuilder()
+                .UseServer((IServerFactory)this)
+                .UseStartup("Microsoft.AspNet.Hosting.Tests")
+                .Build();
+
+            var lifetime = app.Services.GetRequiredService<IApplicationLifetime>();
+
+            var cts = new CancellationTokenSource();
+
+            Task.Run(() => app.Run(cts.Token));
+
+            // Wait on the app to be started
+            lifetime.ApplicationStarted.WaitHandle.WaitOne();
+
+            Assert.Equal(1, _startInstances.Count);
+            Assert.Equal(0, _startInstances[0].DisposeCalls);
+
+            cts.Cancel();
+
+            // Wait on the app to shutdown
+            lifetime.ApplicationStopped.WaitHandle.WaitOne();
+
+            Assert.Equal(1, _startInstances[0].DisposeCalls);
+        }
+
         [Fact]
         public void WebApplicationDisposesServiceProvider()
         {