Add exit call back to custom test host process

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -313,8 +313,8 @@ public void OnClientProcessExit(string stdError)
             this.clientExitErrorMessage = stdError;
             this.clientExited.Set();
 
-            // Note that we're not explicitly disconnecting the communication channel; wait for the
-            // channel to determine the disconnection on its own.
+            // Break communication loop. In somecases(E.g: When tests creates child processes to testhost) communication channel won't break if testhost exits.
+            this.communicationEndpoint.Stop();
         }
 
         /// <inheritdoc />

src/Microsoft.TestPlatform.PlatformAbstractions/Interfaces/System/IProcessHelper.cs

@@ -90,7 +90,7 @@ public interface IProcessHelper
         /// <param name="callbackAction">
         /// Callback on process exit.
         /// </param>
-        void SetExitCallback(int processId, Action callbackAction);
+        void SetExitCallback(int processId, Action<object> callbackAction);
 
         /// <summary>
         /// Terminates a process.

src/Microsoft.TestPlatform.PlatformAbstractions/common/System/ProcessHelper.cs

@@ -51,10 +51,10 @@ public object LaunchProcess(string processPath, string arguments, string working
                     process.Exited += (sender, args) =>
                     {
                         // Call WaitForExit without again to ensure all streams are flushed,
-                        // Add timeout to avoid indefinite waiting on child process exist.
                         var p = sender as Process;
                         try
                         {
+                            // Add timeout to avoid indefinite waiting on child process exit.
                             p.WaitForExit(500);
                         }
                         catch (InvalidOperationException)
@@ -131,15 +131,20 @@ public bool TryGetExitCode(object process, out int exitCode)
         }
 
         /// <inheritdoc/>
-        public void SetExitCallback(int processId, Action callbackAction)
+        public void SetExitCallback(int processId, Action<object> callbackAction)
         {
-            var process = Process.GetProcessById(processId);
-
-            process.EnableRaisingEvents = true;
-            process.Exited += (sender, args) =>
+            try
+            {
+                var process = Process.GetProcessById(processId);
+                process.EnableRaisingEvents = true;
+                process.Exited += (sender, args) => callbackAction?.Invoke(sender);
+            }
+            catch (ArgumentException)
             {
-                callbackAction.Invoke();
-            };
+                // Process.GetProcessById() throws ArgumentException if process is not running(identifier might be expired).
+                // Invoke callback immediately.
+                callbackAction?.Invoke(null);
+            }
         }
 
         /// <inheritdoc/>

src/Microsoft.TestPlatform.PlatformAbstractions/netstandard1.0/System/ProcessHelper.cs

@@ -62,7 +62,7 @@ public bool TryGetExitCode(object process, out int exitCode)
         }
 
         /// <inheritdoc/>
-        public void SetExitCallback(int parentProcessId, Action callbackAction)
+        public void SetExitCallback(int parentProcessId, Action<object> callbackAction)
         {
             throw new NotImplementedException();
         }

src/Microsoft.TestPlatform.PlatformAbstractions/uap10.0/System/ProcessHelper.cs

@@ -62,7 +62,7 @@ public bool TryGetExitCode(object process, out int exitCode)
         }
 
         /// <inheritdoc/>
-        public void SetExitCallback(int parentProcessId, Action callbackAction)
+        public void SetExitCallback(int parentProcessId, Action<object> callbackAction)
         {
         }
 

src/Microsoft.TestPlatform.TestHostProvider/Hosting/DefaultTestHostManager.cs

@@ -385,6 +385,7 @@ private bool LaunchHost(TestProcessStartInfo testHostStartInfo, CancellationToke
                 {
                     int processId = this.customTestHostLauncher.LaunchTestHost(testHostStartInfo);
                     this.testHostProcess = Process.GetProcessById(processId);
+                    this.processHelper.SetExitCallback(processId, this.ExitCallBack);
                 }
             }
             catch (OperationCanceledException ex)

src/Microsoft.TestPlatform.TestHostProvider/Hosting/DotnetTestHostManager.cs

@@ -340,6 +340,7 @@ private bool LaunchHost(TestProcessStartInfo testHostStartInfo, CancellationToke
                 {
                     var processId = this.testHostLauncher.LaunchTestHost(testHostStartInfo);
                     this.testHostProcess = Process.GetProcessById(processId);
+                    this.processHelper.SetExitCallback(processId, this.ExitCallBack);
                 }
             }
             catch (OperationCanceledException ex)

src/datacollector/Program.cs

