[IBuildEngine callbacks] Stage 2: RequestCores/ReleaseCores

src/Build.UnitTests/BackEnd/RequestCoresTask.cs

@@ -0,0 +1,52 @@
+// 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.Framework;
+using Microsoft.Build.Utilities;
+
+namespace Microsoft.Build.UnitTests.BackEnd
+{
+    /// <summary>
+    /// A test task that calls IBuildEngine9.RequestCores and optionally ReleaseCores.
+    /// Used by TaskHostCallback_Tests (in-process) and NetTaskHost_E2E_Tests (cross-runtime).
+    /// The E2E project includes this file via linked compile to avoid duplication.
+    /// </summary>
+    public class RequestCoresTask : Task
+    {
+        /// <summary>
+        /// Number of cores to request. Defaults to 1.
+        /// </summary>
+        public int CoreCount { get; set; } = 1;
+
+        /// <summary>
+        /// Whether to release the granted cores after requesting them.
+        /// </summary>
+        public bool ReleaseAfter { get; set; }
+
+        /// <summary>
+        /// Number of cores granted by the build engine.
+        /// </summary>
+        [Output]
+        public int GrantedCores { get; set; }
+
+        public override bool Execute()
+        {
+            if (BuildEngine is IBuildEngine9 engine9)
+            {
+                GrantedCores = engine9.RequestCores(CoreCount);
+                Log.LogMessage(MessageImportance.High, $"RequestCores({CoreCount}) = {GrantedCores}");
+
+                if (ReleaseAfter && GrantedCores > 0)
+                {
+                    engine9.ReleaseCores(GrantedCores);
+                    Log.LogMessage(MessageImportance.High, $"ReleaseCores({GrantedCores}) completed");
+                }
+
+                return true;
+            }
+
+            Log.LogError("BuildEngine does not implement IBuildEngine9");
+            return false;
+        }
+    }
+}

src/Build.UnitTests/BackEnd/TaskHostCallbackPacket_Tests.cs

@@ -46,5 +46,43 @@ public void TaskHostIsRunningMultipleNodesResponse_RoundTrip_Serialization(bool
             deserialized.IsRunningMultipleNodes.ShouldBe(isRunningMultipleNodes);
             deserialized.Type.ShouldBe(NodePacketType.TaskHostIsRunningMultipleNodesResponse);
         }
+
+        [Theory]
+        [InlineData(4, false)]  // RequestCores(4)
+        [InlineData(2, true)]   // ReleaseCores(2)
+        public void TaskHostCoresRequest_RoundTrip_Serialization(int cores, bool isRelease)
+        {
+            var request = new TaskHostCoresRequest(cores, isRelease);
+            request.RequestId = 77;
+
+            ITranslator writeTranslator = TranslationHelpers.GetWriteTranslator();
+            request.Translate(writeTranslator);
+
+            ITranslator readTranslator = TranslationHelpers.GetReadTranslator();
+            var deserialized = (TaskHostCoresRequest)TaskHostCoresRequest.FactoryForDeserialization(readTranslator);
+
+            deserialized.RequestId.ShouldBe(77);
+            deserialized.RequestedCores.ShouldBe(cores);
+            deserialized.IsRelease.ShouldBe(isRelease);
+            deserialized.Type.ShouldBe(NodePacketType.TaskHostCoresRequest);
+        }
+
+        [Theory]
+        [InlineData(0)]   // ReleaseCores acknowledgment
+        [InlineData(3)]   // RequestCores granted 3
+        public void TaskHostCoresResponse_RoundTrip_Serialization(int grantedCores)
+        {
+            var response = new TaskHostCoresResponse(99, grantedCores);
+
+            ITranslator writeTranslator = TranslationHelpers.GetWriteTranslator();
+            response.Translate(writeTranslator);
+
+            ITranslator readTranslator = TranslationHelpers.GetReadTranslator();
+            var deserialized = (TaskHostCoresResponse)TaskHostCoresResponse.FactoryForDeserialization(readTranslator);
+
+            deserialized.RequestId.ShouldBe(99);
+            deserialized.GrantedCores.ShouldBe(grantedCores);
+            deserialized.Type.ShouldBe(NodePacketType.TaskHostCoresResponse);
+        }
     }
 }

