Detect and correct worker node over-provisioning

src/Build.UnitTests/BackEnd/NodeProviderOutOfProc_Tests.cs

@@ -0,0 +1,169 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using Microsoft.Build.BackEnd;
+using Microsoft.Build.Shared;
+using Shouldly;
+using Xunit;
+
+#nullable disable
+
+namespace Microsoft.Build.UnitTests.BackEnd
+{
+    /// <summary>
+    /// Tests for NodeProviderOutOfProc, specifically the node over-provisioning detection feature.
+    /// </summary>
+    public class NodeProviderOutOfProc_Tests
+    {
+        /// <summary>
+        /// Test helper class to expose protected methods for testing.
+        /// Uses configurable overrides for testing.
+        /// </summary>
+        private sealed class TestableNodeProviderOutOfProcBase : NodeProviderOutOfProcBase
+        {
+            private readonly int _systemWideNodeCount;
+            private readonly int? _thresholdOverride;
+
+            public TestableNodeProviderOutOfProcBase(int systemWideNodeCount, int? thresholdOverride = null)
+            {
+                _systemWideNodeCount = systemWideNodeCount;
+                _thresholdOverride = thresholdOverride;
+            }
+
+            protected override int GetNodeReuseThreshold()
+            {
+                // If threshold is overridden, use it; otherwise use base implementation
+                return _thresholdOverride ?? base.GetNodeReuseThreshold();
+            }
+
+            protected override int CountSystemWideActiveNodes()
+            {
+                return _systemWideNodeCount;
+            }
+
+            public bool[] TestDetermineNodesForReuse(int nodeCount, bool enableReuse)
+            {
+                return DetermineNodesForReuse(nodeCount, enableReuse);
+            }
+
+            public int TestGetNodeReuseThreshold()
+            {
+                return GetNodeReuseThreshold();
+            }
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenReuseDisabled_AllNodesShouldTerminate()
+        {
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 10, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 3, enableReuse: false);
+
+            result.Length.ShouldBe(3);
+            result.ShouldAllBe(shouldReuse => shouldReuse == false);
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenThresholdIsZero_AllNodesShouldTerminate()
+        {
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 10, thresholdOverride: 0);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 3, enableReuse: true);
+
+            result.Length.ShouldBe(3);
+            result.ShouldAllBe(shouldReuse => shouldReuse == false);
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenUnderThreshold_AllNodesShouldBeReused()
+        {
+            // System has 3 nodes total, threshold is 4, so we're under the limit
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 3, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 3, enableReuse: true);
+
+            result.Length.ShouldBe(3);
+            result.ShouldAllBe(shouldReuse => shouldReuse == true);
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenAtThreshold_AllNodesShouldBeReused()
+        {
+            // System has 4 nodes total, threshold is 4, so we're at the limit
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 4, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 4, enableReuse: true);
+
+            result.Length.ShouldBe(4);
+            result.ShouldAllBe(shouldReuse => shouldReuse == true);
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenOverThreshold_ExcessNodesShouldTerminate()
+        {
+            // System has 10 nodes total, threshold is 4
+            // This instance has 3 nodes
+            // We should keep 0 nodes from this instance (since 10 - 3 = 7, which is already > threshold)
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 10, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 3, enableReuse: true);
+
+            result.Length.ShouldBe(3);
+            result.ShouldAllBe(shouldReuse => shouldReuse == false);
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WhenSlightlyOverThreshold_SomeNodesShouldBeReused()
+        {
+            // System has 6 nodes total, threshold is 4
+            // This instance has 3 nodes
+            // Other instances have 6 - 3 = 3 nodes
+            // We need to reduce by 2 nodes to reach threshold
+            // So we should keep 1 node from this instance
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 6, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 3, enableReuse: true);
+
+            result.Length.ShouldBe(3);
+            // First node should be reused, others should terminate
+            result[0].ShouldBeTrue();
+            result[1].ShouldBeFalse();
+            result[2].ShouldBeFalse();
+        }
+
+        [Fact]
+        public void DetermineNodesForReuse_WithSingleNode_BehavesCorrectly()
+        {
+            // System has 5 nodes total, threshold is 4
+            // This instance has 1 node
+            // We're over threshold, but only by 1
+            // We should terminate this node since others already meet threshold
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 5, thresholdOverride: 4);
+
+            bool[] result = provider.TestDetermineNodesForReuse(nodeCount: 1, enableReuse: true);
+
+            result.Length.ShouldBe(1);
+            result[0].ShouldBeFalse();
+        }
+
+        [Fact]
+        public void GetNodeReuseThreshold_DefaultImplementation_ReturnsHalfOfCoreCount()
+        {
+            // Test the default implementation by not providing a threshold override
+            // Note: This test uses the actual system core count, so results vary by machine,
+            // but the mathematical relationship (threshold = max(1, cores/2)) should hold on all systems
+            int coreCount = NativeMethodsShared.GetLogicalCoreCount();
+            int expectedThreshold = Math.Max(1, coreCount / 2);
+
+            // Create a provider WITHOUT threshold override to test the base class implementation
+            var provider = new TestableNodeProviderOutOfProcBase(systemWideNodeCount: 0, thresholdOverride: null);
+
+            // The threshold from the provider should match our expected calculation
+            int actualThreshold = provider.TestGetNodeReuseThreshold();
+            actualThreshold.ShouldBe(expectedThreshold);
+            actualThreshold.ShouldBeGreaterThanOrEqualTo(1);
+            actualThreshold.ShouldBeLessThanOrEqualTo(coreCount);
+        }
+    }
+}

