Implement IBuildEngine3 Yield/Reacquire for out-of-proc task host

src/Build.UnitTests/BackEnd/TaskHost_Tests.cs

@@ -546,6 +546,18 @@ public void IsRunningMultipleNodes4Nodes()
             Assert.True(_taskHost.IsRunningMultipleNodes); // "Expect IsRunningMultipleNodes to be true with 4 nodes"
         }
 
+        /// <summary>
+        /// Test that Yield and Reacquire can be called without throwing exceptions.
+        /// The mock request callback should handle these calls appropriately.
+        /// </summary>
+        [Fact]
+        public void TestYieldAndReacquire()
+        {
+            // Should not throw exceptions when called
+            _taskHost.Yield();
+            _taskHost.Reacquire();
+        }
+
 #if FEATURE_CODETASKFACTORY
         /// <summary>
         /// Task logging after it's done should not crash us.

src/Build/Instance/TaskFactories/TaskHostTask.cs

@@ -191,6 +191,8 @@ public TaskHostTask(
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.LogMessage, LogMessagePacket.FactoryForDeserialization, this);
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostTaskComplete, TaskHostTaskComplete.FactoryForDeserialization, this);
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.NodeShutdown, NodeShutdown.FactoryForDeserialization, this);
+            (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostYield, TaskHostYield.FactoryForDeserialization, this);
+            (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostReacquire, TaskHostReacquire.FactoryForDeserialization, this);
 
             _packetReceivedEvent = new AutoResetEvent(false);
             _receivedPackets = new ConcurrentQueue<INodePacket>();
@@ -508,6 +510,12 @@ private void HandlePacket(INodePacket packet, out bool taskFinished)
                 case NodePacketType.LogMessage:
                     HandleLoggedMessage(packet as LogMessagePacket);
                     break;
+                case NodePacketType.TaskHostYield:
+                    HandleTaskHostYield(packet as TaskHostYield);
+                    break;
+                case NodePacketType.TaskHostReacquire:
+                    HandleTaskHostReacquire(packet as TaskHostReacquire);
+                    break;
                 default:
                     ErrorUtilities.ThrowInternalErrorUnreachable();
                     break;
@@ -647,6 +655,28 @@ private void HandleLoggedMessage(LogMessagePacket logMessagePacket)
             }
         }
 
+        /// <summary>
+        /// Handle yield notification from the task host.
+        /// The task host is indicating it's about to perform long-running operations
+        /// and the node can do other work. We acknowledge receipt immediately.
+        /// </summary>
+        private void HandleTaskHostYield(TaskHostYield yieldPacket)
+        {
+            // Send acknowledgment back to task host
+            _taskHostProvider.SendData(_taskHostNodeId, new TaskHostYield());
+        }
+
+        /// <summary>
+        /// Handle reacquire notification from the task host.
+        /// The task host wants to regain control after previously yielding.
+        /// We acknowledge receipt immediately.
+        /// </summary>
+        private void HandleTaskHostReacquire(TaskHostReacquire reacquirePacket)
+        {
+            // Send acknowledgment back to task host
+            _taskHostProvider.SendData(_taskHostNodeId, new TaskHostReacquire());
+        }
+
         /// <summary>
         /// Since we log that we weren't able to connect to the task host in a couple of different places,
         /// extract it out into a separate method.

src/Build/Microsoft.Build.csproj

@@ -100,6 +100,8 @@
     <Compile Include="..\Shared\TaskHostConfiguration.cs" />
     <Compile Include="..\Shared\TaskHostTaskCancelled.cs" />
     <Compile Include="..\Shared\TaskHostTaskComplete.cs" />
+    <Compile Include="..\Shared\TaskHostYield.cs" />
+    <Compile Include="..\Shared\TaskHostReacquire.cs" />
     <Compile Include="..\Shared\OutOfProcTaskHostTaskResult.cs" />
     <Compile Include="..\Shared\TaskLoader.cs" />
     <Compile Include="..\Shared\NodeEngineShutdownReason.cs" />

src/MSBuild/MSBuild.csproj

@@ -113,6 +113,8 @@
     <Compile Include="..\Shared\TaskHostTaskComplete.cs" />
     <Compile Include="..\Shared\OutOfProcTaskHostTaskResult.cs" />
     <Compile Include="..\Shared\TaskHostTaskCancelled.cs" />
+    <Compile Include="..\Shared\TaskHostYield.cs" />
+    <Compile Include="..\Shared\TaskHostReacquire.cs" />
     <Compile Include="..\Shared\TaskLoader.cs" />
     <Compile Include="..\Shared\MSBuildLoadContext.cs" Condition="'$(TargetFrameworkIdentifier)'!='.NETFramework'" />
     <Compile Include="..\Shared\TypeLoader.cs" />

