Allow NodeProviderOutOfProcTaskHost to manage multiple nodes instead of one per arch

src/Build.UnitTests/BackEnd/AssemblyTaskFactory_Tests.cs

@@ -251,6 +251,7 @@ public void VerifyGoodTaskInstantiation()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 createdTask.ShouldNotBeNull();
                 createdTask.ShouldNotBeOfType<TaskHostTask>();
@@ -283,6 +284,7 @@ public void VerifyMatchingTaskParametersDontLaunchTaskHost1()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.False(createdTask is TaskHostTask);
@@ -315,6 +317,7 @@ public void VerifyMatchingTaskParametersDontLaunchTaskHost2()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.False(createdTask is TaskHostTask);
@@ -349,6 +352,7 @@ public void VerifyMatchingUsingTaskParametersDontLaunchTaskHost1()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.False(createdTask is TaskHostTask);
@@ -383,6 +387,7 @@ public void VerifyMatchingUsingTaskParametersDontLaunchTaskHost2()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.False(createdTask is TaskHostTask);
@@ -419,6 +424,7 @@ public void VerifyMatchingParametersDontLaunchTaskHost()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.False(createdTask is TaskHostTask);
@@ -453,6 +459,7 @@ public void VerifyNonmatchingUsingTaskParametersLaunchTaskHost()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -485,6 +492,7 @@ public void VerifyNonmatchingTaskParametersLaunchTaskHost()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -521,6 +529,7 @@ public void VerifyNonmatchingParametersLaunchTaskHost()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -551,6 +560,7 @@ public void VerifyExplicitlyLaunchTaskHost()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -585,6 +595,7 @@ public void VerifyExplicitlyLaunchTaskHostEvenIfParametersMatch1()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -619,6 +630,7 @@ public void VerifyExplicitlyLaunchTaskHostEvenIfParametersMatch2()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);
@@ -654,6 +666,7 @@ public void VerifySameFactoryCanGenerateDifferentTaskInstances()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsNotType<TaskHostTask>(createdTask);
@@ -678,6 +691,7 @@ public void VerifySameFactoryCanGenerateDifferentTaskInstances()
                     new AppDomainSetup(),
 #endif
                     false,
+                    scheduledNodeId: 1,
                     (string propName) => ProjectPropertyInstance.Create("test", "test"));
                 Assert.NotNull(createdTask);
                 Assert.IsType<TaskHostTask>(createdTask);

src/Build.UnitTests/BackEnd/TaskExecutionHost_Tests.cs

@@ -998,7 +998,7 @@ public void TestTaskResolutionFailureWithUsingTask()
                     false,
                     CancellationToken.None);
                 _host.FindTask(null);
-                _host.InitializeForBatch(new TaskLoggingContext(_loggingService, tlc.BuildEventContext), _bucket, null);
+                _host.InitializeForBatch(new TaskLoggingContext(_loggingService, tlc.BuildEventContext), _bucket, null, scheduledNodeId: 1);
             });
         }
         /// <summary>
@@ -1027,7 +1027,7 @@ public void TestTaskResolutionFailureWithNoUsingTask()
                 CancellationToken.None);
 
             _host.FindTask(null);
-            _host.InitializeForBatch(new TaskLoggingContext(_loggingService, tlc.BuildEventContext), _bucket, null);
+            _host.InitializeForBatch(new TaskLoggingContext(_loggingService, tlc.BuildEventContext), _bucket, null, scheduledNodeId: 1);
             _logger.AssertLogContains("MSB4036");
         }
 
@@ -1286,7 +1286,7 @@ private void InitializeHost(bool throwOnExecute)
             _bucket = new ItemBucket(FrozenSet<string>.Empty, new Dictionary<string, string>(), new Lookup(itemsByName, new PropertyDictionary<ProjectPropertyInstance>()), 0);
             _bucket.Initialize(null);
             _host.FindTask(null);
-            _host.InitializeForBatch(talc, _bucket, null);
+            _host.InitializeForBatch(talc, _bucket, null, scheduledNodeId: 1);
             _parametersSetOnTask = new Dictionary<string, object>(StringComparer.OrdinalIgnoreCase);
             _outputsReadFromTask = new Dictionary<string, object>(StringComparer.OrdinalIgnoreCase);
         }

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

