MTP dotnet test: Better handling for pipe failures (#50571)

Author: Youssef1313

src/Cli/dotnet/Commands/CliCommandStrings.resx

@@ -2648,4 +2648,19 @@ Proceed?</value>
   <data name="CmdMinimumExpectedTestsDescription" xml:space="preserve">
     <value>Specifies the minimum number of tests that are expected to run.</value>
   </data>
-</root>
+  <data name="DotnetTestPipeOverlapping" xml:space="preserve">
+    <value>'dotnet' unexpectedly received overlapping messages from the 'dotnet test' named pipe.</value>
+  </data>
+  <data name="DotnetTestPipeIncompleteSize" xml:space="preserve">
+    <value>'dotnet' unexpectedly received less than 4 bytes from the 'dotnet test' named pipe.</value>
+  </data>
+  <data name="InternalLoopAsyncDidNotExitSuccessfullyErrorMessage" xml:space="preserve">
+    <value>Method '{0}' did not exit successfully</value>
+  </data>
+  <data name="DotnetTestPipeFailureHasHandshake" xml:space="preserve">
+    <value>Error disposing 'NamedPipeServer' corresponding to handshake:</value>
+  </data>
+  <data name="DotnetTestPipeFailureWithoutHandshake" xml:space="preserve">
+    <value>Error disposing 'NamedPipeServer', and no handshake was found.</value>
+  </data>
+</root>

src/Cli/dotnet/Commands/Test/IPC/NamedPipeServer.cs

@@ -4,14 +4,15 @@
 #nullable disable
 
 using System.Buffers;
+using System.Diagnostics;
 using System.Globalization;
 using System.IO.Pipes;
 
 namespace Microsoft.DotNet.Cli.Commands.Test.IPC;
 
 internal sealed class NamedPipeServer : NamedPipeBase
 {
-    private readonly Func<IRequest, Task<IResponse>> _callback;
+    private readonly Func<NamedPipeServer, IRequest, Task<IResponse>> _callback;
     private readonly NamedPipeServerStream _namedPipeServerStream;
     private readonly CancellationToken _cancellationToken;
     private readonly MemoryStream _serializationBuffer = new();
@@ -24,7 +25,7 @@ internal sealed class NamedPipeServer : NamedPipeBase
 
     public NamedPipeServer(
         PipeNameDescription pipeNameDescription,
-        Func<IRequest, Task<IResponse>> callback,
+        Func<NamedPipeServer, IRequest, Task<IResponse>> callback,
         int maxNumberOfServerInstances,
         CancellationToken cancellationToken,
         bool skipUnknownMessages)
@@ -55,14 +56,6 @@ public async Task WaitConnectionAsync(CancellationToken cancellationToken)
                     // We are being cancelled, so we don't need to wait anymore
                     return;
                 }
-                catch (Exception ex)
-                {
-                    // TODO: it might be better idea to let this exception bubble to
-                    // TestApplication. There, we should handle it by printing the exception without failing
-                    // the whole dotnet process, and provide more info about the specific test app that failed.
-                    // Potentially even knowing whether it's test host or test host controller that failed?
-                    Environment.FailFast($"[NamedPipeServer] Unhandled exception:{Environment.NewLine}{ex}", ex);
-                }
             }, cancellationToken);
     }
 
@@ -94,6 +87,11 @@ private async Task InternalLoopAsync(CancellationToken cancellationToken)
             if (currentMessageSize == 0)
             {
                 // We need to read the message size, first 4 bytes
+                if (currentReadBytes < sizeof(int))
+                {
+                    throw new UnreachableException(CliCommandStrings.DotnetTestPipeIncompleteSize);
+                }
+
                 currentMessageSize = BitConverter.ToInt32(_readBuffer, 0);
                 missingBytesToReadOfCurrentChunk = currentReadBytes - sizeof(int);
                 missingBytesToReadOfWholeMessage = currentMessageSize;
@@ -107,6 +105,11 @@ private async Task InternalLoopAsync(CancellationToken cancellationToken)
                 missingBytesToReadOfWholeMessage -= missingBytesToReadOfCurrentChunk;
             }
 
+            if (missingBytesToReadOfWholeMessage < 0)
+            {
+                throw new UnreachableException(CliCommandStrings.DotnetTestPipeOverlapping);
+            }
+
             // If we have read all the message, we can deserialize it
             if (missingBytesToReadOfWholeMessage == 0)
             {
@@ -124,7 +127,7 @@ private async Task InternalLoopAsync(CancellationToken cancellationToken)
                 var deserializedObject = (IRequest)requestNamedPipeSerializer.Deserialize(_messageBuffer);
 
                 // Call the callback
-                IResponse response = await _callback(deserializedObject);
+                IResponse response = await _callback(this, deserializedObject);
 
                 // Write the message size
                 _messageBuffer.Position = 0;
@@ -205,21 +208,28 @@ public void Dispose()
             return;
         }
 
-        if (WasConnected)
+        try
         {
-            // If the loop task is null at this point we have race condition, means that the task didn't start yet and we already dispose.
-            // This is unexpected and we throw an exception.
-
-            // To close gracefully we need to ensure that the client closed the stream line 103.
-            if (!_loopTask.Wait(TimeSpan.FromSeconds(90)))
+            if (WasConnected)
             {
-                throw new InvalidOperationException("InternalLoopAsyncDidNotExitSuccessfullyErrorMessage");
+                // If the loop task is null at this point we have race condition, means that the task didn't start yet and we already dispose.
+                // This is unexpected and we throw an exception.
+
+                // To close gracefully we need to ensure that the client closed the stream line 103.
+                if (!_loopTask.Wait(TimeSpan.FromSeconds(90)))
+                {
+                    throw new InvalidOperationException(CliCommandStrings.InternalLoopAsyncDidNotExitSuccessfullyErrorMessage);
+                }
             }
         }