src/Build/BackEnd/Components/Communications/NodeProviderOutOfProcBase.cs

@@ -116,23 +116,39 @@ protected void ShutdownConnectedNodes(List<NodeContext> contextsToShutDown, bool
                                 !Console.IsInputRedirected &&
                                 Traits.Instance.EscapeHatches.EnsureStdOutForChildNodesIsPrimaryStdout;
 
+            // Determine which nodes should actually be reused based on system-wide node count
+            bool[] shouldReuseNode = DetermineNodesForReuse(contextsToShutDown.Count, enableReuse);
+
             Task[] waitForExitTasks = waitForExit && contextsToShutDown.Count > 0 ? new Task[contextsToShutDown.Count] : null;
             int i = 0;
+            int contextIndex = 0;
             var loggingService = _componentHost.LoggingService;
             foreach (NodeContext nodeContext in contextsToShutDown)
             {
                 if (nodeContext is null)
                 {
+                    contextIndex++;
                     continue;
                 }
-                nodeContext.SendData(new NodeBuildComplete(enableReuse));
-                if (waitForExit)
+
+                // Use the per-node reuse decision
+                bool reuseThisNode = shouldReuseNode[contextIndex++];
+                nodeContext.SendData(new NodeBuildComplete(reuseThisNode));
+
+                if (!reuseThisNode || waitForExit)
                 {
-                    waitForExitTasks[i++] = nodeContext.WaitForExitAsync(loggingService);
+                    if (i < (waitForExitTasks?.Length ?? 0))
+                    {
+                        waitForExitTasks[i++] = nodeContext.WaitForExitAsync(loggingService);
+                    }
                 }
             }
-            if (waitForExitTasks != null)
+            if (waitForExitTasks != null && i > 0)
             {
+                if (i < waitForExitTasks.Length)
+                {
+                    Array.Resize(ref waitForExitTasks, i);
+                }
                 Task.WaitAll(waitForExitTasks);
             }
         }
@@ -511,6 +527,170 @@ private string GetProcessesToIgnoreKey(Handshake hostHandshake, int nodeProcessI
 #endif
         }
 