src/MSBuild/OutOfProcTaskHostNode.cs

@@ -130,6 +130,16 @@ internal class OutOfProcTaskHostNode :
         /// </summary>
         private ManualResetEvent _taskCancelledEvent;
 
+        /// <summary>
+        /// The event which is set when the parent node acknowledges a yield request.
+        /// </summary>
+        private AutoResetEvent _yieldAcknowledgedEvent;
+
+        /// <summary>
+        /// The event which is set when the parent node acknowledges a reacquire request.
+        /// </summary>
+        private AutoResetEvent _reacquireAcknowledgedEvent;
+
         /// <summary>
         /// The thread currently executing user task in the TaskRunner
         /// </summary>
@@ -195,6 +205,8 @@ public OutOfProcTaskHostNode()
             _shutdownEvent = new ManualResetEvent(false);
             _taskCompleteEvent = new AutoResetEvent(false);
             _taskCancelledEvent = new ManualResetEvent(false);
+            _yieldAcknowledgedEvent = new AutoResetEvent(false);
+            _reacquireAcknowledgedEvent = new AutoResetEvent(false);
 
             _packetFactory = new NodePacketFactory();
 
@@ -203,6 +215,8 @@ public OutOfProcTaskHostNode()
             thisINodePacketFactory.RegisterPacketHandler(NodePacketType.TaskHostConfiguration, TaskHostConfiguration.FactoryForDeserialization, this);
             thisINodePacketFactory.RegisterPacketHandler(NodePacketType.TaskHostTaskCancelled, TaskHostTaskCancelled.FactoryForDeserialization, this);
             thisINodePacketFactory.RegisterPacketHandler(NodePacketType.NodeBuildComplete, NodeBuildComplete.FactoryForDeserialization, this);
+            thisINodePacketFactory.RegisterPacketHandler(NodePacketType.TaskHostYield, TaskHostYield.FactoryForDeserialization, this);
+            thisINodePacketFactory.RegisterPacketHandler(NodePacketType.TaskHostReacquire, TaskHostReacquire.FactoryForDeserialization, this);
 
 #if !CLR2COMPATIBILITY
             EngineServices = new EngineServicesImpl(this);
@@ -405,21 +419,36 @@ public BuildEngineResult BuildProjectFilesInParallel(string[] projectFileNames,
         }
 
         /// <summary>
-        /// Stub implementation of IBuildEngine3.Yield.  The task host does not support yielding, so just go ahead and silently
-        /// return, letting the task continue.
+        /// Implementation of IBuildEngine3.Yield. Sends a yield packet to the parent node
+        /// indicating that the task host is yielding control and the node can do other work
+        /// while the task performs long-running out-of-process operations.
         /// </summary>
         public void Yield()
         {
-            return;
+            if (_nodeEndpoint?.LinkStatus == LinkStatus.Active)
+            {
+                TaskHostYield yieldPacket = new TaskHostYield();
+                _nodeEndpoint.SendData(yieldPacket);
+
+                // Wait for the parent node to acknowledge the yield
+                _yieldAcknowledgedEvent.WaitOne();
+            }
         }
 
         /// <summary>
-        /// Stub implementation of IBuildEngine3.Reacquire. The task host does not support yielding, so just go ahead and silently
-        /// return, letting the task continue.
+        /// Implementation of IBuildEngine3.Reacquire. Sends a reacquire packet to the parent node
+        /// indicating that the task host wants to regain control after previously yielding.
         /// </summary>
         public void Reacquire()
         {
-            return;
+            if (_nodeEndpoint?.LinkStatus == LinkStatus.Active)
+            {
+                TaskHostReacquire reacquirePacket = new TaskHostReacquire();
+                _nodeEndpoint.SendData(reacquirePacket);
+
+                // Wait for the parent node to acknowledge the reacquire
+                _reacquireAcknowledgedEvent.WaitOne();
+            }
         }
 
         #endregion // IBuildEngine3 Implementation
@@ -725,6 +754,14 @@ private void HandlePacket(INodePacket packet)
                 case NodePacketType.NodeBuildComplete:
                     HandleNodeBuildComplete(packet as NodeBuildComplete);
                     break;
+                case NodePacketType.TaskHostYield:
+                    // Acknowledgment from parent node
+                    _yieldAcknowledgedEvent.Set();
+                    break;
+                case NodePacketType.TaskHostReacquire:
+                    // Acknowledgment from parent node
+                    _reacquireAcknowledgedEvent.Set();
+                    break;
             }
         }
 