src/Build.UnitTests/BackEnd/TaskHostCallback_Tests.cs

@@ -140,5 +140,143 @@ public void IsRunningMultipleNodes_LogsErrorWhenCallbacksNotSupported()
             logger.ErrorCount.ShouldBeGreaterThan(0);
             logger.FullLog.ShouldContain("MSB5022");
         }
+
+        /// <summary>
+        /// Verifies RequestCores callback works when task is explicitly run in TaskHost via TaskHostFactory.
+        /// The first RequestCores call should always return at least 1 (the implicit core).
+        /// </summary>
+        [Fact]
+        public void RequestCores_WorksWithExplicitTaskHostFactory()
+        {
+            using TestEnvironment env = TestEnvironment.Create(_output);
+            env.SetEnvironmentVariable("MSBUILDENABLETASKHOSTCALLBACKS", "1");
+
+            string projectContents = $@"
+<Project>
+    <UsingTask TaskName=""{nameof(RequestCoresTask)}"" AssemblyFile=""{typeof(RequestCoresTask).Assembly.Location}"" TaskFactory=""TaskHostFactory"" />
+    <Target Name=""Test"">
+        <{nameof(RequestCoresTask)} CoreCount=""2"">
+            <Output PropertyName=""Result"" TaskParameter=""GrantedCores"" />
+        </{nameof(RequestCoresTask)}>
+    </Target>
+</Project>";
+
+            TransientTestProjectWithFiles project = env.CreateTestProjectWithFiles(projectContents);
+            ProjectInstance projectInstance = new(project.ProjectFile);
+
+            var logger = new MockLogger(_output);
+            BuildResult buildResult = BuildManager.DefaultBuildManager.Build(
+                new BuildParameters { MaxNodeCount = 4, EnableNodeReuse = false, Loggers = [logger] },
+                new BuildRequestData(projectInstance, targetsToBuild: ["Test"]));
+
+            buildResult.OverallResult.ShouldBe(BuildResultCode.Success);
+            // First RequestCores call gets at least the implicit core
+            int.Parse(projectInstance.GetPropertyValue("Result")).ShouldBeGreaterThanOrEqualTo(1);
+        }
+
+        /// <summary>
+        /// Verifies RequestCores + ReleaseCores works end-to-end when task runs in TaskHost.
+        /// </summary>
+        [Fact]
+        public void RequestAndReleaseCores_WorksWithExplicitTaskHostFactory()
+        {
+            using TestEnvironment env = TestEnvironment.Create(_output);
+            env.SetEnvironmentVariable("MSBUILDENABLETASKHOSTCALLBACKS", "1");
+
+            string projectContents = $@"
+<Project>
+    <UsingTask TaskName=""{nameof(RequestCoresTask)}"" AssemblyFile=""{typeof(RequestCoresTask).Assembly.Location}"" TaskFactory=""TaskHostFactory"" />
+    <Target Name=""Test"">
+        <{nameof(RequestCoresTask)} CoreCount=""2"" ReleaseAfter=""true"">
+            <Output PropertyName=""Result"" TaskParameter=""GrantedCores"" />
+        </{nameof(RequestCoresTask)}>
+    </Target>
+</Project>";
+
+            TransientTestProjectWithFiles project = env.CreateTestProjectWithFiles(projectContents);
+            ProjectInstance projectInstance = new(project.ProjectFile);
+
+            var logger = new MockLogger(_output);
+            BuildResult buildResult = BuildManager.DefaultBuildManager.Build(
+                new BuildParameters { MaxNodeCount = 4, EnableNodeReuse = false, Loggers = [logger] },
+                new BuildRequestData(projectInstance, targetsToBuild: ["Test"]));
+
+            buildResult.OverallResult.ShouldBe(BuildResultCode.Success);
+            // Verify both RequestCores and ReleaseCores ran without error
+            logger.AssertNoErrors();
+            logger.FullLog.ShouldContain("ReleaseCores(");
+        }
+
+        /// <summary>
+        /// Verifies RequestCores callback works when task is auto-ejected to TaskHost in multithreaded mode.
+        /// </summary>
+        [Fact]
+        public void RequestCores_WorksWhenAutoEjectedInMultiThreadedMode()
+        {
+            using TestEnvironment env = TestEnvironment.Create(_output);
+            env.SetEnvironmentVariable("MSBUILDENABLETASKHOSTCALLBACKS", "1");
+            string testDir = env.CreateFolder().Path;
+
+            // RequestCoresTask lacks MSBuildMultiThreadableTask attribute, so it's auto-ejected to TaskHost in MT mode
+            string projectContents = $@"
+<Project>
+    <UsingTask TaskName=""{nameof(RequestCoresTask)}"" AssemblyFile=""{typeof(RequestCoresTask).Assembly.Location}"" />
+    <Target Name=""Test"">
+        <{nameof(RequestCoresTask)} CoreCount=""1"">
+            <Output PropertyName=""Result"" TaskParameter=""GrantedCores"" />
+        </{nameof(RequestCoresTask)}>
+    </Target>
+</Project>";
+
+            string projectFile = Path.Combine(testDir, "Test.proj");
+            File.WriteAllText(projectFile, projectContents);
+
+            var logger = new MockLogger(_output);
+            BuildResult buildResult = BuildManager.DefaultBuildManager.Build(
+                new BuildParameters
+                {
+                    MultiThreaded = true,
+                    MaxNodeCount = 4,
+                    Loggers = [logger],
+                    EnableNodeReuse = false
+                },
+                new BuildRequestData(projectFile, new Dictionary<string, string?>(), null, ["Test"], null));
+
+            buildResult.OverallResult.ShouldBe(BuildResultCode.Success);
+            logger.FullLog.ShouldContain("external task host");
+            logger.FullLog.ShouldContain("RequestCores(1) =");
+        }
+
+        /// <summary>
+        /// Verifies that RequestCores when callbacks are disabled logs error MSB5022 and returns 0.
+        /// </summary>
+        [Fact]
+        public void RequestCores_LogsErrorWhenCallbacksNotSupported()
+        {
+            using TestEnvironment env = TestEnvironment.Create(_output);
+
+            // Explicitly do NOT set MSBUILDENABLETASKHOSTCALLBACKS
+            string projectContents = $@"
+<Project>
+    <UsingTask TaskName=""{nameof(RequestCoresTask)}"" AssemblyFile=""{typeof(RequestCoresTask).Assembly.Location}"" TaskFactory=""TaskHostFactory"" />
+    <Target Name=""Test"">
+        <{nameof(RequestCoresTask)} CoreCount=""2"">
+            <Output PropertyName=""Result"" TaskParameter=""GrantedCores"" />
+        </{nameof(RequestCoresTask)}>
+    </Target>
+</Project>";
+
+            TransientTestProjectWithFiles project = env.CreateTestProjectWithFiles(projectContents);
+            ProjectInstance projectInstance = new(project.ProjectFile);
+
+            var logger = new MockLogger(_output);
+            BuildResult buildResult = BuildManager.DefaultBuildManager.Build(
+                new BuildParameters { MaxNodeCount = 4, EnableNodeReuse = false, Loggers = [logger] },
+                new BuildRequestData(projectInstance, targetsToBuild: ["Test"]));
+
+            // MSB5022 error should be logged — the callback was not forwarded
+            logger.ErrorCount.ShouldBeGreaterThan(0);
+            logger.FullLog.ShouldContain("MSB5022");
+        }
     }
 }

