Cancel/abort not honored when sent before custom host launch (#1543)

* Cancellation in custom host launcher, test host, while ongoing communication

* passing cancellation to data collectors.

* Passing operation canceled instead of test platform exception to translation layer.

samples/Microsoft.TestPlatform.TranslationLayer.E2ETest/Program.cs

@@ -168,7 +168,7 @@ public CustomTestHostLauncher(Action callback)
 
         public bool IsDebug => false;
 
-        public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo)
+        public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo, CancellationToken cancellationToken)
         {
             var processInfo = new ProcessStartInfo(
                                   defaultTestHostStartInfo.FileName,
@@ -192,6 +192,11 @@ public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo)
 
             throw new Exception("Process in invalid state.");
         }
+
+        public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo)
+        {
+            return this.LaunchTestHost(defaultTestHostStartInfo, CancellationToken.None);
+        }
     }
 
     public class DiscoveryEventHandler : ITestDiscoveryEventsHandler

src/Microsoft.TestPlatform.Client/DesignMode/DesignModeClient.cs

@@ -235,10 +235,13 @@ private void ProcessRequests(ITestRequestManager testRequestManager)
         /// <param name="testProcessStartInfo">
         /// The test Process Start Info.
         /// </param>
+        /// <param name="cancellationToken">
+        /// The cancellation token.
+        /// </param>
         /// <returns>
         /// The <see cref="int"/>.
         /// </returns>