@@ -875,11 +912,15 @@ private NodeEngineShutdownReason HandleShutdown()
             _shutdownEvent.Close();
             _taskCompleteEvent.Close();
             _taskCancelledEvent.Close();
+            _yieldAcknowledgedEvent.Close();
+            _reacquireAcknowledgedEvent.Close();
 #else
             _packetReceivedEvent.Dispose();
             _shutdownEvent.Dispose();
             _taskCompleteEvent.Dispose();
             _taskCancelledEvent.Dispose();
+            _yieldAcknowledgedEvent.Dispose();
+            _reacquireAcknowledgedEvent.Dispose();
 #endif
 
             return _shutdownReason;

src/Shared/INodePacket.cs

@@ -208,7 +208,6 @@ internal enum NodePacketType : byte
         /// </summary>
         ProcessReport, // 0x13
 
-
         /// Notifies the RAR node to set a configuration for a particular build.
         RarNodeEndpointConfiguration,
 
@@ -222,6 +221,26 @@ internal enum NodePacketType : byte
         /// </summary>
         RarNodeExecuteResponse, // 0x15
 
+        /// <summary>
+        /// Message sent from the task host to the parent node when a task
+        /// calls IBuildEngine3.Yield() to indicate it is yielding control
+        /// and the node can do other work.
+        ///
+        /// Contents:
+        /// (nothing)
+        /// </summary>
+        TaskHostYield, // 0x16
+
+        /// <summary>
+        /// Message sent from the task host to the parent node when a task
+        /// calls IBuildEngine3.Reacquire() to indicate it wants to regain
+        /// control after yielding.
+        ///
+        /// Contents:
+        /// (nothing)
+        /// </summary>
+        TaskHostReacquire, // 0x17
+
         // Reserve space for future core packet types (0x16-0x3B available for expansion)
 
         // Server command packets placed at end of safe range to maintain separation from core packets

src/Shared/TaskHostReacquire.cs

@@ -0,0 +1,48 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#nullable disable
+
+namespace Microsoft.Build.BackEnd
+{
+    /// <summary>
+    /// TaskHostReacquire informs the parent node that the task host
+    /// wants to reacquire control after previously yielding via Yield().
+    /// </summary>
+    internal class TaskHostReacquire : INodePacket
+    {
+        /// <summary>
+        /// Constructor
+        /// </summary>
+        public TaskHostReacquire()
+        {
+        }
+
+        /// <summary>
+        /// The type of this NodePacket
+        /// </summary>
+        public NodePacketType Type
+        {
+            get { return NodePacketType.TaskHostReacquire; }
+        }
+
+        /// <summary>
+        /// Translates the packet to/from binary form.
+        /// </summary>
+        /// <param name="translator">The translator to use.</param>
+        public void Translate(ITranslator translator)
+        {
+            // Do nothing -- this packet doesn't contain any parameters.
+        }
+
+        /// <summary>
+        /// Factory for deserialization.
+        /// </summary>
+        internal static INodePacket FactoryForDeserialization(ITranslator translator)
+        {
+            TaskHostReacquire taskReacquire = new TaskHostReacquire();
+            taskReacquire.Translate(translator);
+            return taskReacquire;
+        }
+    }
+}

src/Shared/TaskHostYield.cs

@@ -0,0 +1,49 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#nullable disable
+
+namespace Microsoft.Build.BackEnd
+{
+    /// <summary>
+    /// TaskHostYield informs the parent node that the task host is
+    /// yielding control and the node can do other work while the task
+    /// performs long-running out-of-process operations.
+    /// </summary>
+    internal class TaskHostYield : INodePacket
+    {
+        /// <summary>
+        /// Constructor
+        /// </summary>
+        public TaskHostYield()
+        {
+        }
+
+        /// <summary>
+        /// The type of this NodePacket
+        /// </summary>
+        public NodePacketType Type
+        {
+            get { return NodePacketType.TaskHostYield; }
+        }
+
+        /// <summary>
+        /// Translates the packet to/from binary form.
+        /// </summary>
+        /// <param name="translator">The translator to use.</param>
+        public void Translate(ITranslator translator)
+        {
+            // Do nothing -- this packet doesn't contain any parameters.
+        }
+
+        /// <summary>
+        /// Factory for deserialization.
+        /// </summary>
+        internal static INodePacket FactoryForDeserialization(ITranslator translator)
+        {
+            TaskHostYield taskYield = new TaskHostYield();
+            taskYield.Translate(translator);
+            return taskYield;
+        }
+    }
+}