@@ -328,7 +328,8 @@ private IList<NodeInfo> AttemptCreateNode(INodeProvider nodeProvider, NodeConfig
 
             // Assign a global ID to the node we are about to create.
             int fromNodeId;
-            if (nodeProvider is NodeProviderInProc && !(_componentHost?.BuildParameters.MultiThreaded ?? false))
+            bool isMultiThreadedModeOn = _componentHost?.BuildParameters.MultiThreaded ?? false;
+            if (nodeProvider is NodeProviderInProc && !isMultiThreadedModeOn)
             {
                 fromNodeId = _inprocNodeId;
             }

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

@@ -357,7 +357,7 @@ private bool InstantiateNode(int nodeId, INodePacketFactory factory)
             nodeContext._inProcNodeEndpoint.OnLinkStatusChanged += new LinkStatusChangedDelegate(InProcNodeEndpoint_OnLinkStatusChanged);
 
             nodeContext._packetFactory = factory;
-            nodeContext._inProcNode = new InProcNode(_componentHost, endpoints.NodeEndpoint);
+            nodeContext._inProcNode = new InProcNode(nodeId, _componentHost, endpoints.NodeEndpoint);
 #if FEATURE_THREAD_CULTURE
             nodeContext._inProcNodeThread = new Thread(() => InProcNodeThreadProc(nodeContext._inProcNode), BuildParameters.ThreadStackSize);
 #else

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
@@ -22,13 +23,6 @@ namespace Microsoft.Build.BackEnd
     /// </summary>
     internal class NodeProviderOutOfProcTaskHost : NodeProviderOutOfProcBase, INodeProvider, INodePacketFactory, INodePacketHandler
     {
-        /// <summary>
-        /// The maximum number of nodes that this provider supports. Should
-        /// always be equivalent to the number of different TaskHostContexts
-        /// that exist.
-        /// </summary>
-        private const int MaxNodeCount = 4;
-
         /// <summary>
         /// Store the path for MSBuild / MSBuildTaskHost so that we don't have to keep recalculating it.
         /// </summary>
@@ -85,9 +79,9 @@ internal class NodeProviderOutOfProcTaskHost : NodeProviderOutOfProcBase, INodeP
         private ManualResetEvent _noNodesActiveEvent;
 
         /// <summary>
-        /// A mapping of all the nodes managed by this provider.
+        /// A mapping of all the task host nodes managed by this provider.
         /// </summary>
-        private Dictionary<HandshakeOptions, NodeContext> _nodeContexts;
+        private ConcurrentDictionary<int, NodeContext> _nodeContexts;
 
         /// <summary>
         /// A mapping of all of the INodePacketFactories wrapped by this provider.
@@ -135,7 +129,7 @@ public int AvailableNodes
         {
             get
             {
-                return MaxNodeCount - _nodeContexts.Count;
+                throw new NotImplementedException("This property is not implemented because available nodes are unlimited.");
             }
         }
 
@@ -175,19 +169,9 @@ public IList<NodeInfo> CreateNodes(int nextNodeId, INodePacketFactory packetFact
         /// <param name="packet">The packet to send.</param>
         public void SendData(int nodeId, INodePacket packet)
         {
-            throw new NotImplementedException("Use the other overload of SendData instead");
-        }
-
-        /// <summary>
-        /// Sends data to the specified node.
-        /// </summary>
-        /// <param name="hostContext">The node to which data shall be sent.</param>
-        /// <param name="packet">The packet to send.</param>
-        public void SendData(HandshakeOptions hostContext, INodePacket packet)
-        {
-            ErrorUtilities.VerifyThrow(_nodeContexts.ContainsKey(hostContext), "Invalid host context specified: {0}.", hostContext.ToString());
+            ErrorUtilities.VerifyThrow(_nodeContexts.TryGetValue(nodeId, out NodeContext context), "Invalid host context specified: {0}.", nodeId);
 
-            SendData(_nodeContexts[hostContext], packet);
+            SendData(context, packet);
         }
 
         /// <summary>
@@ -197,12 +181,7 @@ public void SendData(HandshakeOptions hostContext, INodePacket packet)
         public void ShutdownConnectedNodes(bool enableReuse)
         {
             // Send the build completion message to the nodes, causing them to shutdown or reset.
-            List<NodeContext> contextsToShutDown;
-
-            lock (_nodeContexts)
-            {
-                contextsToShutDown = new List<NodeContext>(_nodeContexts.Values);
-            }
+            List<NodeContext> contextsToShutDown = [.. _nodeContexts.Values];
 
             ShutdownConnectedNodes(contextsToShutDown, enableReuse);
 
@@ -227,7 +206,7 @@ public void ShutdownAllNodes()
         public void InitializeComponent(IBuildComponentHost host)
         {
             this.ComponentHost = host;
-            _nodeContexts = new Dictionary<HandshakeOptions, NodeContext>();
+            _nodeContexts = new ConcurrentDictionary<int, NodeContext>();
             _nodeIdToPacketFactory = new Dictionary<int, INodePacketFactory>();
             _nodeIdToPacketHandler = new Dictionary<int, INodePacketHandler>();
             _activeNodes = new HashSet<int>();
@@ -589,15 +568,16 @@ private static string GetPathFromEnvironmentOrDefault(string environmentVariable
         /// </summary>
         internal bool AcquireAndSetUpHost(
             HandshakeOptions hostContext,
+            int taskHostNodeId,
             INodePacketFactory factory,
             INodePacketHandler handler,
             TaskHostConfiguration configuration,
             Dictionary<string, string> taskHostParameters)
         {
             bool nodeCreationSucceeded;
-            if (!_nodeContexts.ContainsKey(hostContext))
+            if (!_nodeContexts.ContainsKey(taskHostNodeId))
             {
-                nodeCreationSucceeded = CreateNode(hostContext, factory, handler, configuration, taskHostParameters);
+                nodeCreationSucceeded = CreateNode(hostContext, taskHostNodeId, factory, handler, configuration, taskHostParameters);
             }
             else
             {
@@ -607,9 +587,9 @@ internal bool AcquireAndSetUpHost(
 
             if (nodeCreationSucceeded)
             {
-                NodeContext context = _nodeContexts[hostContext];
-                _nodeIdToPacketFactory[(int)hostContext] = factory;
-                _nodeIdToPacketHandler[(int)hostContext] = handler;
+                NodeContext context = _nodeContexts[taskHostNodeId];
+                _nodeIdToPacketFactory[taskHostNodeId] = factory;
+                _nodeIdToPacketHandler[taskHostNodeId] = handler;
 
                 // Configure the node.
                 context.SendData(configuration);
@@ -622,33 +602,26 @@ internal bool AcquireAndSetUpHost(
         /// <summary>
         /// Expected to be called when TaskHostTask is done with host of the given context.
         /// </summary>
-        internal void DisconnectFromHost(HandshakeOptions hostContext)
+        internal void DisconnectFromHost(int nodeId)
         {
-            ErrorUtilities.VerifyThrow(_nodeIdToPacketFactory.ContainsKey((int)hostContext) && _nodeIdToPacketHandler.ContainsKey((int)hostContext), "Why are we trying to disconnect from a context that we already disconnected from?  Did we call DisconnectFromHost twice?");
+            ErrorUtilities.VerifyThrow(_nodeIdToPacketFactory.ContainsKey(nodeId) && _nodeIdToPacketHandler.ContainsKey(nodeId), "Why are we trying to disconnect from a context that we already disconnected from?  Did we call DisconnectFromHost twice?");
 
-            _nodeIdToPacketFactory.Remove((int)hostContext);
-            _nodeIdToPacketHandler.Remove((int)hostContext);
+            _nodeIdToPacketFactory.Remove(nodeId);
+            _nodeIdToPacketHandler.Remove(nodeId);
         }
 
         /// <summary>
         /// Instantiates a new MSBuild or MSBuildTaskHost process acting as a child node.
         /// </summary>
-        internal bool CreateNode(HandshakeOptions hostContext, INodePacketFactory factory, INodePacketHandler handler, TaskHostConfiguration configuration, Dictionary<string, string> taskHostParameters)
+        internal bool CreateNode(HandshakeOptions hostContext, int taskHostNodeId, INodePacketFactory factory, INodePacketHandler handler, TaskHostConfiguration configuration, Dictionary<string, string> taskHostParameters)
         {
             ErrorUtilities.VerifyThrowArgumentNull(factory);
-            ErrorUtilities.VerifyThrow(!_nodeIdToPacketFactory.ContainsKey((int)hostContext), "We should not already have a factory for this context!  Did we forget to call DisconnectFromHost somewhere?");
-
-            if (AvailableNodes <= 0)
-            {
-                ErrorUtilities.ThrowInternalError("All allowable nodes already created ({0}).", _nodeContexts.Count);
-                return false;
-            }
+            ErrorUtilities.VerifyThrow(!_nodeIdToPacketFactory.ContainsKey(taskHostNodeId), "We should not already have a factory for this context!  Did we forget to call DisconnectFromHost somewhere?");
 
             // if runtime host path is null it means we don't have MSBuild.dll path resolved and there is no need to include it in the command line arguments.
             string commandLineArgsPlaceholder = "{0} /nologo /nodemode:2 /nodereuse:{1} /low:{2} ";
 
             IList<NodeContext> nodeContexts;
-            int nodeId = (int)hostContext;
 
             // Handle .NET task host context
 #if NETFRAMEWORK
@@ -664,7 +637,7 @@ internal bool CreateNode(HandshakeOptions hostContext, INodePacketFactory factor
                 nodeContexts = GetNodes(
                     runtimeHostPath,
                     string.Format(commandLineArgsPlaceholder, Path.Combine(msbuildAssemblyPath, Constants.MSBuildAssemblyName), NodeReuseIsEnabled(hostContext), ComponentHost.BuildParameters.LowPriority),
-                    nodeId,
+                    taskHostNodeId,
                     this,
                     handshake,
                     NodeContextCreated,
@@ -688,7 +661,7 @@ internal bool CreateNode(HandshakeOptions hostContext, INodePacketFactory factor
             nodeContexts = GetNodes(
                 msbuildLocation,
                 string.Format(commandLineArgsPlaceholder, string.Empty, NodeReuseIsEnabled(hostContext), ComponentHost.BuildParameters.LowPriority),
-                nodeId,
+                taskHostNodeId,
                 this,
                 new Handshake(hostContext),
                 NodeContextCreated,
@@ -714,7 +687,7 @@ bool NodeReuseIsEnabled(HandshakeOptions hostContext)
         /// </summary>
         private void NodeContextCreated(NodeContext context)
         {
-            _nodeContexts[(HandshakeOptions)context.NodeId] = context;
+            _nodeContexts[context.NodeId] = context;
 
             // Start the asynchronous read.
             context.BeginAsyncPacketRead();
@@ -731,10 +704,7 @@ private void NodeContextCreated(NodeContext context)
         /// </summary>
         private void NodeContextTerminated(int nodeId)
         {
-            lock (_nodeContexts)
-            {
-                _nodeContexts.Remove((HandshakeOptions)nodeId);
-            }
+            _nodeContexts.TryRemove(nodeId, out _);
 
             // May also be removed by unnatural termination, so don't assume it's there
             lock (_activeNodes)

src/Build/BackEnd/Components/RequestBuilder/TaskBuilder.cs

@@ -658,7 +658,7 @@ private void ExecuteIntrinsicTask(ItemBucket bucket)
         /// </summary>
         private async Task<WorkUnitResult> InitializeAndExecuteTask(TaskLoggingContext taskLoggingContext, ItemBucket bucket, IDictionary<string, string> taskIdentityParameters, TaskHost taskHost, TaskExecutionMode howToExecuteTask)
         {
-            if (!_taskExecutionHost.InitializeForBatch(taskLoggingContext, bucket, taskIdentityParameters))
+            if (!_taskExecutionHost.InitializeForBatch(taskLoggingContext, bucket, taskIdentityParameters, _buildRequestEntry.Request.ScheduledNodeId))
             {
                 ProjectErrorUtilities.ThrowInvalidProject(_targetChildInstance.Location, "TaskDeclarationOrUsageError", _taskNode.Name);
             }