src/Build.UnitTests/NetTaskHost_E2E_Tests.cs

@@ -162,6 +162,33 @@ public void NetTaskHost_CallbackIsRunningMultipleNodesTest()
             testTaskOutput.ShouldContain("CallbackResult: IsRunningMultipleNodes = False");
         }
 
+        [WindowsFullFrameworkOnlyFact]
+        public void NetTaskHost_CallbackRequestCoresTest()
+        {
+            using TestEnvironment env = TestEnvironment.Create(_output);
+            env.SetEnvironmentVariable("MSBUILDENABLETASKHOSTCALLBACKS", "1");
+
+            // Point dotnet resolution to the bootstrap layout so the .NET Core TaskHost
+            // uses the locally-built MSBuild.dll (with callback support) instead of the system SDK.
+            string bootstrapCorePath = Path.Combine(RunnerUtilities.BootstrapRootPath, "core");
+            string bootstrapDotnet = Path.Combine(bootstrapCorePath, "dotnet.exe");
+            env.SetEnvironmentVariable("DOTNET_HOST_PATH", bootstrapDotnet);
+            env.SetEnvironmentVariable("DOTNET_ROOT", bootstrapCorePath);
+            env.SetEnvironmentVariable("DOTNET_MSBUILD_SDK_RESOLVER_CLI_DIR", bootstrapCorePath);
+
+            string testProjectPath = Path.Combine(TestAssetsRootPath, "ExampleNetTask", "TestNetTaskResourceCallback", "TestNetTaskResourceCallback.csproj");
+
+            string testTaskOutput = RunnerUtilities.ExecBootstrapedMSBuild($"{testProjectPath} -v:n -p:LatestDotNetCoreForMSBuild={RunnerUtilities.LatestDotNetCoreForMSBuild} -t:TestTask", out bool successTestTask);
+
+            if (!successTestTask)
+            {
+                _output.WriteLine(testTaskOutput);
+            }
+
+            successTestTask.ShouldBeTrue();
+            testTaskOutput.ShouldContain("CallbackResult: RequestCores(2) =");
+        }
+
         [WindowsFullFrameworkOnlyFact]
         public void NetTaskWithImplicitHostParamsTest()
         {

src/Build.UnitTests/TestAssets/ExampleNetTask/ExampleTask/ExampleTask.csproj

@@ -16,6 +16,7 @@
 
   <ItemGroup>
     <Compile Include="..\..\..\BackEnd\IsRunningMultipleNodesTask.cs" Link="IsRunningMultipleNodesTask.cs" />
+    <Compile Include="..\..\..\BackEnd\RequestCoresTask.cs" Link="RequestCoresTask.cs" />
   </ItemGroup>
 
 </Project>

src/Build.UnitTests/TestAssets/ExampleNetTask/TestNetTaskResourceCallback/TestNetTaskResourceCallback.csproj

@@ -0,0 +1,25 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>$(LatestDotNetCoreForMSBuild)</TargetFramework>
+  </PropertyGroup>
+
+  <PropertyGroup>
+    <TestProjectFolder>$([System.IO.Path]::GetFullPath('$([System.IO.Path]::Combine('$(MSBuildThisFileDirectory)', '..', '..', '..', '..'))'))</TestProjectFolder>
+    <ExampleTaskPath>$([System.IO.Path]::Combine('$(TestProjectFolder)', '$(TargetFramework)', 'ExampleTask.dll'))</ExampleTaskPath>
+  </PropertyGroup>
+
+  <UsingTask
+      TaskName="Microsoft.Build.UnitTests.BackEnd.RequestCoresTask"
+      AssemblyFile="$(ExampleTaskPath)"
+      TaskFactory="TaskHostFactory"
+      Runtime="NET"/>
+
+  <Target Name="TestTask" BeforeTargets="Build">
+    <Microsoft.Build.UnitTests.BackEnd.RequestCoresTask CoreCount="2" ReleaseAfter="true">
+      <Output PropertyName="Granted" TaskParameter="GrantedCores" />
+    </Microsoft.Build.UnitTests.BackEnd.RequestCoresTask>
+    <Message Text="CallbackResult: RequestCores(2) = $(Granted)" Importance="high" />
+  </Target>
+
+</Project>

src/Build.UnitTests/TestAssets/ExampleNetTask/TestNetTaskResourceCallback/global.json

@@ -0,0 +1,2 @@
+{
+}

src/Build/Instance/TaskFactories/TaskHostTask.cs

@@ -204,6 +204,7 @@ public TaskHostTask(
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostTaskComplete, TaskHostTaskComplete.FactoryForDeserialization, this);
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.NodeShutdown, NodeShutdown.FactoryForDeserialization, this);
             (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostIsRunningMultipleNodesRequest, TaskHostIsRunningMultipleNodesRequest.FactoryForDeserialization, this);
