Exit worker nodes on Ctrl+C / Restart Manager shutdown (#12776)

src/Build/BackEnd/Node/OutOfProcNode.cs

@@ -33,6 +33,9 @@ namespace Microsoft.Build.Execution
     /// </summary>
     public class OutOfProcNode : INode, IBuildComponentHost, INodePacketFactory, INodePacketHandler
     {
+        private static readonly object s_activeNodeLock = new();
+        private static OutOfProcNode s_activeNode;
+
         /// <summary>
         /// Whether the current appdomain has an out of proc node.
         /// For diagnostics.
@@ -254,35 +257,75 @@ public NodeEngineShutdownReason Run(bool enableReuse, bool lowPriority, out Exce
             _nodeEndpoint.OnLinkStatusChanged += OnLinkStatusChanged;
             _nodeEndpoint.Listen(this);
 
-            WaitHandle[] waitHandles = [_shutdownEvent, _packetReceivedEvent];
+            lock (s_activeNodeLock)
+            {
+                s_activeNode = this;
+            }
 
-            // Get the current directory before doing any work. We need this so we can restore the directory when the node shutsdown.
-            while (true)
+            try
             {
-                int index = WaitHandle.WaitAny(waitHandles);
-                switch (index)
+                WaitHandle[] waitHandles = [_shutdownEvent, _packetReceivedEvent];
+
+                // Get the current directory before doing any work. We need this so we can restore the directory when the node shutsdown.
+                while (true)
                 {
-                    case 0:
-                        NodeEngineShutdownReason shutdownReason = HandleShutdown(out shutdownException);
-                        return shutdownReason;
+                    int index = WaitHandle.WaitAny(waitHandles);
+                    switch (index)
+                    {
+                        case 0:
+                            NodeEngineShutdownReason shutdownReason = HandleShutdown(out shutdownException);
+                            return shutdownReason;
 
-                    case 1:
+                        case 1:
 
-                        while (_receivedPackets.TryDequeue(out INodePacket packet))
-                        {
-                            if (packet != null)
+                            while (_receivedPackets.TryDequeue(out INodePacket packet))
                             {
-                                HandlePacket(packet);
+                                if (packet != null)
+                                {
+                                    HandlePacket(packet);
+                                }
                             }
-                        }
 
-                        break;
+                            break;
+                    }
+                }
+            }
+            finally
+            {
+                lock (s_activeNodeLock)
+                {
+                    if (s_activeNode == this)
+                    {
+                        s_activeNode = null;
+                    }
                 }
             }
 
             // UNREACHABLE
         }
 
+        /// <summary>
+        /// Requests the active out-of-proc worker node (if any) to shut down in response to an external stop signal.
+        /// Windows Restart Manager sends <c>CTRL_C_EVENT</c> to console processes; the MSBuild console handler must exit
+        /// worker nodes that never set the main-build "started" flag used by the MSBuild console entry point.
+        /// </summary>
+        public static void RequestExternalShutdown()
+        {
+            OutOfProcNode node;
+            lock (s_activeNodeLock)
+            {
+                node = s_activeNode;
+            }
+
+            node?.SignalShutdownFromExternalStopRequest();
+        }
+
+        private void SignalShutdownFromExternalStopRequest()
+        {
+            _shutdownReason = NodeEngineShutdownReason.BuildComplete;
+            _shutdownEvent.Set();
+        }
+
         #endregion
 
         #region IBuildComponentHost Members

src/MSBuild/OutOfProcTaskHostNode.cs

@@ -33,6 +33,9 @@ internal class OutOfProcTaskHostNode :
 #endif
         INodePacketFactory, INodePacketHandler, IBuildEngine10
     {
+        private static readonly object s_activeInstanceLock = new();
+        private static OutOfProcTaskHostNode s_activeInstance;
+
         /// <summary>
         /// Keeps a record of all environment variables that, on startup of the task host, have a different
         /// value from those that are passed to the task host in the configuration packet for the first task.
@@ -703,54 +706,94 @@ public NodeEngineShutdownReason Run(out Exception shutdownException, bool nodeRe
             _nodeEndpoint.OnLinkStatusChanged += new LinkStatusChangedDelegate(OnLinkStatusChanged);
             _nodeEndpoint.Listen(this);
 
-            WaitHandle[] waitHandles = [_shutdownEvent, _packetReceivedEvent, _taskCompleteEvent, _taskCancelledEvent];
+            lock (s_activeInstanceLock)
+            {
+                s_activeInstance = this;
+            }
 
-            while (true)
+            try
             {
-                int index = WaitHandle.WaitAny(waitHandles);
-                switch (index)
+                WaitHandle[] waitHandles = [_shutdownEvent, _packetReceivedEvent, _taskCompleteEvent, _taskCancelledEvent];
+
+                while (true)
                 {
-                    case 0: // shutdownEvent
-                        NodeEngineShutdownReason shutdownReason = HandleShutdown();
-                        return shutdownReason;
+                    int index = WaitHandle.WaitAny(waitHandles);
+                    switch (index)
+                    {
+                        case 0: // shutdownEvent
+                            NodeEngineShutdownReason shutdownReason = HandleShutdown();
+                            return shutdownReason;
 
-                    case 1: // packetReceivedEvent
-                        INodePacket packet = null;
+                        case 1: // packetReceivedEvent
+                            INodePacket packet = null;
 
-                        int packetCount = _receivedPackets.Count;
+                            int packetCount = _receivedPackets.Count;
 
-                        while (packetCount > 0)
-                        {
-                            lock (_receivedPackets)
+                            while (packetCount > 0)
                             {
-                                if (_receivedPackets.Count > 0)
+                                lock (_receivedPackets)
                                 {
-                                    packet = _receivedPackets.Dequeue();
+                                    if (_receivedPackets.Count > 0)
+                                    {
+                                        packet = _receivedPackets.Dequeue();
+                                    }
+                                    else
+                                    {
+                                        break;
+                                    }
                                 }
-                                else
+
+                                if (packet != null)
                                 {
-                                    break;
+                                    HandlePacket(packet);
                                 }
                             }
 
-                            if (packet != null)
-                            {
-                                HandlePacket(packet);
-                            }
-                        }
-
-                        break;
-                    case 2: // taskCompleteEvent
-                        CompleteTask();
-                        break;
-                    case 3: // taskCancelledEvent
-                        CancelTask();
-                        break;
+                            break;
+                        case 2: // taskCompleteEvent
+                            CompleteTask();
+                            break;
+                        case 3: // taskCancelledEvent
+                            CancelTask();
+                            break;
+                    }
+                }
+            }
+            finally
+            {
+                lock (s_activeInstanceLock)
+                {
+                    if (s_activeInstance == this)
+                    {
+                        s_activeInstance = null;
+                    }
                 }
             }
 
             // UNREACHABLE
         }
+
+        /// <summary>
+        /// Requests the active out-of-proc task host node (if any) to shut down in response to an external stop signal
+        /// (for example <c>CTRL_C_EVENT</c> from Windows Restart Manager).
+        /// </summary>
+        internal static void RequestExternalShutdown()
+        {
+            OutOfProcTaskHostNode instance;
+            lock (s_activeInstanceLock)
+            {
+                instance = s_activeInstance;
+            }
+
+            instance?.SignalShutdownFromExternalStopRequest();
+        }
+
+        private void SignalShutdownFromExternalStopRequest()
+        {
+            _shutdownReason = NodeEngineShutdownReason.BuildComplete;
+            _shutdownEvent.Set();
+        }
+
         #endregion
 
         /// <summary>

src/MSBuild/XMake.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
@@ -13,6 +13,7 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.InteropServices;
 using System.Security;
 using System.Text;
 using System.Text.Json;
@@ -64,7 +65,7 @@ namespace Microsoft.Build.CommandLine
     /// This class implements the MSBuild.exe command-line application. It processes
     /// command-line arguments and invokes the build engine.
     /// </summary>
-    public static class MSBuildApp
+    public static partial class MSBuildApp
     {
         /// <summary>
         /// Enumeration of the various ways in which the MSBuild.exe application can exit.
@@ -317,6 +318,7 @@ public static int Main(string[] args)
             {
                 Console.CancelKeyPress += Console_CancelKeyPress;
 
+                EnsureWindowsConsoleShutdownHandlersRegistered();
 
                 // Use the client app to execute build in msbuild server. Opt-in feature.
                 exitCode = ((s_initialized && MSBuildClientApp.Execute(args, s_buildCancellationSource.Token) == ExitType.Success) ? 0 : 1);
@@ -677,6 +679,8 @@ public static ExitType Execute(string[] commandLine)
 
                 Console.CancelKeyPress += cancelHandler;
 
+                EnsureWindowsConsoleShutdownHandlersRegistered();
+
                 // check the operating system the code is running on
                 VerifyThrowSupportedOS();
 
@@ -1185,14 +1189,23 @@ private static void Console_CancelKeyPress(object sender, ConsoleCancelEventArgs
 
             e.Cancel = true; // do not terminate rudely
 
+            QueueGracefulShutdownFromExternalConsoleSignal();
+        }
+
+        /// <summary>
+        /// Begins the same graceful cancellation path as Ctrl+C, for signals that do not go through
+        /// <see cref="Console.CancelKeyPress"/> (e.g. Windows Restart Manager / installer shutdown via
+        /// <c>CTRL_CLOSE_EVENT</c>, <c>CTRL_LOGOFF_EVENT</c>, <c>CTRL_SHUTDOWN_EVENT</c>).
+        /// </summary>
+        private static void QueueGracefulShutdownFromExternalConsoleSignal()
+        {
             if (s_buildCancellationSource.IsCancellationRequested)
             {
                 return;
             }
 
             s_buildCancellationSource.Cancel();
 
-
             // The OS takes a lock in
             // kernel32.dll!_SetConsoleCtrlHandler, so if a task
             // waits for that lock somehow before quitting, it would hang
@@ -1243,6 +1256,43 @@ private static void Console_CancelKeyPress(object sender, ConsoleCancelEventArgs
             ThreadPoolExtensions.QueueThreadPoolWorkItemWithCulture(callback, CultureInfo.CurrentCulture, CultureInfo.CurrentUICulture);
         }
 
+        /// <summary>
+        /// Registers a console ctrl handler on Windows so that Restart Manager / installer shutdown
+        /// (<c>CTRL_CLOSE_EVENT</c>, <c>CTRL_LOGOFF_EVENT</c>, <c>CTRL_SHUTDOWN_EVENT</c>) triggers the
+        /// same graceful cancellation path as Ctrl+C. No-op on non-Windows.
+        /// </summary>
+        private static void EnsureWindowsConsoleShutdownHandlersRegistered()
+        {
+            if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            {
+                return;
+            }
+
+            s_consoleCtrlHandler = ConsoleCtrlHandler;
+            SetConsoleCtrlHandler(s_consoleCtrlHandler, Add: true);
+        }
+
+#if NETFRAMEWORK || NETCOREAPP
+        private static ConsoleCtrlHandlerDelegate s_consoleCtrlHandler;
+
+        private delegate bool ConsoleCtrlHandlerDelegate(uint dwCtrlType);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        private static extern bool SetConsoleCtrlHandler(ConsoleCtrlHandlerDelegate handler, bool Add);
+
+        private static bool ConsoleCtrlHandler(uint dwCtrlType)
+        {
+            // CTRL_CLOSE_EVENT (2), CTRL_LOGOFF_EVENT (5), CTRL_SHUTDOWN_EVENT (6) - Restart Manager / installer
+            if (dwCtrlType is 2 or 5 or 6)
+            {
+                QueueGracefulShutdownFromExternalConsoleSignal();
+                return true; // We handled it
+            }
+
+            return false;
+        }
+#endif
+
         /// <summary>
         /// Clears out any state accumulated from previous builds, and resets
         /// member data in preparation for a new build.