+        finally
+        {
+            // Ensure we are still disposing the resouces correctly, even if _loopTask completes with
+            // an exception, or if the task doesn't complete within the 90 seconds limit.
+            _namedPipeServerStream.Dispose();
+            PipeName.Dispose();
 
-        _namedPipeServerStream.Dispose();
-        PipeName.Dispose();
-
-        _disposed = true;
+            _disposed = true;
+        }
     }
 }

src/Cli/dotnet/Commands/Test/TestApplication.cs

@@ -23,6 +23,7 @@ internal sealed class TestApplication(TestModule module, BuildOptions buildOptio
 
     private Task _testAppPipeConnectionLoop;
     private readonly List<NamedPipeServer> _testAppPipeConnections = [];
+    private readonly Dictionary<NamedPipeServer, HandshakeMessage> _handshakes = new();
 
     public event EventHandler<HandshakeArgs> HandshakeReceived;
     public event EventHandler<HelpEventArgs> HelpRequested;
@@ -35,6 +36,8 @@ internal sealed class TestApplication(TestModule module, BuildOptions buildOptio
 
     public TestModule Module { get; } = module;
 
+    public bool HasFailureDuringDispose { get; private set; }
+
     public async Task<int> RunAsync(TestOptions testOptions)
     {
         if (testOptions.HasFilterMode && !ModulePathExists())
@@ -146,7 +149,7 @@ private async Task WaitConnectionAsync(CancellationToken token)
         {
             while (!token.IsCancellationRequested)
             {
-                NamedPipeServer pipeConnection = new(_pipeNameDescription, OnRequest, NamedPipeServerStream.MaxAllowedServerInstances, token, skipUnknownMessages: true);
+                var pipeConnection = new NamedPipeServer(_pipeNameDescription, OnRequest, NamedPipeServerStream.MaxAllowedServerInstances, token, skipUnknownMessages: true);
                 pipeConnection.RegisterAllSerializers();
 
                 await pipeConnection.WaitConnectionAsync(token);
@@ -173,20 +176,16 @@ private async Task WaitConnectionAsync(CancellationToken token)
         }
     }
 
-    private Task<IResponse> OnRequest(IRequest request)
+    private Task<IResponse> OnRequest(NamedPipeServer server, IRequest request)
     {
         try
         {
             switch (request)
             {
                 case HandshakeMessage handshakeMessage:
-                    if (handshakeMessage.Properties.TryGetValue(HandshakeMessagePropertyNames.ModulePath, out string value))
-                    {
-                        OnHandshakeMessage(handshakeMessage);
-
-                        return Task.FromResult((IResponse)CreateHandshakeMessage(GetSupportedProtocolVersion(handshakeMessage)));
-                    }
-                    break;
+                    _handshakes.Add(server, handshakeMessage);
+                    OnHandshakeMessage(handshakeMessage);
+                    return Task.FromResult((IResponse)CreateHandshakeMessage(GetSupportedProtocolVersion(handshakeMessage)));
 
                 case CommandLineOptionMessages commandLineOptionMessages:
                     OnCommandLineOptionMessages(commandLineOptionMessages);
@@ -387,7 +386,35 @@ public void Dispose()
     {
         foreach (var namedPipeServer in _testAppPipeConnections)
         {
-            namedPipeServer.Dispose();
+            try
+            {
+                namedPipeServer.Dispose();
+            }
+            catch (Exception ex)
+            {
+                StringBuilder messageBuilder;
+                if (_handshakes.TryGetValue(namedPipeServer, out var handshake))
+                {
+                    messageBuilder = new StringBuilder(CliCommandStrings.DotnetTestPipeFailureHasHandshake);
+                    messageBuilder.AppendLine();
+                    foreach (var kvp in handshake.Properties)
+                    {
+                        messageBuilder.AppendLine($"{kvp.Key}: {kvp.Value}");
+                    }                    
+                }
+                else
+                {
+                    messageBuilder = new StringBuilder(CliCommandStrings.DotnetTestPipeFailureWithoutHandshake);
+                    messageBuilder.AppendLine();
+                }
+
+                messageBuilder.AppendLine($"RunCommand: {Module.RunProperties.Command}");
+                messageBuilder.AppendLine($"RunArguments: {Module.RunProperties.Arguments}");
+                messageBuilder.AppendLine(ex.ToString());
+
+                HasFailureDuringDispose = true;
+                Reporter.Error.WriteLine(messageBuilder.ToString());
+            }
         }
 
         WaitOnTestApplicationPipeConnectionLoop();

src/Cli/dotnet/Commands/Test/TestApplicationActionQueue.cs

@@ -56,9 +56,18 @@ private async Task Read(Func<TestApplication, Task<int>> action, BuildOptions bu
         {
             foreach (var module in nonParallelizedGroup)
             {
+                int result = ExitCode.GenericFailure;
                 var testApp = new TestApplication(module, buildOptions);
-                int result = await action(testApp);
-                testApp.Dispose();
+                using (testApp)
+                {
+                    result = await action(testApp);
+                }
+
+                if (result == ExitCode.Success && testApp.HasFailureDuringDispose)
+                {
+                    result = ExitCode.GenericFailure;
+                }
+                
                 lock (_lock)
                 {
                     if (_aggregateExitCode is null)