Avoid action delegate + async state machine allocations on packet header read

Author: ccastanedaucf

src/Build/BackEnd/Client/MSBuildClientPacketPump.cs

@@ -203,7 +203,7 @@ private void RunReadLoop(Stream localStream, ManualResetEvent localPacketPumpShu
 
                 // Use a separate reuseable wait handle to avoid allocating on Task.AsyncWaitHandle.
                 using AutoResetEvent readTaskEvent = new(false);
-                ValueTask<int> readTask = CommunicationsUtilities.ReadAsync(localStream, headerByte, headerByte.Length, readTaskEvent);
+                ValueTask<int> readTask = CommunicationsUtilities.ReadExactlyAsync(localStream, headerByte, headerByte.Length, readTaskEvent);
 
                 // Ordering of the wait handles is important. The first signalled wait handle in the array
                 // will be returned by WaitAny if multiple wait handles are signalled. We prefer to have the
@@ -295,7 +295,7 @@ private void RunReadLoop(Stream localStream, ManualResetEvent localPacketPumpShu
                                 else
                                 {
                                     // Start reading the next package header.
-                                    readTask = CommunicationsUtilities.ReadAsync(localStream, headerByte, headerByte.Length, readTaskEvent);
+                                    readTask = CommunicationsUtilities.ReadExactlyAsync(localStream, headerByte, headerByte.Length, readTaskEvent);
                                 }
                             }
                             break;

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

@@ -646,7 +646,7 @@ public async Task RunPacketReadLoopAsync()
                 {
                     try
                     {
-                        int bytesRead = await CommunicationsUtilities.ReadAsync(_clientToServerStream, _headerByte, _headerByte.Length);
+                        int bytesRead = await CommunicationsUtilities.ReadExactlyAsync(_clientToServerStream, _headerByte, _headerByte.Length);
                         if (!ProcessHeaderBytesRead(bytesRead))
                         {
                             return;

src/Shared/CommunicationsUtilities.cs

@@ -588,16 +588,30 @@ internal static int ReadIntForHandshake(this PipeStream stream, byte? byteToAcce
 
 #if !TASKHOST
         /// <summary>
-        /// Allow interop with EAP / Event-based wait handles without additional allocations.
-        /// By signalling an external reset event, this allows use of WaitHandle.WaitAny() in non-async/await contexts.
+        /// This is intended to get around state-machine and wait handle allocations on .NET Framework for async reads.
+        /// Prefer ReadAsync() when the read is expected to complete synchronously, or if the bytes to read are greater
+        /// than the stream's buffer and will require multiple reads (e.g. the packet body).
+        /// By signalling an external reset event, this also allows use of WaitHandle.WaitAny() in non-async/await contexts.
         /// </summary>
-        internal static async ValueTask<int> ReadAsync(Stream stream, byte[] buffer, int bytesToRead, AutoResetEvent autoResetEvent)
+        internal static ValueTask<int> ReadExactlyAsync(Stream stream, byte[] buffer, int bytesToRead, AutoResetEvent autoResetEvent = null)
         {
-            // Signal to the caller only after the read is complete.
-            int result = await ReadAsync(stream, buffer, bytesToRead).ConfigureAwait(false);
-            _ = autoResetEvent.Set();
+            Task<int> readTask = stream.ReadAsync(buffer, 0, bytesToRead);
 
-            return result;
+            // If the task completed synchronously, directly return the result.
+            if (readTask.IsCompleted)
+            {
+                _ = autoResetEvent?.Set();
+                return new ValueTask<int>(readTask.Result);
+            }
+
+            // Otherwise, a Task has been allocated and we'll need to set a callback.
+            readTask = readTask.ContinueWith(static (completedTask, state) =>
+            {
+                _ = ((AutoResetEvent)state)?.Set();
+                return completedTask.Result;
+            }, autoResetEvent, TaskContinuationOptions.ExecuteSynchronously);
+
+            return new ValueTask<int>(readTask);
         }
 
         internal static async ValueTask<int> ReadAsync(Stream stream, byte[] buffer, int bytesToRead)

src/Shared/NodeEndpointOutOfProcBase.cs

@@ -524,7 +524,7 @@ private void RunReadLoop(BufferedReadStream localReadPipe, NamedPipeServerStream
 #if !TASKHOST
             // Use a separate reuseable wait handle to avoid allocating on Task.AsyncWaitHandle.
             using AutoResetEvent readTaskEvent = new(false);
-            ValueTask<int> readTask = CommunicationsUtilities.ReadAsync(localReadPipe, headerByte, headerByte.Length, readTaskEvent);
+            ValueTask<int> readTask = CommunicationsUtilities.ReadExactlyAsync(localReadPipe, headerByte, headerByte.Length, readTaskEvent);
 #else
             IAsyncResult result = localReadPipe.BeginRead(headerByte, 0, headerByte.Length, null, null);
 #endif
@@ -617,7 +617,7 @@ private void RunReadLoop(BufferedReadStream localReadPipe, NamedPipeServerStream
                             }
 
 #if !TASKHOST
-                            readTask = CommunicationsUtilities.ReadAsync(localReadPipe, headerByte, headerByte.Length, readTaskEvent);
+                            readTask = CommunicationsUtilities.ReadExactlyAsync(localReadPipe, headerByte, headerByte.Length, readTaskEvent);
 #else
                             result = localReadPipe.BeginRead(headerByte, 0, headerByte.Length, null, null);
                             handles[0] = result.AsyncWaitHandle;