Cancel/abort not honored when sent before custom host launch

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)
             {
@@ -250,6 +253,9 @@ public int LaunchCustomHost(TestProcessStartInfo testProcessStartInfo)
                     waitHandle.Set();
                 };
 
+                // Registering cancellationToken to set waitHandle (whenever request is cancelled).
+                var cancellationTokenRegistration = cancellationToken.Register(() => waitHandle.Set());
+
                 this.communicationManager.SendMessage(MessageType.CustomTestHostLaunch, testProcessStartInfo);
 
                 // LifeCycle of the TP through DesignModeClient is maintained by the IDEs or user-facing-clients like LUTs, who call TestPlatform
@@ -258,6 +264,10 @@ public int LaunchCustomHost(TestProcessStartInfo testProcessStartInfo)
                 // 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();
+
+                cancellationTokenRegistration.Dispose();
+                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,16 @@ public class TestRunRequest : ITestRunRequest, ITestRunEventsHandler
         /// </summary>
         private object syncObject = new Object();
 
+        /// <summary>
+        /// Sync object for cancel operation
+        /// </summary>
+        private object cancelSyncObject = new Object();
+
+        /// <summary>
+        /// Sync object for abort operation
+        /// </summary>
+        private object abortSyncObject = new Object();
+
         /// <summary>
         /// The run completion event which will be signalled on completion of test run.
         /// </summary>
