Implement graceful shutdown timeout period

.vscode/launch.json

@@ -13,7 +13,7 @@
             },
             "args": "${input:funcArgs} --script-root ${input:scriptRootArg}",
             "cwd": "${workspaceFolder}/src/Cli/func",
-            "console": "internalConsole",
+            "console": "integratedTerminal",
             "stopAtEntry": false
         },
         {

release_notes.md

@@ -10,4 +10,5 @@
 - Fix dotnet templates installation (#4538)
 - Disable diagnostic events in local development by replacing the `IDiagnosticEventRepository` with a `DiagnosticEventNullRepository` (#4542)
 - Add `func pack` support for in-proc functions (#4529)
-- Default to remote build for `func pack` for python apps (#4530)
+- Default to remote build for `func pack` for python apps (#4530)
+- Implement (2 second) graceful timeout period for the CLI shutdown (#4540)

src/Cli/func/Actions/HostActions/StartHostAction.cs

@@ -442,22 +442,25 @@ public override async Task RunAsync()
             var runTask = host.RunAsync();
             var hostService = host.Services.GetRequiredService<WebJobsScriptHostService>();
 
-            await hostService.DelayUntilHostReady();
-
-            var scriptHost = hostService.Services.GetRequiredService<IScriptJobHost>();
-            var httpOptions = hostService.Services.GetRequiredService<IOptions<HttpOptions>>();
-            if (scriptHost != null && scriptHost.Functions.Any())
+            if (hostService.State is not ScriptHostState.Stopping && hostService.State is not ScriptHostState.Stopped)
             {
-                // Checking if in Limelight - it should have a `AzureDevSessionsRemoteHostName` value in local.settings.json.
-                var forwardedHttpUrl = _secretsManager.GetSecrets().FirstOrDefault(
-                    s => s.Key.Equals(Constants.AzureDevSessionsRemoteHostName, StringComparison.OrdinalIgnoreCase)).Value;
-                if (forwardedHttpUrl != null)
+                await hostService.DelayUntilHostReady();
+
+                var scriptHost = hostService.Services.GetRequiredService<IScriptJobHost>();
+                var httpOptions = hostService.Services.GetRequiredService<IOptions<HttpOptions>>();
+                if (scriptHost != null && scriptHost.Functions.Any())
                 {
-                    var baseUrl = forwardedHttpUrl.Replace(Constants.AzureDevSessionsPortSuffixPlaceholder, Port.ToString(), StringComparison.OrdinalIgnoreCase);
-                    baseUri = new Uri(baseUrl);
-                }
+                    // Checking if in Limelight - it should have a `AzureDevSessionsRemoteHostName` value in local.settings.json.
+                    var forwardedHttpUrl = _secretsManager.GetSecrets().FirstOrDefault(
+                        s => s.Key.Equals(Constants.AzureDevSessionsRemoteHostName, StringComparison.OrdinalIgnoreCase)).Value;
+                    if (forwardedHttpUrl != null)
+                    {
+                        var baseUrl = forwardedHttpUrl.Replace(Constants.AzureDevSessionsPortSuffixPlaceholder, Port.ToString(), StringComparison.OrdinalIgnoreCase);
+                        baseUri = new Uri(baseUrl);
+                    }
 
-                DisplayFunctionsInfoUtilities.DisplayFunctionsInfo(scriptHost.Functions, httpOptions.Value, baseUri);
+                    DisplayFunctionsInfoUtilities.DisplayFunctionsInfo(scriptHost.Functions, httpOptions.Value, baseUri);
+                }
             }
 
             if (VerboseLogging == null || !VerboseLogging.Value)

src/Cli/func/CancelKeyHandler.cs

@@ -0,0 +1,54 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+using Azure.Functions.Cli.Interfaces;
+
+namespace Azure.Functions.Cli;
+
+internal static class CancelKeyHandler
+{
+    private static readonly TimeSpan _gracefulShutdownPeriod = TimeSpan.FromSeconds(2);
+    private static IProcessManager _processManager = null!;
+    private static Action _onShuttingDown;
+    private static Action _onGracePeriodTimeout;
+    private static bool _registered = false;
+
+    public static void Register(
+        IProcessManager processManager,
+        Action onShuttingDown,
+        Action onGracePeriodTimeout = null)
+    {
+        if (_registered)
+        {
+            return;
+        }
+
+        _processManager = processManager ?? throw new ArgumentNullException(nameof(processManager));
+        _onShuttingDown = onShuttingDown ?? throw new ArgumentNullException(nameof(onShuttingDown));
+        _onGracePeriodTimeout = onGracePeriodTimeout ?? (() => { });
+
+        Console.CancelKeyPress += HandleCancelKeyPress;
+        _registered = true;
+    }
+
+    internal static void HandleCancelKeyPress(object sender, ConsoleCancelEventArgs e)
+    {
+        _processManager.KillChildProcesses();
+        _onShuttingDown?.Invoke();
+
+        _ = Task.Run(async () =>
+        {
+            await Task.Delay(_gracefulShutdownPeriod);
+            _onGracePeriodTimeout?.Invoke();
+        });
+    }
+
+    internal static void Dispose()
+    {
+        if (_registered)
+        {
+            Console.CancelKeyPress -= HandleCancelKeyPress;
+            _registered = false;
+        }
+    }
+}

src/Cli/func/Common/ProcessManager.cs

@@ -26,6 +26,11 @@ public IEnumerable<IProcessInfo> GetProcessesByName(string processName)
                 .Select(p => new ProcessInfo(p));
         }
 
+        public void KillMainProcess()
+        {
+            Process.GetCurrentProcess().Kill();
+        }
+
         public void KillChildProcesses()
         {
             if (_childProcesses == null)

src/Cli/func/ConsoleApp.cs

@@ -43,12 +43,12 @@ internal ConsoleApp(string[] args, Assembly assembly, IContainer container)
             GlobalCoreToolsSettings.Init(container.Resolve<ISecretsManager>(), args);
         }
 
-        public static void Run<T>(string[] args, IContainer container)
+        public static void Run<T>(string[] args, IContainer container, CancellationToken cancellationToken = default)
         {
             Task.Run(() => RunAsync<T>(args, container)).Wait();
         }
 
-        public static async Task RunAsync<T>(string[] args, IContainer container)
+        public static async Task RunAsync<T>(string[] args, IContainer container, CancellationToken cancellationToken = default)
         {
             var stopWatch = Stopwatch.StartNew();
 

src/Cli/func/Interfaces/IProcessManager.cs

@@ -22,5 +22,8 @@ internal interface IProcessManager
 
         // Kill all child processes spawned by the current process.
         internal void KillChildProcesses();
+
+        // Kill the main process.
+        internal void KillMainProcess();
     }
 }

src/Cli/func/Program.cs

@@ -12,9 +12,11 @@ namespace Azure.Functions.Cli
     internal class Program
     {
         private static readonly string[] _versionArgs = ["version", "v"];
+        private static readonly CancellationTokenSource _shuttingDownCts = new CancellationTokenSource();
+        private static readonly CancellationTokenSource _forceShutdownCts = new CancellationTokenSource();
         private static IContainer _container;
 
-        internal static void Main(string[] args)
+        internal static async Task Main(string[] args)
         {
             // Configure console encoding
             ConsoleHelper.ConfigureConsoleOutputEncoding();
@@ -32,14 +34,26 @@ internal static void Main(string[] args)
             SetupGlobalExceptionHandler();
             SetCoreToolsEnvironmentVariables(args);
             _container = InitializeAutofacContainer();
+            var processManager = _container.Resolve<IProcessManager>();
             AppDomain.CurrentDomain.ProcessExit += CurrentDomain_ProcessExit;
 
-            Console.CancelKeyPress += (s, e) =>
-            {
-                _container.Resolve<IProcessManager>()?.KillChildProcesses();
-            };
+            CancelKeyHandler.Register(
+                processManager: processManager,
+                onShuttingDown: _shuttingDownCts.Cancel,
+                onGracePeriodTimeout: _forceShutdownCts.Cancel);
 
-            ConsoleApp.Run<Program>(args, _container);
+            try
+            {
+                await ConsoleApp.RunAsync<Program>(args, _container, _shuttingDownCts.Token).WaitAsync(_forceShutdownCts.Token);
+            }
+            catch (OperationCanceledException)
+            {
+                processManager.KillMainProcess();
+            }
+            catch (Exception ex)
+            {
+                ColoredConsole.WriteLine($"Unexpected error: {ex.Message}");
+            }
         }
 
         private static void CurrentDomain_ProcessExit(object sender, EventArgs e)

test/Cli/Func.UnitTests/CancelKeyHandlerTests.cs

@@ -0,0 +1,52 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+using Autofac;
+using Azure.Functions.Cli;
+using Azure.Functions.Cli.UnitTests.Helpers;
+using Xunit;
+
+public class CancelKeyHandlerTests
+{
+    [Fact]
+    public async Task CtrlC_KillsChildImmediately_AndInvokesGracePeriodTimeout()
+    {
+        // Arrange
+        var ctsShuttingDown = new CancellationTokenSource();
+        var ctsGracePeriod = new CancellationTokenSource();
+        var mockProcessManager = new TestProcessManager();
+
+        CancelKeyHandler.Register(
+            processManager: mockProcessManager,
+            onShuttingDown: ctsShuttingDown.Cancel,
+            onGracePeriodTimeout: ctsGracePeriod.Cancel);
+
+        try
+        {
+            // Act
+            CancelKeyHandler.HandleCancelKeyPress(null, CreateFakeCancelEventArgs());
+
+            // Assert immediate child process kill
+            Assert.True(mockProcessManager.KillChildProcessesCalled);
+            Assert.False(ctsGracePeriod.IsCancellationRequested);
+
+            // Assert graceful timeout is called after ~2 seconds
+            var gracePeriodInvoked = await Task.Run(() =>
+                ctsGracePeriod.Token.WaitHandle.WaitOne(3000)); // Slight buffer
+
+            Assert.True(gracePeriodInvoked, "Grace period timeout was not triggered.");
+        }
+        finally
+        {
+            CancelKeyHandler.Dispose();
+        }
+    }
+
+    internal static ConsoleCancelEventArgs CreateFakeCancelEventArgs()
+    {
+        var constructor = typeof(ConsoleCancelEventArgs)
+            .GetConstructors(System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)[0];
+
+        return (ConsoleCancelEventArgs)constructor.Invoke([ConsoleSpecialKey.ControlC]);
+    }
+}

test/Cli/Func.UnitTests/Helpers/TestProcessManager.cs

@@ -0,0 +1,26 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+using Azure.Functions.Cli.Interfaces;
+
+namespace Azure.Functions.Cli.UnitTests.Helpers;
+
+internal class TestProcessManager : IProcessManager
+{
+    public bool KillChildProcessesCalled { get; private set; }
+
+    public bool KillMainProcessCalled { get; private set; }
+
+    public void KillChildProcesses() => KillChildProcessesCalled = true;
+
+    public void KillMainProcess() => KillMainProcessCalled = true;
+
+    // Other interface members omitted for this test
+    public IEnumerable<IProcessInfo> GetProcessesByName(string processName) => null;
+
+    public IProcessInfo GetCurrentProcess() => null;
+
+    public IProcessInfo GetProcessById(int processId) => null;
+
+    public bool RegisterChildProcess(System.Diagnostics.Process childProcess) => false;
+}