-        public int LaunchCustomHost(TestProcessStartInfo testProcessStartInfo)
+        public int LaunchCustomHost(TestProcessStartInfo testProcessStartInfo, CancellationToken cancellationToken)
         {
             lock (ackLockObject)
             {
@@ -257,7 +260,10 @@ public int LaunchCustomHost(TestProcessStartInfo testProcessStartInfo)
                 // Even if TP has a timeout here, there is no way TP can abort or stop the thread/task that is hung in IDE or LUT
                 // Even if TP can abort the API somehow, TP is essentially putting IDEs or Clients in inconsistent state without having info on
                 // Since the IDEs own user-UI-experience here, TP will let the custom host launch as much time as IDEs define it for their users
-                waitHandle.WaitOne();
+                WaitHandle.WaitAny(new WaitHandle[] { waitHandle, cancellationToken.WaitHandle });
+
+                cancellationToken.ThrowIfCancellationRequested();
+
                 this.onAckMessageReceived = null;
 
                 var ackPayload = this.dataSerializer.DeserializePayload<CustomHostLaunchAckPayload>(ackMessage);

src/Microsoft.TestPlatform.Client/DesignMode/DesignModeTestHostLauncher.cs

@@ -3,6 +3,7 @@
 
 namespace Microsoft.VisualStudio.TestPlatform.Client.DesignMode
 {
+    using System.Threading;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.Interfaces;
 
@@ -28,7 +29,13 @@ public DesignModeTestHostLauncher(IDesignModeClient designModeClient)
         /// <inheritdoc/>
         public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo)
         {
-            return this.designModeClient.LaunchCustomHost(defaultTestHostStartInfo);
+            return this.designModeClient.LaunchCustomHost(defaultTestHostStartInfo, CancellationToken.None);
+        }
+
+        /// <inheritdoc/>
+        public int LaunchTestHost(TestProcessStartInfo defaultTestHostStartInfo, CancellationToken cancellationToken)
+        {
+            return this.designModeClient.LaunchCustomHost(defaultTestHostStartInfo, cancellationToken);
         }
     }
 

src/Microsoft.TestPlatform.Client/DesignMode/IDesignModeClient.cs

@@ -3,9 +3,10 @@
 
 namespace Microsoft.VisualStudio.TestPlatform.Client.DesignMode
 {
+    using System;
+    using System.Threading;
     using Microsoft.VisualStudio.TestPlatform.Client.RequestHelper;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel;
-    using System;
 
     /// <summary>
     /// The interface for design mode client.
@@ -28,7 +29,9 @@ public interface IDesignModeClient : IDisposable
         /// Send a custom host launch message to IDE
         /// </summary>
         /// <param name="defaultTestHostStartInfo">Default TestHost Start Info</param>
-        int LaunchCustomHost(TestProcessStartInfo defaultTestHostStartInfo);
+        /// <param name="cancellationToken">The cancellation Token.</param>
+        /// <returns>Process id of the launched test host.</returns>
+        int LaunchCustomHost(TestProcessStartInfo defaultTestHostStartInfo, CancellationToken cancellationToken);
 
         /// <summary>
         /// Handles parent process exit

src/Microsoft.TestPlatform.Client/Execution/TestRunRequest.cs

@@ -41,6 +41,11 @@ public class TestRunRequest : ITestRunRequest, ITestRunEventsHandler
         /// </summary>
         private object syncObject = new Object();
 
+        /// <summary>
+        /// Sync object for cancel operation
+        /// </summary>
+        private object cancelSyncObject = new Object();
+
         /// <summary>
         /// The run completion event which will be signalled on completion of test run.
         /// </summary>
@@ -236,7 +241,7 @@ public void CancelAsync()
         {
             EqtTrace.Verbose("TestRunRequest.CancelAsync: Canceling.");
 
-            lock (this.syncObject)
+            lock (this.cancelSyncObject)
             {
                 if (this.disposed)
                 {
@@ -265,7 +270,7 @@ public void Abort()
         {
             EqtTrace.Verbose("TestRunRequest.Abort: Aborting.");
 
-            lock (this.syncObject)
+            lock (this.cancelSyncObject)
             {
                 if (this.disposed)
                 {

src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionRequestSender.cs

@@ -160,7 +160,7 @@ public Collection<AttachmentSet> SendAfterTestRunStartAndGetResult(ITestMessageE
 
             // Cycle through the messages that the datacollector sends.
             // Currently each of the operations are not separate tasks since they should not each take much time. This is just a notification.
-            while (!isDataCollectionComplete)
+            while (!isDataCollectionComplete && !isCancelled)
             {
                 var message = this.communicationManager.ReceiveMessage();
 

src/Microsoft.TestPlatform.CommunicationUtilities/Interfaces/ITestRequestSender.cs

@@ -5,6 +5,7 @@ namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
 
     using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Client;
@@ -29,8 +30,9 @@ public interface ITestRequestSender : IDisposable
         /// Waits for Request Handler to be connected
         /// </summary>
         /// <param name="connectionTimeout">Time to wait for connection</param>
+        /// <param name="cancellationToken">Cancellation token</param>
         /// <returns>True, if Handler is connected</returns>
-        bool WaitForRequestHandlerConnection(int connectionTimeout);
+        bool WaitForRequestHandlerConnection(int connectionTimeout, CancellationToken cancellationToken);
 
         /// <summary>
         /// Close the Sender

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -135,14 +135,17 @@ public int InitializeCommunication()
         }
 
         /// <inheritdoc />
-        public bool WaitForRequestHandlerConnection(int connectionTimeout)
+        public bool WaitForRequestHandlerConnection(int connectionTimeout, CancellationToken cancellationToken)
         {
             if (EqtTrace.IsVerboseEnabled)
             {
                 EqtTrace.Verbose("TestRequestSender.WaitForRequestHandlerConnection: waiting for connection with timeout: {0}", connectionTimeout);
             }
 
-            return this.connected.Wait(connectionTimeout);
+            var waitIndex = WaitHandle.WaitAny(new WaitHandle[] { this.connected.WaitHandle, cancellationToken.WaitHandle }, connectionTimeout);
+
+            // Return true if wait is because of waitHandle.
+            return waitIndex == 0;
         }
 
         /// <inheritdoc />
@@ -560,8 +563,10 @@ private string GetAbortErrorMessage(Exception exception, bool getClientError)
                     EqtTrace.Info("TestRequestSender: GetAbortErrorMessage: Received test host error message.");
                     reason = this.clientExitErrorMessage;
                 }
-
-                EqtTrace.Info("TestRequestSender: GetAbortErrorMessage: Timed out waiting for test host error message.");
+                else
+                {
+                    EqtTrace.Info("TestRequestSender: GetAbortErrorMessage: Timed out waiting for test host error message.");
+                }
             }
 
             return reason;

src/Microsoft.TestPlatform.CrossPlatEngine/Client/ProxyDiscoveryManager.cs

@@ -27,7 +27,6 @@ public class ProxyDiscoveryManager : ProxyOperationManager, IProxyDiscoveryManag
     {
         private readonly ITestRuntimeProvider testHostManager;
         private IDataSerializer dataSerializer;
-        private CancellationTokenSource cancellationTokenSource;
         private bool isCommunicationEstablished;
         private IRequestData requestData;
         private ITestDiscoveryEventsHandler2 baseTestDiscoveryEventsHandler;
@@ -67,7 +66,6 @@ internal ProxyDiscoveryManager(IRequestData requestData,
         {
             this.dataSerializer = dataSerializer;
             this.testHostManager = testHostManager;
-            this.cancellationTokenSource = new CancellationTokenSource();
             this.isCommunicationEstablished = false;
             this.requestData = requestData;
         }
@@ -95,7 +93,7 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
             this.baseTestDiscoveryEventsHandler = eventHandler;
             try
             {
-                this.isCommunicationEstablished = this.SetupChannel(discoveryCriteria.Sources, this.cancellationTokenSource.Token);
+                this.isCommunicationEstablished = this.SetupChannel(discoveryCriteria.Sources);
 
                 if (this.isCommunicationEstablished)
                 {

src/Microsoft.TestPlatform.CrossPlatEngine/Client/ProxyExecutionManager.cs

@@ -29,7 +29,6 @@ internal class ProxyExecutionManager : ProxyOperationManager, IProxyExecutionMan
     {
         private readonly ITestRuntimeProvider testHostManager;
         private IDataSerializer dataSerializer;
-        private CancellationTokenSource cancellationTokenSource;
         private bool isCommunicationEstablished;
         private IRequestData requestData;
         private ITestRunEventsHandler baseTestRunEventsHandler;
@@ -65,7 +64,6 @@ internal ProxyExecutionManager(IRequestData requestData, ITestRequestSender requ
         {
             this.testHostManager = testHostManager;
             this.dataSerializer = dataSerializer;
-            this.cancellationTokenSource = new CancellationTokenSource();
             this.isCommunicationEstablished = false;
             this.requestData = requestData;
         }
@@ -100,16 +98,15 @@ public virtual int StartTestRun(TestRunCriteria testRunCriteria, ITestRunEventsH
                 {
                     EqtTrace.Verbose("ProxyExecutionManager: Test host is always Lazy initialize.");
                 }
-
                 var testPackages = new List<string>(testRunCriteria.HasSpecificSources ? testRunCriteria.Sources :
                                                     // If the test execution is with a test filter, group them by sources
                                                     testRunCriteria.Tests.GroupBy(tc => tc.Source).Select(g => g.Key));
 
-                this.isCommunicationEstablished = this.SetupChannel(testPackages, this.cancellationTokenSource.Token);
+                this.isCommunicationEstablished = this.SetupChannel(testPackages);
 
                 if (this.isCommunicationEstablished)
                 {
-                    if (this.cancellationTokenSource.IsCancellationRequested)
+                    if (this.CancellationTokenSource.IsCancellationRequested)
                     {
                         if (EqtTrace.IsVerboseEnabled)
                         {
@@ -139,13 +136,11 @@ public virtual int StartTestRun(TestRunCriteria testRunCriteria, ITestRunEventsH
                     if (testRunCriteria.HasSpecificSources)
                     {
                         var runRequest = testRunCriteria.CreateTestRunCriteriaForSources(testHostManager, runsettings, executionContext, testPackages);
-
                         this.RequestSender.StartTestRun(runRequest, this);
                     }
                     else
                     {
                         var runRequest = testRunCriteria.CreateTestRunCriteriaForTests(testHostManager, runsettings, executionContext, testPackages);
-
                         this.RequestSender.StartTestRun(runRequest, this);
                     }
                 }
@@ -185,7 +180,7 @@ public virtual void Cancel(ITestRunEventsHandler eventHandler)
             }
 
             // Cancel fast, try to stop testhost deployment/launch
-            this.cancellationTokenSource.Cancel();
+            this.CancellationTokenSource.Cancel();
             if (this.isCommunicationEstablished)
             {
                 this.RequestSender.SendTestRunCancel();
@@ -210,7 +205,13 @@ public void Abort(ITestRunEventsHandler eventHandler)
                 this.baseTestRunEventsHandler = eventHandler;
             }
 
-            this.RequestSender.SendTestRunAbort();
+            // Cancel fast, try to stop testhost deployment/launch
+            this.CancellationTokenSource.Cancel();
+
+            if (this.isCommunicationEstablished)
+            {
+                this.RequestSender.SendTestRunAbort();
+            }
         }
 
         /// <inheritdoc/>

src/Microsoft.TestPlatform.CrossPlatEngine/Client/ProxyExecutionManagerWithDataCollection.cs

@@ -112,7 +112,7 @@ public override int StartTestRun(TestRunCriteria testRunCriteria, ITestRunEvents
             var currentEventHandler = eventHandler;
             if (this.ProxyDataCollectionManager != null)
             {
-                currentEventHandler = new DataCollectionTestRunEventsHandler(eventHandler, this.ProxyDataCollectionManager);
+                currentEventHandler = new DataCollectionTestRunEventsHandler(eventHandler, this.ProxyDataCollectionManager, this.CancellationTokenSource.Token);
             }
 
             // Log all the messages that are reported while initializing DataCollectionClient
@@ -129,19 +129,6 @@ public override int StartTestRun(TestRunCriteria testRunCriteria, ITestRunEvents
             return base.StartTestRun(testRunCriteria, currentEventHandler);
         }
 
-        /// <inheritdoc/>
-        public override void Cancel(ITestRunEventsHandler eventHandler)
-        {
-            try
-            {
-                this.ProxyDataCollectionManager.AfterTestRunEnd(isCanceled: true, runEventsHandler: this.DataCollectionRunEventsHandler);
-            }
-            finally
-            {
-                base.Cancel(eventHandler);
-            }
-        }
-
         public override int LaunchProcessWithDebuggerAttached(TestProcessStartInfo testProcessStartInfo)
         {
             if (this.dataCollectionEnvironmentVariables != null)

src/Microsoft.TestPlatform.CrossPlatEngine/Client/ProxyOperationManager.cs

@@ -56,6 +56,7 @@ public abstract class ProxyOperationManager
         protected ProxyOperationManager(IRequestData requestData, ITestRequestSender requestSender, ITestRuntimeProvider testHostManager)
         {
             this.RequestSender = requestSender;
+            this.CancellationTokenSource = new CancellationTokenSource();
             this.testHostManager = testHostManager;
             this.processHelper = new ProcessHelper();
             this.initialized = false;
@@ -75,6 +76,11 @@ protected ProxyOperationManager(IRequestData requestData, ITestRequestSender req
         /// </summary>
         protected ITestRequestSender RequestSender { get; set; }
 
+        /// <summary>
+        /// Gets or sets the cancellation token source.
+        /// </summary>
+        protected CancellationTokenSource CancellationTokenSource { get; set; }
+
         #endregion
 
         #region IProxyOperationManager implementation.
@@ -91,8 +97,9 @@ protected ProxyOperationManager(IRequestData requestData, ITestRequestSender req
         /// <returns>
         /// Returns true if Communation is established b/w runner and host
         /// </returns>
-        public virtual bool SetupChannel(IEnumerable<string> sources, CancellationToken cancellationToken)
+        public virtual bool SetupChannel(IEnumerable<string> sources)
         {
+            this.CancellationTokenSource.Token.ThrowIfCancellationRequested();
             var connTimeout = EnvironmentHelper.GetConnectionTimeout();
 
             if (!this.initialized)
@@ -108,7 +115,6 @@ public virtual bool SetupChannel(IEnumerable<string> sources, CancellationToken
                 }
 
                 var processId = this.processHelper.GetCurrentProcessId();
-
                 var connectionInfo = new TestRunnerConnectionInfo { Port = portNumber, ConnectionInfo = testHostConnectionInfo, RunnerProcessId = processId, LogFile = this.GetTimestampedLogFile(EqtTrace.LogFile) };
 
                 // Subscribe to TestHost Event
@@ -120,7 +126,7 @@ public virtual bool SetupChannel(IEnumerable<string> sources, CancellationToken
                 try
                 {
                     // Launch the test host.
-                    var hostLaunchedTask = this.testHostManager.LaunchTestHostAsync(testHostStartInfo, cancellationToken);
+                    var hostLaunchedTask = this.testHostManager.LaunchTestHostAsync(testHostStartInfo, this.CancellationTokenSource.Token);
                     this.testHostLaunched = hostLaunchedTask.Result;
 
                     if (this.testHostLaunched && testHostConnectionInfo.Role == ConnectionRole.Host)
@@ -150,7 +156,7 @@ public virtual bool SetupChannel(IEnumerable<string> sources, CancellationToken
 
                 // Wait for a timeout for the client to connect.
                 if (!this.testHostLaunched ||
-                    !this.RequestSender.WaitForRequestHandlerConnection(connTimeout * 1000))
+                    !this.RequestSender.WaitForRequestHandlerConnection(connTimeout * 1000, this.CancellationTokenSource.Token))
                 {
                     this.ThrowExceptionOnConnectionFailure(connTimeout);
                 }