diff --git a/src/Build.UnitTests/BackEnd/TaskHostNodeKey_Tests.cs b/src/Build.UnitTests/BackEnd/TaskHostNodeKey_Tests.cs new file mode 100644 index 00000000000..c02ea2c094e --- /dev/null +++ b/src/Build.UnitTests/BackEnd/TaskHostNodeKey_Tests.cs @@ -0,0 +1,116 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using Microsoft.Build.Internal; +using Shouldly; +using Xunit; + +namespace Microsoft.Build.Engine.UnitTests.BackEnd +{ + /// + /// Tests for TaskHostNodeKey record struct functionality. + /// + public class TaskHostNodeKey_Tests + { + [Fact] + public void TaskHostNodeKey_Equality_SameValues_AreEqual() + { + var key1 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + var key2 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + + key1.ShouldBe(key2); + (key1 == key2).ShouldBeTrue(); + key1.GetHashCode().ShouldBe(key2.GetHashCode()); + } + + [Fact] + public void TaskHostNodeKey_Equality_DifferentNodeId_AreNotEqual() + { + var key1 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + var key2 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 2); + + key1.ShouldNotBe(key2); + (key1 != key2).ShouldBeTrue(); + } + + [Fact] + public void TaskHostNodeKey_Equality_DifferentHandshakeOptions_AreNotEqual() + { + var key1 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + var key2 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.X64, 1); + + key1.ShouldNotBe(key2); + (key1 != key2).ShouldBeTrue(); + } + + [Fact] + public void TaskHostNodeKey_CanBeUsedAsDictionaryKey() + { + var dict = new System.Collections.Generic.Dictionary(); + var key1 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + var key2 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.X64, 2); + + dict[key1] = "value1"; + dict[key2] = "value2"; + + dict[key1].ShouldBe("value1"); + dict[key2].ShouldBe("value2"); + + // Create a new key with same values as key1 + var key1Copy = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1); + dict[key1Copy].ShouldBe("value1"); + } + + [Fact] + public void TaskHostNodeKey_LargeNodeId_Works() + { + // Test that we can use node IDs greater than 255 (the previous limit) + var key1 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 256); + var key2 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, 1000); + var key3 = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, int.MaxValue); + + key1.NodeId.ShouldBe(256); + key2.NodeId.ShouldBe(1000); + key3.NodeId.ShouldBe(int.MaxValue); + + // Ensure they are all different + key1.ShouldNotBe(key2); + key2.ShouldNotBe(key3); + key1.ShouldNotBe(key3); + } + + [Fact] + public void TaskHostNodeKey_NegativeNodeId_Works() + { + // Traditional multi-proc builds use -1 for node ID + var key = new TaskHostNodeKey(HandshakeOptions.TaskHost | HandshakeOptions.NET, -1); + + key.NodeId.ShouldBe(-1); + key.HandshakeOptions.ShouldBe(HandshakeOptions.TaskHost | HandshakeOptions.NET); + } + + [Fact] + public void TaskHostNodeKey_AllHandshakeOptions_Work() + { + // Test various HandshakeOptions combinations + HandshakeOptions[] optionsList = + [ + HandshakeOptions.None, + HandshakeOptions.TaskHost, + HandshakeOptions.TaskHost | HandshakeOptions.NET, + HandshakeOptions.TaskHost | HandshakeOptions.X64, + HandshakeOptions.TaskHost | HandshakeOptions.NET | HandshakeOptions.NodeReuse, + HandshakeOptions.TaskHost | HandshakeOptions.CLR2, + HandshakeOptions.TaskHost | HandshakeOptions.Arm64 + ]; + + foreach (var options in optionsList) + { + var key = new TaskHostNodeKey(options, 42); + + key.HandshakeOptions.ShouldBe(options); + key.NodeId.ShouldBe(42); + } + } + } +} diff --git a/src/Build/BackEnd/Components/Communications/NodeProviderOutOfProcTaskHost.cs b/src/Build/BackEnd/Components/Communications/NodeProviderOutOfProcTaskHost.cs index 4741ebbbb08..baff1943fa8 100644 --- a/src/Build/BackEnd/Components/Communications/NodeProviderOutOfProcTaskHost.cs +++ b/src/Build/BackEnd/Components/Communications/NodeProviderOutOfProcTaskHost.cs @@ -81,11 +81,19 @@ internal class NodeProviderOutOfProcTaskHost : NodeProviderOutOfProcBase, INodeP /// /// A mapping of all the task host nodes managed by this provider. + /// The key is a TaskHostNodeKey combining HandshakeOptions and scheduled node ID. /// - private ConcurrentDictionary _nodeContexts; + private ConcurrentDictionary _nodeContexts; + + /// + /// Reverse mapping from communication node ID to TaskHostNodeKey. + /// Used for O(1) lookup when handling node termination from ShutdownAllNodes. + /// + private ConcurrentDictionary _nodeIdToNodeKey; /// /// A mapping of all of the INodePacketFactories wrapped by this provider. + /// Keyed by the communication node ID (NodeContext.NodeId) for O(1) packet routing. /// Thread-safe to support parallel taskhost creation in /mt mode where multiple thread nodes /// can simultaneously create their own taskhosts. /// @@ -93,16 +101,23 @@ internal class NodeProviderOutOfProcTaskHost : NodeProviderOutOfProcBase, INodeP /// /// A mapping of all of the INodePacketHandlers wrapped by this provider. + /// Keyed by the communication node ID (NodeContext.NodeId) for O(1) packet routing. /// Thread-safe to support parallel taskhost creation in /mt mode where multiple thread nodes /// can simultaneously create their own taskhosts. /// private ConcurrentDictionary _nodeIdToPacketHandler; /// - /// Keeps track of the set of nodes for which we have not yet received shutdown notification. + /// Keeps track of the set of node IDs for which we have not yet received shutdown notification. /// private HashSet _activeNodes; + /// + /// Counter for generating unique communication node IDs. + /// Incremented atomically for each new node created. + /// + private int _nextNodeId; + /// /// Packet factory we use if there's not already one associated with a particular context. /// @@ -169,12 +184,23 @@ public IList CreateNodes(int nextNodeId, INodePacketFactory packetFact /// /// Sends data to the specified node. + /// Note: For task hosts, use the overload that takes TaskHostNodeKey instead. /// /// The node to which data shall be sent. /// The packet to send. public void SendData(int nodeId, INodePacket packet) { - ErrorUtilities.VerifyThrow(_nodeContexts.TryGetValue(nodeId, out NodeContext context), "Invalid host context specified: {0}.", nodeId); + throw new NotImplementedException("For task hosts, use the overload that takes TaskHostNodeKey."); + } + + /// + /// Sends data to the specified task host node. + /// + /// The task host node key identifying the target node. + /// The packet to send. + internal void SendData(TaskHostNodeKey nodeKey, INodePacket packet) + { + ErrorUtilities.VerifyThrow(_nodeContexts.TryGetValue(nodeKey, out NodeContext context), "Invalid host context specified: {0}.", nodeKey); SendData(context, packet); } @@ -211,10 +237,12 @@ public void ShutdownAllNodes() public void InitializeComponent(IBuildComponentHost host) { this.ComponentHost = host; - _nodeContexts = new ConcurrentDictionary(); + _nodeContexts = new ConcurrentDictionary(); + _nodeIdToNodeKey = new ConcurrentDictionary(); _nodeIdToPacketFactory = new ConcurrentDictionary(); _nodeIdToPacketHandler = new ConcurrentDictionary(); - _activeNodes = new HashSet(); + _activeNodes = []; + _nextNodeId = 0; _noNodesActiveEvent = new ManualResetEvent(true); _localPacketFactory = new NodePacketFactory(); @@ -569,17 +597,16 @@ private static string GetPathFromEnvironmentOrDefault(string environmentVariable /// Make sure a node in the requested context exists. /// internal bool AcquireAndSetUpHost( - HandshakeOptions hostContext, - int taskHostNodeId, + TaskHostNodeKey nodeKey, INodePacketFactory factory, INodePacketHandler handler, TaskHostConfiguration configuration, in TaskHostParameters taskHostParameters) { bool nodeCreationSucceeded; - if (!_nodeContexts.ContainsKey(taskHostNodeId)) + if (!_nodeContexts.ContainsKey(nodeKey)) { - nodeCreationSucceeded = CreateNode(hostContext, taskHostNodeId, factory, handler, configuration, taskHostParameters); + nodeCreationSucceeded = CreateNode(nodeKey, factory, handler, configuration, taskHostParameters); } else { @@ -589,9 +616,10 @@ internal bool AcquireAndSetUpHost( if (nodeCreationSucceeded) { - NodeContext context = _nodeContexts[taskHostNodeId]; - _nodeIdToPacketFactory[taskHostNodeId] = factory; - _nodeIdToPacketHandler[taskHostNodeId] = handler; + NodeContext context = _nodeContexts[nodeKey]; + // Map the transport ID directly to the handlers for O(1) packet routing + _nodeIdToPacketFactory[context.NodeId] = factory; + _nodeIdToPacketHandler[context.NodeId] = handler; // Configure the node. context.SendData(configuration); @@ -604,10 +632,12 @@ internal bool AcquireAndSetUpHost( /// /// Expected to be called when TaskHostTask is done with host of the given context. /// - internal void DisconnectFromHost(int nodeId) + internal void DisconnectFromHost(TaskHostNodeKey nodeKey) { - bool successRemoveFactory = _nodeIdToPacketFactory.TryRemove(nodeId, out _); - bool successRemoveHandler = _nodeIdToPacketHandler.TryRemove(nodeId, out _); + ErrorUtilities.VerifyThrow(_nodeContexts.TryGetValue(nodeKey, out NodeContext context), "Node context not found for key: {0}. Was the node created?", nodeKey); + + bool successRemoveFactory = _nodeIdToPacketFactory.TryRemove(context.NodeId, out _); + bool successRemoveHandler = _nodeIdToPacketHandler.TryRemove(context.NodeId, out _); ErrorUtilities.VerifyThrow(successRemoveFactory && successRemoveHandler, "Why are we trying to disconnect from a context that we already disconnected from? Did we call DisconnectFromHost twice?"); } @@ -615,14 +645,22 @@ internal void DisconnectFromHost(int nodeId) /// /// Instantiates a new MSBuild or MSBuildTaskHost process acting as a child node. /// - internal bool CreateNode(HandshakeOptions hostContext, int taskHostNodeId, INodePacketFactory factory, INodePacketHandler handler, TaskHostConfiguration configuration, in TaskHostParameters taskHostParameters) + internal bool CreateNode(TaskHostNodeKey nodeKey, INodePacketFactory factory, INodePacketHandler handler, TaskHostConfiguration configuration, in TaskHostParameters taskHostParameters) { ErrorUtilities.VerifyThrowArgumentNull(factory); - ErrorUtilities.VerifyThrow(!_nodeIdToPacketFactory.ContainsKey(taskHostNodeId), "We should not already have a factory for this context! Did we forget to call DisconnectFromHost somewhere?"); + ErrorUtilities.VerifyThrow(!_nodeContexts.ContainsKey(nodeKey), "We should not already have a node for this context! Did we forget to call DisconnectFromHost somewhere?"); + + HandshakeOptions hostContext = nodeKey.HandshakeOptions; // 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} "; + // Generate a unique node ID for communication purposes using atomic increment. + int communicationNodeId = Interlocked.Increment(ref _nextNodeId); + + // Create callbacks that capture the TaskHostNodeKey + void OnNodeContextCreated(NodeContext context) => NodeContextCreated(context, nodeKey); + IList nodeContexts; // Handle .NET task host context @@ -639,10 +677,10 @@ internal bool CreateNode(HandshakeOptions hostContext, int taskHostNodeId, INode nodeContexts = GetNodes( runtimeHostPath, string.Format(commandLineArgsPlaceholder, Path.Combine(msbuildAssemblyPath, Constants.MSBuildAssemblyName), NodeReuseIsEnabled(hostContext), ComponentHost.BuildParameters.LowPriority), - taskHostNodeId, + communicationNodeId, this, handshake, - NodeContextCreated, + OnNodeContextCreated, NodeContextTerminated, 1); @@ -663,10 +701,10 @@ internal bool CreateNode(HandshakeOptions hostContext, int taskHostNodeId, INode nodeContexts = GetNodes( msbuildLocation, string.Format(commandLineArgsPlaceholder, string.Empty, NodeReuseIsEnabled(hostContext), ComponentHost.BuildParameters.LowPriority), - taskHostNodeId, + communicationNodeId, this, new Handshake(hostContext), - NodeContextCreated, + OnNodeContextCreated, NodeContextTerminated, 1); @@ -687,9 +725,10 @@ bool NodeReuseIsEnabled(HandshakeOptions hostContext) /// /// Method called when a context created. /// - private void NodeContextCreated(NodeContext context) + private void NodeContextCreated(NodeContext context, TaskHostNodeKey nodeKey) { - _nodeContexts[context.NodeId] = context; + _nodeContexts[nodeKey] = context; + _nodeIdToNodeKey[context.NodeId] = nodeKey; // Start the asynchronous read. context.BeginAsyncPacketRead(); @@ -702,19 +741,20 @@ private void NodeContextCreated(NodeContext context) } /// - /// Method called when a context terminates. + /// Method called when a context terminates (called from CreateNode callbacks or ShutdownAllNodes). /// private void NodeContextTerminated(int nodeId) { - _nodeContexts.TryRemove(nodeId, out _); + // Remove from nodeKey-based lookup if we have it + if (_nodeIdToNodeKey.TryRemove(nodeId, out TaskHostNodeKey nodeKey)) + { + _nodeContexts.TryRemove(nodeKey, out _); + } // May also be removed by unnatural termination, so don't assume it's there lock (_activeNodes) { - if (_activeNodes.Contains(nodeId)) - { - _activeNodes.Remove(nodeId); - } + _activeNodes.Remove(nodeId); if (_activeNodes.Count == 0) { diff --git a/src/Build/Instance/TaskFactories/TaskHostTask.cs b/src/Build/Instance/TaskFactories/TaskHostTask.cs index 3d76fba1e1e..4f13c635e9d 100644 --- a/src/Build/Instance/TaskFactories/TaskHostTask.cs +++ b/src/Build/Instance/TaskFactories/TaskHostTask.cs @@ -28,12 +28,6 @@ namespace Microsoft.Build.BackEnd /// internal class TaskHostTask : IGeneratedTask, ICancelableTask, INodePacketFactory, INodePacketHandler { - private const int HANDSHAKE_OPTIONS_BITS = 9; - - private const int HANDSHAKE_OPTIONS_MASK = 0x1FF; - - private const int NODE_ID_MAX_VALUE_FOR_MULTITHREADED = 255; - /// /// The IBuildEngine callback object. /// @@ -100,9 +94,9 @@ internal class TaskHostTask : IGeneratedTask, ICancelableTask, INodePacketFactor private HandshakeOptions _requiredContext = HandshakeOptions.None; /// - /// The task host node ID of the task host we're launching. + /// The task host node key identifying the task host we're launching. /// - private int _taskHostNodeId; + private TaskHostNodeKey _taskHostNodeKey; /// /// The ID of the node on which this task is scheduled to run. @@ -270,7 +264,7 @@ public void Cancel() { if (_taskHostProvider != null && _connectedToTaskHost) { - _taskHostProvider.SendData(_taskHostNodeId, new TaskHostTaskCancelled()); + _taskHostProvider.SendData(_taskHostNodeKey, new TaskHostTaskCancelled()); } } @@ -339,8 +333,8 @@ public bool Execute() nodeReuse: _buildComponentHost.BuildParameters.EnableNodeReuse && _useSidecarTaskHost, taskHostParameters: _taskHostParameters); - _taskHostNodeId = GenerateTaskHostNodeId(_scheduledNodeId, _requiredContext); - _connectedToTaskHost = _taskHostProvider.AcquireAndSetUpHost(_requiredContext, _taskHostNodeId, this, this, hostConfiguration, _taskHostParameters); + _taskHostNodeKey = new TaskHostNodeKey(_requiredContext, _scheduledNodeId); + _connectedToTaskHost = _taskHostProvider.AcquireAndSetUpHost(_taskHostNodeKey, this, this, hostConfiguration, _taskHostParameters); } if (_connectedToTaskHost) @@ -369,7 +363,7 @@ public bool Execute() { lock (_taskHostLock) { - _taskHostProvider.DisconnectFromHost(_taskHostNodeId); + _taskHostProvider.DisconnectFromHost(_taskHostNodeKey); _connectedToTaskHost = false; } } @@ -391,22 +385,6 @@ public bool Execute() return _taskExecutionSucceeded; } - private static int GenerateTaskHostNodeId(int scheduledNodeId, HandshakeOptions handshakeOptions) - { - // For traditional multi-proc builds, the task host id is just (int)HandshakeOptions that represents the runtime / architecture. - // For new multi-threaded mode, NodeProviderOutOfProcTaskHost needs to distinguish task hosts not only by HandshakeOptions (runtime / architecture), - // but also by which node they were requested. This is because NodeProviderOutOfProcTaskHost runs on the same process as multiple in-proc nodes, - // as opposed to the traditional multi-proc case where each node and single NodeProviderOutOfProcTaskHost runs on its own process. - // nodeId: [1, 255] (8 bits more than enough) (max is number of processors, usually 8. Let's assume max is 256 processors) - // HandshakeOptions: [0, 511] (9 bits) - // Pack nodeId into upper bits, handshakeOptions into lower bits - ErrorUtilities.VerifyThrowArgumentOutOfRange(scheduledNodeId == -1 || (scheduledNodeId >= 1 && scheduledNodeId <= NODE_ID_MAX_VALUE_FOR_MULTITHREADED), nameof(scheduledNodeId)); - - return scheduledNodeId == -1 ? - (int)handshakeOptions : - (scheduledNodeId << HANDSHAKE_OPTIONS_BITS) | ((int)handshakeOptions & HANDSHAKE_OPTIONS_MASK); - } - /// /// Registers the specified handler for a particular packet type. /// diff --git a/src/Shared/CommunicationsUtilities.cs b/src/Shared/CommunicationsUtilities.cs index 530b23b164b..dd6c5ba02ac 100644 --- a/src/Shared/CommunicationsUtilities.cs +++ b/src/Shared/CommunicationsUtilities.cs @@ -37,7 +37,6 @@ namespace Microsoft.Build.Internal /// /// Enumeration of all possible (currently supported) options for handshakes. /// - /// In case of adding new options, please remember to update the generation of unique task host node id in NodeProviderOutOfProcTaskHost. [Flags] internal enum HandshakeOptions { @@ -89,6 +88,19 @@ internal enum HandshakeOptions SidecarTaskHost = 256, } + /// + /// Represents a unique key for identifying task host nodes. + /// Combines HandshakeOptions (which specify runtime/architecture configuration) with + /// the scheduled node ID to uniquely identify task hosts in multi-threaded mode. + /// + /// The handshake options specifying runtime and architecture configuration. + /// + /// The scheduled node ID. In traditional multi-proc builds, this is -1 (meaning the task host + /// is identified by HandshakeOptions alone). In multi-threaded mode, each in-proc node has + /// its own task host, so the node ID is used to distinguish them. + /// + internal readonly record struct TaskHostNodeKey(HandshakeOptions HandshakeOptions, int NodeId); + /// /// Status codes for the handshake process. /// It aggregates return values across several functions so we use an aggregate instead of a separate class for each method.