+        /// <summary>
+        /// Determines which nodes should be reused based on system-wide node count to avoid over-provisioning.
+        /// </summary>
+        /// <param name="nodeCount">The number of nodes in this MSBuild instance</param>
+        /// <param name="enableReuse">Whether reuse is enabled at all</param>
+        /// <returns>Array indicating which nodes should be reused (true) or terminated (false)</returns>
+        protected virtual bool[] DetermineNodesForReuse(int nodeCount, bool enableReuse)
+        {
+            bool[] shouldReuse = new bool[nodeCount];
+
+            // If reuse is disabled, no nodes should be reused
+            if (!enableReuse)
+            {
+                return shouldReuse; // All false
+            }
+
+            // Get threshold for this node type
+            int maxNodesToKeep = GetNodeReuseThreshold();
+
+            // If threshold is 0, terminate all nodes in this instance
+            if (maxNodesToKeep == 0)
+            {
+                CommunicationsUtilities.Trace("Node reuse threshold is 0, terminating all {0} nodes", nodeCount);
+                return shouldReuse; // All false
+            }
+
+            // Count system-wide active nodes of the same type
+            int systemWideNodeCount = CountSystemWideActiveNodes();
+
+            CommunicationsUtilities.Trace("System-wide node count: {0}, threshold: {1}, this instance has: {2} nodes",
+                systemWideNodeCount, maxNodesToKeep, nodeCount);
+
+            // If we're already under the threshold system-wide, keep all our nodes
+            if (systemWideNodeCount <= maxNodesToKeep)
+            {
+                for (int i = 0; i < nodeCount; i++)
+                {
+                    shouldReuse[i] = true;
+                }
+                return shouldReuse;
+            }
+
+            // We're over-provisioned. Determine how many of our nodes to keep.
+            // Strategy: Keep nodes up to the threshold, terminate the rest.
+            // This instance's contribution is limited to help reach the threshold.
+            int nodesToKeepInThisInstance = Math.Max(0, maxNodesToKeep - (systemWideNodeCount - nodeCount));
+
+            CommunicationsUtilities.Trace("Keeping {0} of {1} nodes in this instance to help meet threshold of {2}",
+                nodesToKeepInThisInstance, nodeCount, maxNodesToKeep);
+
+            // Mark the first N nodes for reuse
+            for (int i = 0; i < Math.Min(nodesToKeepInThisInstance, nodeCount); i++)
+            {
+                shouldReuse[i] = true;
+            }
+
+            return shouldReuse;
+        }
+
+        /// <summary>
+        /// Gets the maximum number of nodes of this type that should remain active system-wide.
+        /// </summary>
+        /// <returns>The threshold for node reuse</returns>
+        protected virtual int GetNodeReuseThreshold()
+        {
+            // Default for worker nodes: NUM_PROCS / 2
+            // Derived classes (Server, RAR) can override to return 0
+            return Math.Max(1, NativeMethodsShared.GetLogicalCoreCount() / 2);
+        }
+
+        /// <summary>
+        /// Counts the number of active MSBuild node processes of the same type system-wide.
+        /// Uses improved node detection logic to filter by NodeMode and handle dotnet processes.
+        /// </summary>
+        /// <returns>The count of active node processes</returns>
+        protected virtual int CountSystemWideActiveNodes()
+            => CountActiveNodesWithMode(NodeMode.OutOfProcNode);
+
+        /// <summary>
+        /// Counts the number of active MSBuild processes running with the specified <see cref="NodeMode"/>.
+        /// Includes the current process in the count if it matches.
+        /// Used by out-of-proc nodes (e.g., server node) to detect over-provisioning at build completion.
+        /// </summary>
+        /// <param name="nodeMode">The node mode to filter for.</param>
+        /// <returns>The number of matching processes, or 0 if enumeration fails or the feature wave is disabled.</returns>
+        internal static int CountActiveNodesWithMode(NodeMode nodeMode)
+        {
+            try
+            {
+                (_, IList<Process> nodeProcesses) = GetPossibleRunningNodes(nodeMode);
+                int count = nodeProcesses.Count;
+                foreach (var process in nodeProcesses)
+                {
+                    process?.Dispose();
+                }
+                return count;
+            }
+            catch (Exception ex)
+            {
+                CommunicationsUtilities.Trace("Error counting system-wide nodes with mode {0}: {1}", nodeMode, ex.Message);
+                return 0;
+            }
+        }
+
+        private static (string expectedProcessName, IList<Process> nodeProcesses) GetPossibleRunningNodes(NodeMode? expectedNodeMode)
+        {
+            string msbuildLocation = Constants.MSBuildExecutableName;
+            var expectedProcessName = Path.GetFileNameWithoutExtension(CurrentHost.GetCurrentHost() ?? msbuildLocation);
+
+            Process[] processes;
+            try
+            {
+                processes = Process.GetProcessesByName(expectedProcessName);
+            }
+            catch
+            {
+                return (expectedProcessName, Array.Empty<Process>());
+            }
+
+            if (expectedNodeMode.HasValue && ChangeWaves.AreFeaturesEnabled(ChangeWaves.Wave18_5))
+            {
+                List<Process> filteredProcesses = [];
+                bool isDotnetProcess = expectedProcessName.Equals(Path.GetFileNameWithoutExtension(Constants.DotnetProcessName), StringComparison.OrdinalIgnoreCase);
+
+                foreach (var process in processes)
+                {
+                    try
+                    {
+                        if (!process.TryGetCommandLine(out string commandLine))
+                        {
+                            continue;
+                        }
+
+                        if (commandLine is null)
+                        {
+                            filteredProcesses.Add(process);
+                            continue;
+                        }
+
+                        if (isDotnetProcess && !commandLine.Contains("MSBuild.dll", StringComparison.OrdinalIgnoreCase))
+                        {
+                            continue;
+                        }
+
+                        NodeMode? processNodeMode = NodeModeHelper.ExtractFromCommandLine(commandLine);
+                        if (processNodeMode.HasValue && processNodeMode.Value == expectedNodeMode.Value)
+                        {
+                            filteredProcesses.Add(process);
+                        }
+                    }
+                    catch
+                    {
+                        continue;
+                    }
+                }
+
+                filteredProcesses.Sort((left, right) => left.Id.CompareTo(right.Id));
+                return (expectedProcessName, filteredProcesses);
+            }
+
+            Array.Sort(processes, (left, right) => left.Id.CompareTo(right.Id));
+            return (expectedProcessName, processes);
+        }
+
 #if !FEATURE_PIPEOPTIONS_CURRENTUSERONLY
         // This code needs to be in a separate method so that we don't try (and fail) to load the Windows-only APIs when JIT-ing the code
         //  on non-Windows operating systems