Merge pull request #57 from mayuki/feature/ShutdownTimeout

Wait for a timeout before shutting down when canceling with Ctrl+C
mayuki · Jan 10, 2022 · 3f3f3ee · 3f3f3ee
2 parents 8f253aa + f8a2928
commit 3f3f3ee
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 18 deletions.
diff --git a/src/Cocona.Lite/CoconaLiteAppOptions.cs b/src/Cocona.Lite/CoconaLiteAppOptions.cs
@@ -32,5 +32,10 @@ public class CoconaLiteAppOptions
         /// Specify enable shell completion support. The default value is false.
         /// </summary>
         public bool EnableShellCompletionSupport { get; set; } = false;
+
+        /// <summary>
+        /// Specify the timeout before the application is shutdown when Ctrl+C is pressed.
+        /// </summary>
+        public TimeSpan ShutdownTimeout { get; set; } = TimeSpan.FromSeconds(5);
     }
 }
diff --git a/src/Cocona.Lite/Lite/Hosting/CoconaLiteAppHost.cs b/src/Cocona.Lite/Lite/Hosting/CoconaLiteAppHost.cs
@@ -15,12 +15,14 @@ public class CoconaLiteAppHost
         private readonly IServiceProvider _serviceProvider;
         private readonly CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();
         private readonly ManualResetEventSlim _waitForShutdown = new ManualResetEventSlim(false);
+        private readonly TimeSpan _shutdownTimeout;
 
         public IServiceProvider Services => _serviceProvider;
 
-        public CoconaLiteAppHost(IServiceProvider serviceProvider)
+        public CoconaLiteAppHost(IServiceProvider serviceProvider, CoconaLiteAppOptions options)
         {
             _serviceProvider = serviceProvider;
+            _shutdownTimeout = options.ShutdownTimeout;
         }
 
         public async Task RunAsync(CancellationToken cancellationToken)
@@ -37,20 +39,28 @@ public async Task RunAsync(CancellationToken cancellationToken)
 
             try
             {
-                var task = bootstrapper.RunAsync(linkedCancellationToken.Token);
-                if (task.IsCompleted)
-                {
-                    Environment.ExitCode = task.Result;
-                }
-                else
+                var cancellationTask = CreateTaskFromCancellationToken(linkedCancellationToken.Token);
+                var runTask = Task.Run(() => bootstrapper.RunAsync(linkedCancellationToken.Token).AsTask());
+
+                var winTask = await Task.WhenAny(cancellationTask, runTask);
+                if (winTask != runTask)
                 {
-                    Environment.ExitCode = await task;
+                    // Wait for shutdown timeout.
+                    var timeoutToken = new CancellationTokenSource(_shutdownTimeout);
+                    var winTask2 = await Task.WhenAny(runTask, CreateTaskFromCancellationToken(timeoutToken.Token));
+                    if (winTask2 != runTask)
+                    {
+                        // Timed out. (throw OperationCanceledException)
+                        await cancellationTask;
+                    }
                 }
+
+                Environment.ExitCode = await runTask;
             }
-            catch (OperationCanceledException ex) when (ex.CancellationToken == _cancellationTokenSource.Token)
+            catch (OperationCanceledException ex) when (ex.CancellationToken == linkedCancellationToken.Token)
             {
                 // NOTE: Ignore OperationCanceledException that was thrown by non-user code.
-                Environment.ExitCode = 0;
+                Environment.ExitCode = 130;
             }
 
             _waitForShutdown.Set();
@@ -77,5 +87,15 @@ private void OnCancelKeyPress(object? sender, ConsoleCancelEventArgs e)
             e.Cancel = true;
             _cancellationTokenSource.Cancel();
         }
+
+        private static Task CreateTaskFromCancellationToken(CancellationToken cancellationToken)
+        {
+            var tsc = new TaskCompletionSource<bool>();
+            cancellationToken.Register(() =>
+            {
+                tsc.TrySetCanceled(cancellationToken);
+            });
+            return tsc.Task;
+        }
     }
 }
diff --git a/src/Cocona.Lite/Lite/Hosting/CoconaLiteAppHostBuilder.cs b/src/Cocona.Lite/Lite/Hosting/CoconaLiteAppHostBuilder.cs
@@ -100,7 +100,7 @@ public CoconaLiteAppHost Build()
                 .UseMiddleware((next, sp) => new InitializeCoconaLiteConsoleAppMiddleware(next, sp.GetRequiredService<ICoconaAppContextAccessor>()))
                 .UseMiddleware<CoconaCommandInvokeMiddleware>();
 
-            return new CoconaLiteAppHost(serviceProvider);
+            return new CoconaLiteAppHost(serviceProvider, options);
         }
 
         private static string[] GetCommandLineArguments()

diff --git a/src/Cocona/Hosting/CoconaHostedService.cs b/src/Cocona/Hosting/CoconaHostedService.cs
@@ -26,7 +26,6 @@ public class CoconaHostedService : IHostedService
 
         private readonly CancellationTokenSource _cancellationTokenSource;
         private Task? _runningCommandTask;
-        private ExceptionDispatchInfo? _capturedException;
 
         public CoconaHostedService(
             ICoconaConsoleProvider console,
@@ -78,24 +77,44 @@ private async Task ExecuteCoconaApplicationAsync(Task waitForApplicationStarted)
                 // NOTE: Ignore OperationCanceledException that was thrown by non-user code.
                 Environment.ExitCode = 0;
             }
-            catch (Exception ex)
+            catch (Exception)
             {
                 Environment.ExitCode = 1;
-                _capturedException = ExceptionDispatchInfo.Capture(ex);
+                throw;
             }
 
             _lifetime.StopApplication();
         }
 
         public async Task StopAsync(CancellationToken cancellationToken)
         {
-            _cancellationTokenSource?.Cancel();
-
-            _capturedException?.Throw();
+            _cancellationTokenSource.Cancel();
 
             if (_runningCommandTask != null && !_runningCommandTask.IsCompleted)
             {
-                await _runningCommandTask;
+                var cancellationTask = CreateTaskFromCancellationToken(cancellationToken);
+                try
+                {
+                    var winTask = await Task.WhenAny(cancellationTask, _runningCommandTask);
+                    if (winTask == _runningCommandTask)
+                    {
+                        await _runningCommandTask;
+                    }
+                }
+                catch (OperationCanceledException e) when (e.CancellationToken == cancellationToken)
+                {
+                    Environment.ExitCode = 130;
+                }
+            }
+
+            static Task CreateTaskFromCancellationToken(CancellationToken cancellationToken)
+            {
+                var tsc = new TaskCompletionSource<bool>();
+                cancellationToken.Register(() =>
+                {
+                    tsc.TrySetCanceled(cancellationToken);
+                });
+                return tsc.Task;
             }
         }
     }