Ensure correct exit code in case of cancellation and add OnExit phase for for IPushOnlyProtocol (#3820)

Co-authored-by: Marco Rossignoli <mrossignol@microsoft.com>

Author: MarcoRossignoli

src/Platform/Microsoft.Testing.Platform/Hosts/CommonTestHost.cs

@@ -26,28 +26,43 @@ public async Task<int> RunAsync()
     {
         CancellationToken testApplicationCancellationToken = ServiceProvider.GetTestApplicationCancellationTokenSource().CancellationToken;
 
-        int exitCode;
+        int exitCode = ExitCodes.GenericFailure;
         try
         {
             if (PushOnlyProtocol is null || PushOnlyProtocol?.IsServerMode == false)
             {
                 exitCode = await RunTestAppAsync(testApplicationCancellationToken);
+
+                if (testApplicationCancellationToken.IsCancellationRequested)
+                {
+                    exitCode = ExitCodes.TestSessionAborted;
+                }
+
                 return exitCode;
             }
 
-            RoslynDebug.Assert(PushOnlyProtocol is not null);
+            try
+            {
+                RoslynDebug.Assert(PushOnlyProtocol is not null);
 
-            ITestApplicationModuleInfo testApplicationModuleInfo = serviceProvider.GetTestApplicationModuleInfo();
-            bool isValidProtocol = await PushOnlyProtocol.IsCompatibleProtocolAsync(GetHostType());
+                ITestApplicationModuleInfo testApplicationModuleInfo = serviceProvider.GetTestApplicationModuleInfo();
+                bool isValidProtocol = await PushOnlyProtocol.IsCompatibleProtocolAsync(GetHostType());
 
-            exitCode = isValidProtocol
-                ? await RunTestAppAsync(testApplicationCancellationToken)
-                : ExitCodes.IncompatibleProtocolVersion;
+                exitCode = isValidProtocol
+                    ? await RunTestAppAsync(testApplicationCancellationToken)
+                    : ExitCodes.IncompatibleProtocolVersion;
+            }
+            finally
+            {
+                if (PushOnlyProtocol is not null)
+                {
+                    await PushOnlyProtocol.OnExitAsync();
+                }
+            }
         }
         catch (OperationCanceledException) when (testApplicationCancellationToken.IsCancellationRequested)
         {
             // We do nothing we're canceling
-            exitCode = ExitCodes.TestSessionAborted;
         }
         finally
         {
@@ -57,6 +72,11 @@ public async Task<int> RunAsync()
             await DisposeHelper.DisposeAsync(ServiceProvider.GetTestApplicationCancellationTokenSource());
         }
 
+        if (testApplicationCancellationToken.IsCancellationRequested)
+        {
+            exitCode = ExitCodes.TestSessionAborted;
+        }
+
         return exitCode;
     }
 
@@ -217,7 +237,8 @@ protected static async Task DisposeServiceProviderAsync(ServiceProvider serviceP
             // We need to ensure that we won't dispose special services till the shutdown
             if (!isProcessShutdown &&
                 service is ITelemetryCollector or
-                 ITestApplicationLifecycleCallbacks)
+                 ITestApplicationLifecycleCallbacks or
+                 IPushOnlyProtocol)
             {
                 continue;
             }

src/Platform/Microsoft.Testing.Platform/ServerMode/DotnetTest/DotnetTestConnection.cs

@@ -142,6 +142,8 @@ public async Task SendMessageAsync(IRequest message)
         }
     }
 
+    public Task OnExitAsync() => Task.CompletedTask;
+
     public void Dispose() => _dotnetTestPipeClient?.Dispose();
 
 #if NETCOREAPP

src/Platform/Microsoft.Testing.Platform/ServerMode/IPushOnlyProtocol.cs

@@ -20,4 +20,6 @@ internal interface IPushOnlyProtocol :
     Task<bool> IsCompatibleProtocolAsync(string testHostType);
 
     Task<IPushOnlyProtocolConsumer> GetDataConsumerAsync();
+
+    Task OnExitAsync();
 }