@@ -236,7 +246,7 @@ public void CancelAsync()
         {
             EqtTrace.Verbose("TestRunRequest.CancelAsync: Canceling.");
 
-            lock (this.syncObject)
+            lock (this.cancelSyncObject)
             {
                 if (this.disposed)
                 {
@@ -265,7 +275,7 @@ public void Abort()
         {
             EqtTrace.Verbose("TestRunRequest.Abort: Aborting.");
 
-            lock (this.syncObject)
+            lock (this.abortSyncObject)
             {
                 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
@@ -61,14 +63,16 @@ public interface ITestRequestSender : IDisposable
         /// </summary>
         /// <param name="runCriteria">RunCriteria for test run</param>
         /// <param name="eventHandler">EventHandler for test run events</param>
-        void StartTestRun(TestRunCriteriaWithSources runCriteria, ITestRunEventsHandler eventHandler);
+        /// <param name="cancellationToken">Cancellation token</param>
+        void StartTestRun(TestRunCriteriaWithSources runCriteria, ITestRunEventsHandler eventHandler, CancellationToken cancellationToken);
 
         /// <summary>
         /// Starts the TestRun with given test cases and criteria
         /// </summary>
         /// <param name="runCriteria">RunCriteria for test run</param>
         /// <param name="eventHandler">EventHandler for test run events</param>
-        void StartTestRun(TestRunCriteriaWithTests runCriteria, ITestRunEventsHandler eventHandler);
+        /// <param name="cancellationToken">Cancellation token</param>
+        void StartTestRun(TestRunCriteriaWithTests runCriteria, ITestRunEventsHandler eventHandler, CancellationToken cancellationToken);
 
         /// <summary>
         /// Ends the Session

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -135,14 +135,18 @@ public int InitializeCommunication()
         }
 
         /// <inheritdoc />
-        public bool WaitForRequestHandlerConnection(int connectionTimeout)
+        public bool WaitForRequestHandlerConnection(int connectionTimeout, CancellationToken cancellationToken)
         {
+            var cancellationTokenRegistration = cancellationToken.Register(() => this.connected.Set());
             if (EqtTrace.IsVerboseEnabled)
             {
                 EqtTrace.Verbose("TestRequestSender.WaitForRequestHandlerConnection: waiting for connection with timeout: {0}", connectionTimeout);
             }
 
-            return this.connected.Wait(connectionTimeout);
+            var waitSuccess = this.connected.Wait(connectionTimeout) && !cancellationToken.IsCancellationRequested;
+            cancellationTokenRegistration.Dispose();
+
+            return waitSuccess;
         }
 
         /// <inheritdoc />
@@ -259,14 +263,14 @@ public void InitializeExecution(IEnumerable<string> pathToAdditionalExtensions)
         }
 
         /// <inheritdoc />
-        public void StartTestRun(TestRunCriteriaWithSources runCriteria, ITestRunEventsHandler eventHandler)
+        public void StartTestRun(TestRunCriteriaWithSources runCriteria, ITestRunEventsHandler eventHandler, CancellationToken cancellationToken)
         {
             this.messageEventHandler = eventHandler;
             this.onDisconnected = (disconnectedEventArgs) =>
                 {
-                    this.OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true);
+                    this.OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true, cancellationToken);
                 };
-            this.onMessageReceived = (sender, args) => this.OnExecutionMessageReceived(sender, args, eventHandler);
+            this.onMessageReceived = (sender, args) => this.OnExecutionMessageReceived(sender, args, eventHandler, cancellationToken);
             this.channel.MessageReceived += this.onMessageReceived;
 
             var message = this.dataSerializer.SerializePayload(
@@ -283,14 +287,14 @@ public void StartTestRun(TestRunCriteriaWithSources runCriteria, ITestRunEventsH
         }
 
         /// <inheritdoc />
-        public void StartTestRun(TestRunCriteriaWithTests runCriteria, ITestRunEventsHandler eventHandler)
+        public void StartTestRun(TestRunCriteriaWithTests runCriteria, ITestRunEventsHandler eventHandler, CancellationToken cancellationToken)
         {
             this.messageEventHandler = eventHandler;
             this.onDisconnected = (disconnectedEventArgs) =>
                 {
-                    this.OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true);
+                    this.OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true, cancellationToken);
                 };
-            this.onMessageReceived = (sender, args) => this.OnExecutionMessageReceived(sender, args, eventHandler);
+            this.onMessageReceived = (sender, args) => this.OnExecutionMessageReceived(sender, args, eventHandler, cancellationToken);
             this.channel.MessageReceived += this.onMessageReceived;
 
             var message = this.dataSerializer.SerializePayload(
@@ -379,7 +383,7 @@ public void Dispose()
             this.communicationEndpoint.Stop();
         }
 
-        private void OnExecutionMessageReceived(object sender, MessageReceivedEventArgs messageReceived, ITestRunEventsHandler testRunEventsHandler)
+        private void OnExecutionMessageReceived(object sender, MessageReceivedEventArgs messageReceived, ITestRunEventsHandler testRunEventsHandler, CancellationToken cancellationToken)
         {
             try
             {
@@ -434,7 +438,7 @@ private void OnExecutionMessageReceived(object sender, MessageReceivedEventArgs
             }
             catch (Exception exception)
             {
-                this.OnTestRunAbort(testRunEventsHandler, exception, false);
+                this.OnTestRunAbort(testRunEventsHandler, exception, false, cancellationToken);
             }
         }
 
@@ -485,7 +489,7 @@ private void OnDiscoveryMessageReceived(ITestDiscoveryEventsHandler2 discoveryEv
             }
         }
 
-        private void OnTestRunAbort(ITestRunEventsHandler testRunEventsHandler, Exception exception, bool getClientError)
+        private void OnTestRunAbort(ITestRunEventsHandler testRunEventsHandler, Exception exception, bool getClientError, CancellationToken cancellationToken)
         {
             if (this.IsOperationComplete())
             {
@@ -496,7 +500,7 @@ private void OnTestRunAbort(ITestRunEventsHandler testRunEventsHandler, Exceptio
             EqtTrace.Verbose("TestRequestSender: OnTestRunAbort: Set operation complete.");
             this.SetOperationComplete();
 
-            var reason = this.GetAbortErrorMessage(exception, getClientError);
+            var reason = this.GetAbortErrorMessage(exception, getClientError, cancellationToken);
             EqtTrace.Error("TestRequestSender: Aborting test run because {0}", reason);
             this.LogErrorMessage(string.Format(CommonResources.AbortedTestRun, reason));
 
@@ -522,7 +526,7 @@ private void OnDiscoveryAbort(ITestDiscoveryEventsHandler2 eventHandler, Excepti
             this.SetOperationComplete();
 
             var discoveryCompleteEventArgs = new DiscoveryCompleteEventArgs(-1, true);
-            var reason = this.GetAbortErrorMessage(exception, getClientError);
+            var reason = this.GetAbortErrorMessage(exception, getClientError, CancellationToken.None);
             EqtTrace.Error("TestRequestSender: Aborting test discovery because {0}", reason);
             this.LogErrorMessage(string.Format(CommonResources.AbortedTestDiscovery, reason));
 
@@ -540,8 +544,9 @@ private void OnDiscoveryAbort(ITestDiscoveryEventsHandler2 eventHandler, Excepti
             eventHandler.HandleDiscoveryComplete(discoveryCompleteEventArgs, null);
         }
 
-        private string GetAbortErrorMessage(Exception exception, bool getClientError)
+        private string GetAbortErrorMessage(Exception exception, bool getClientError, CancellationToken cancellationToken)
         {
+            var cancellationTokenRegistration = cancellationToken.Register(() => this.clientExited.Set());
             EqtTrace.Verbose("TestRequestSender: GetAbortErrorMessage: Exception: " + exception);
 
             // It is also possible for an operation to abort even if client has not
@@ -554,15 +559,18 @@ private string GetAbortErrorMessage(Exception exception, bool getClientError)
 
                 // Set a default message and wait for test host to exit for a moment
                 reason = CommonResources.UnableToCommunicateToTestHost;
-                if (this.clientExited.Wait(this.clientExitedWaitTime))
+                if (this.clientExited.Wait(this.clientExitedWaitTime) && !cancellationToken.IsCancellationRequested)
                 {
                     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.");
+                }
             }
 
+            cancellationTokenRegistration.Dispose();
             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;
@@ -71,7 +70,6 @@ internal ProxyDiscoveryManager(
         {
             this.dataSerializer = dataSerializer;
             this.testHostManager = testHostManager;
-            this.cancellationTokenSource = new CancellationTokenSource();
             this.isCommunicationEstablished = false;
             this.requestData = requestData;
         }
@@ -99,7 +97,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)
                 {