@@ -77,7 +77,7 @@ private static void Run(string[] args)
             var processHelper = new ProcessHelper();
             processHelper.SetExitCallback(
                 parentProcessId,
-                () =>
+                (obj) =>
                     {
                         EqtTrace.Info("DataCollector: ParentProcess '{0}' Exited.", parentProcessId);
                         Environment.Exit(1);

src/testhost.x86/DefaultEngineInvoker.cs

@@ -90,7 +90,7 @@ public void Invoke(IDictionary<string, string> argsDictionary)
                     var processHelper = new ProcessHelper();
                     processHelper.SetExitCallback(
                         parentProcessId,
-                        () =>
+                        (obj) =>
                             {
                                 EqtTrace.Info("DefaultEngineInvoker: ParentProcess '{0}' Exited.", parentProcessId);
                                 new PlatformEnvironment().Exit(1);

src/vstest.console/Processors/PortArgumentProcessor.cs

@@ -203,7 +203,7 @@ private static IDesignModeClient InitializeDesignMode(int parentProcessId, IProc
         {
             if (parentProcessId > 0)
             {
-                processHelper.SetExitCallback(parentProcessId, () =>
+                processHelper.SetExitCallback(parentProcessId, (obj) =>
                 {
                     EqtTrace.Info($"PortArgumentProcessor: parent process:{parentProcessId} exited.");
                     DesignModeClient.Instance?.HandleParentProcessExit();

test/Microsoft.TestPlatform.TestHostProvider.UnitTests/Hosting/DefaultTestHostManagerTests.cs

@@ -380,6 +380,18 @@ public void LaunchTestHostShouldUseCustomHostIfSet()
             Assert.AreEqual(currentProcess.Id, this.testHostId);
         }
 
+        [TestMethod]
+        public void LaunchTestHostShouldSetExitCallbackInCaseCustomHost()
+        {
+            var mockCustomLauncher = new Mock<ITestHostLauncher>();
+            this.testHostManager.SetCustomLauncher(mockCustomLauncher.Object);
+            var currentProcess = Process.GetCurrentProcess();
+            mockCustomLauncher.Setup(mc => mc.LaunchTestHost(It.IsAny<TestProcessStartInfo>())).Returns(currentProcess.Id);
+            this.testHostManager.LaunchTestHostAsync(this.startInfo, CancellationToken.None).Wait();
+
+            this.mockProcessHelper.Verify(ph => ph.SetExitCallback(currentProcess.Id, It.IsAny<Action<object>>()));
+        }
+
         [TestMethod]
         public void GetTestSourcesShouldReturnAppropriateSourceIfAppxRecipeIsProvided()
         {

test/Microsoft.TestPlatform.TestHostProvider.UnitTests/Hosting/DotnetTestHostManagerTests.cs

@@ -369,6 +369,20 @@ public async Task LaunchTestHostShouldLaunchProcessWithConnectionInfo()
             this.mockTestHostLauncher.Verify(thl => thl.LaunchTestHost(It.Is<TestProcessStartInfo>(x => x.Arguments.Equals(expectedArgs))), Times.Once);
         }
 
+        [TestMethod]
+        public void LaunchTestHostShouldSetExitCallBackInCaseCustomHost()
+        {
+            var expectedProcessId = Process.GetCurrentProcess().Id;
+            this.mockTestHostLauncher.Setup(thl => thl.LaunchTestHost(It.IsAny<TestProcessStartInfo>())).Returns(expectedProcessId);
+            this.mockFileHelper.Setup(ph => ph.Exists("testhost.dll")).Returns(true);
+
+            var startInfo = this.GetDefaultStartInfo();
+            this.dotnetHostManager.SetCustomLauncher(this.mockTestHostLauncher.Object);
+            this.dotnetHostManager.LaunchTestHostAsync(startInfo, CancellationToken.None).Wait();
+
+            this.mockProcessHelper.Verify(ph => ph.SetExitCallback(expectedProcessId, It.IsAny<Action<object>>()));
+        }
+
         [TestMethod]
         public void GetTestHostProcessStartInfoShouldIncludeTestHostPathFromSourceDirectoryIfDepsFileNotFound()
         {

test/vstest.console.UnitTests/Processors/PortArgumentProcessorTests.cs

@@ -124,7 +124,7 @@ public void ExecutorInitializeShouldSetProcessExitCallback()
 
             this.executor.Initialize(port.ToString());
 
-            this.mockProcessHelper.Verify(ph => ph.SetExitCallback(processId, It.IsAny<Action>()), Times.Once);
+            this.mockProcessHelper.Verify(ph => ph.SetExitCallback(processId, It.IsAny<Action<object>>()), Times.Once);
         }
 
         [TestMethod]