+            (this as INodePacketFactory).RegisterPacketHandler(NodePacketType.TaskHostCoresRequest, TaskHostCoresRequest.FactoryForDeserialization, this);
 
             _packetReceivedEvent = new AutoResetEvent(false);
             _receivedPackets = new ConcurrentQueue<INodePacket>();
@@ -513,6 +514,9 @@ private void HandlePacket(INodePacket packet, out bool taskFinished)
                 case NodePacketType.TaskHostIsRunningMultipleNodesRequest:
                     HandleIsRunningMultipleNodesRequest(packet as TaskHostIsRunningMultipleNodesRequest);
                     break;
+                case NodePacketType.TaskHostCoresRequest:
+                    HandleCoresRequest(packet as TaskHostCoresRequest);
+                    break;
                 default:
                     ErrorUtilities.ThrowInternalErrorUnreachable();
                     break;
@@ -662,6 +666,34 @@ private void HandleIsRunningMultipleNodesRequest(TaskHostIsRunningMultipleNodesR
             _taskHostProvider.SendData(_taskHostNodeKey, response);
         }
 
+        /// <summary>
+        /// Handle RequestCores/ReleaseCores request from the TaskHost.
+        /// Forwards the call to the in-process TaskHost's IBuildEngine9 implementation,
+        /// which handles implicit core accounting and scheduler communication.
+        /// </summary>
+        private void HandleCoresRequest(TaskHostCoresRequest request)
+        {
+            int grantedCores = 0;
+
+            if (request.IsRelease)
+            {
+                if (_buildEngine is IBuildEngine9 engine9)
+                {
+                    engine9.ReleaseCores(request.RequestedCores);
+                }
+            }
+            else
+            {
+                if (_buildEngine is IBuildEngine9 engine9)
+                {
+                    grantedCores = engine9.RequestCores(request.RequestedCores);
+                }
+            }
+
+            var response = new TaskHostCoresResponse(request.RequestId, grantedCores);
+            _taskHostProvider.SendData(_taskHostNodeKey, response);
+        }
+
         /// <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

