Detatch CancelKeyPress and ProcessExit event handlers when IWebHost.RunAsync is completed. (#6638)

Author: CBaud

src/Hosting/Hosting/src/Internal/WebHostLifetime.cs

@@ -0,0 +1,70 @@
+// 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;
+
+namespace Microsoft.AspNetCore.Hosting.Internal
+{
+    internal class WebHostLifetime : IDisposable
+    {
+        private readonly CancellationTokenSource _cts;
+        private readonly ManualResetEventSlim _resetEvent;
+        private readonly string _shutdownMessage;
+
+        private bool _disposed = false;
+
+        public WebHostLifetime(CancellationTokenSource cts, ManualResetEventSlim resetEvent, string shutdownMessage)
+        {
+            _cts = cts;
+            _resetEvent = resetEvent;
+            _shutdownMessage = shutdownMessage;
+
+            AppDomain.CurrentDomain.ProcessExit += ProcessExit;
+            Console.CancelKeyPress += CancelKeyPress;
+        }
+
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            _disposed = true;
+            AppDomain.CurrentDomain.ProcessExit -= ProcessExit;
+            Console.CancelKeyPress -= CancelKeyPress;
+        }
+
+        private void CancelKeyPress(object sender, ConsoleCancelEventArgs eventArgs)
+        {
+            Shutdown();
+            // Don't terminate the process immediately, wait for the Main thread to exit gracefully.
+            eventArgs.Cancel = true;
+        }
+
+        private void ProcessExit(object sender, EventArgs eventArgs)
+        {
+            Shutdown();
+        }
+
+        private 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();
+        }
+    }
+}

src/Hosting/Hosting/src/WebHostExtensions.cs

@@ -43,15 +43,16 @@ public static async Task WaitForShutdownAsync(this IWebHost host, CancellationTo
             var done = new ManualResetEventSlim(false);
             using (var cts = CancellationTokenSource.CreateLinkedTokenSource(token))
             {
-                AttachCtrlcSigtermShutdown(cts, done, shutdownMessage: string.Empty);
-
-                try
-                {
-                    await host.WaitForTokenShutdownAsync(cts.Token);
-                }
-                finally
+                using (var lifetime = new WebHostLifetime(cts, done, shutdownMessage: string.Empty))
                 {
-                    done.Set();
+                    try
+                    {
+                        await host.WaitForTokenShutdownAsync(cts.Token);
+                    }
+                    finally
+                    {
+                        done.Set();
+                    }
                 }
             }
         }
@@ -84,15 +85,16 @@ public static async Task RunAsync(this IWebHost host, CancellationToken token =
             using (var cts = new CancellationTokenSource())
             {
                 var shutdownMessage = host.Services.GetRequiredService<WebHostOptions>().SuppressStatusMessages ? string.Empty : "Application is shutting down...";
-                AttachCtrlcSigtermShutdown(cts, done, shutdownMessage: shutdownMessage);
-
-                try
-                {
-                    await host.RunAsync(cts.Token, "Application started. Press Ctrl+C to shut down.");
-                }
-                finally
+                using (var lifetime = new WebHostLifetime(cts, done, shutdownMessage: shutdownMessage))
                 {
-                    done.Set();
+                    try
+                    {
+                        await host.RunAsync(cts.Token, "Application started. Press Ctrl+C to shut down.");
+                    }
+                    finally
+                    {
+                        done.Set();
+                    }
                 }
             }
         }
@@ -131,36 +133,6 @@ private static async Task RunAsync(this IWebHost host, CancellationToken token,
             }
         }
 
-        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>();
@@ -184,4 +156,4 @@ private static async Task WaitForTokenShutdownAsync(this IWebHost host, Cancella
             await host.StopAsync();
         }
     }
-}
+}