@@ -125,6 +125,8 @@
     <Compile Include="..\Shared\ITaskHostCallbackPacket.cs" />
     <Compile Include="..\Shared\TaskHostIsRunningMultipleNodesRequest.cs" />
     <Compile Include="..\Shared\TaskHostIsRunningMultipleNodesResponse.cs" />
+    <Compile Include="..\Shared\TaskHostCoresRequest.cs" />
+    <Compile Include="..\Shared\TaskHostCoresResponse.cs" />
     <Compile Include="..\Shared\OutOfProcTaskHostTaskResult.cs" />
     <Compile Include="..\Shared\TaskLoader.cs" />
     <Compile Include="..\Shared\NodeEngineShutdownReason.cs" />

src/MSBuild/MSBuild.csproj

@@ -120,6 +120,8 @@
     <Compile Include="..\Shared\ITaskHostCallbackPacket.cs" />
     <Compile Include="..\Shared\TaskHostIsRunningMultipleNodesRequest.cs" />
     <Compile Include="..\Shared\TaskHostIsRunningMultipleNodesResponse.cs" />
+    <Compile Include="..\Shared\TaskHostCoresRequest.cs" />
+    <Compile Include="..\Shared\TaskHostCoresResponse.cs" />
     <Compile Include="..\Shared\TaskLoader.cs" />
     <Compile Include="..\Shared\MSBuildLoadContext.cs" Condition="'$(TargetFrameworkIdentifier)'!='.NETFramework'" />
     <Compile Include="..\Shared\TypeLoader.cs" />