Implemented the cancellation of discovery request

TestPlatform.sln

@@ -1,7 +1,7 @@
 
 Microsoft Visual Studio Solution File, Format Version 12.00
-# Visual Studio 15
-VisualStudioVersion = 15.0.26730.16
+# Visual Studio Version 16
+VisualStudioVersion = 16.0.29025.244
 MinimumVisualStudioVersion = 10.0.40219.1
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "src", "src", "{ED0C35EB-7F31-4841-A24F-8EB708FFA959}"
 EndProject
@@ -170,6 +170,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "SettingsMigrator", "src\Set
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "SettingsMigrator.UnitTests", "test\SettingsMigrator.UnitTests\SettingsMigrator.UnitTests.csproj", "{E7D4921C-F12D-4E1C-85AC-8B7F91C59B0E}"
 EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "DiscoveryTestProject", "test\TestAssets\DiscoveryTestProject\DiscoveryTestProject.csproj", "{D16ACC60-52F8-4912-8870-5733A9F6852D}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -852,6 +854,18 @@ Global
 		{E7D4921C-F12D-4E1C-85AC-8B7F91C59B0E}.Release|x64.Build.0 = Release|Any CPU
 		{E7D4921C-F12D-4E1C-85AC-8B7F91C59B0E}.Release|x86.ActiveCfg = Release|Any CPU
 		{E7D4921C-F12D-4E1C-85AC-8B7F91C59B0E}.Release|x86.Build.0 = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|x64.ActiveCfg = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|x64.Build.0 = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|x86.ActiveCfg = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Debug|x86.Build.0 = Debug|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|Any CPU.Build.0 = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|x64.ActiveCfg = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|x64.Build.0 = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|x86.ActiveCfg = Release|Any CPU
+		{D16ACC60-52F8-4912-8870-5733A9F6852D}.Release|x86.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -925,6 +939,7 @@ Global
 		{A7E2261B-B2E6-4CBF-983F-E3A3FF8E52E3} = {D9A30E32-D466-4EC5-B4F2-62E17562279B}
 		{69F5FF81-5615-4F06-B83C-FCF979BB84CA} = {ED0C35EB-7F31-4841-A24F-8EB708FFA959}
 		{E7D4921C-F12D-4E1C-85AC-8B7F91C59B0E} = {B27FAFDF-2DBA-4AB0-BA85-FD5F21D359D6}
+		{D16ACC60-52F8-4912-8870-5733A9F6852D} = {8DA7CBD9-F17E-41B6-90C4-CFF55848A25A}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {0541B30C-FF51-4E28-B172-83F5F3934BCD}

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

@@ -196,6 +196,12 @@ private void ProcessRequests(ITestRequestManager testRequestManager)
                                 break;
                             }
 
+                        case MessageType.CancelDiscovery:
+                            {
+                                testRequestManager.CancelDiscovery();
+                                break;
+                            }
+
                         case MessageType.CancelTestRun:
                             {
                                 testRequestManager.CancelTestRun();

src/Microsoft.TestPlatform.Client/RequestHelper/ITestRequestManager.cs

@@ -53,5 +53,10 @@ public interface ITestRequestManager : IDisposable
         /// Abort the current TestRun
         /// </summary>
         void AbortTestRun();
+
+        /// <summary>
+        /// Cancels the current discovery request
+        /// </summary>
+        void CancelDiscovery();
     }
 }

src/Microsoft.TestPlatform.CommunicationUtilities/Messages/MessageType.cs

@@ -63,6 +63,11 @@ public static class MessageType
         /// </summary>
         public const string DiscoveryComplete = "TestDiscovery.Completed";
 
+        /// <summary>
+        /// Cancel Test Discovery
+        /// </summary>
+        public const string CancelDiscovery = "TestDiscovery.Cancel";
+
         /// <summary>
         /// The session start.
         /// </summary>

src/Microsoft.TestPlatform.CrossPlatEngine/Discovery/DiscovererEnumerator.cs

@@ -10,7 +10,7 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Discovery
     using System.IO;
     using System.Linq;
     using System.Reflection;
-
+    using System.Threading;
     using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework;
     using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework.Utilities;
     using Microsoft.VisualStudio.TestPlatform.Common.Filtering;
@@ -33,59 +33,80 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Discovery
     /// </summary>
     internal class DiscovererEnumerator
     {
-        private DiscoveryResultCache discoveryResultCache;
-        private ITestPlatformEventSource testPlatformEventSource;
-        private IRequestData requestData;
-        private IAssemblyProperties assemblyProperties;
+        private readonly DiscoveryResultCache discoveryResultCache;
+        private readonly ITestPlatformEventSource testPlatformEventSource;
+        private readonly IRequestData requestData;
+        private readonly IAssemblyProperties assemblyProperties;
+        private readonly CancellationToken cancellationToken;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="DiscovererEnumerator"/> class.
         /// </summary>
         /// <param name="requestData">The request data for providing discovery services and data.</param>
         /// <param name="discoveryResultCache"> The discovery result cache. </param>
-        /// <param name="metricsCollector">Metric Collector</param>
-        public DiscovererEnumerator(IRequestData requestData, DiscoveryResultCache discoveryResultCache) : this(requestData, discoveryResultCache, TestPlatformEventSource.Instance)
+        public DiscovererEnumerator(IRequestData requestData, DiscoveryResultCache discoveryResultCache, CancellationToken token)
+            : this(requestData, discoveryResultCache, TestPlatformEventSource.Instance, token)
         {
         }
 
-        internal DiscovererEnumerator(IRequestData requestData,
+        /// <summary>
+        /// Initializes a new instance of the <see cref="DiscovererEnumerator"/> class.
+        /// </summary>
+        /// <param name="requestData">The request data for providing discovery services and data.</param>
+        /// <param name="discoveryResultCache"> The discovery result cache. </param>
+        /// <param name="testPlatformEventSource">Telemetry events receiver</param>
+        /// <param name="token">Cancellation Token to abort discovery</param>
+        public DiscovererEnumerator(IRequestData requestData,
             DiscoveryResultCache discoveryResultCache,
-            ITestPlatformEventSource testPlatformEventSource)
+            ITestPlatformEventSource testPlatformEventSource,
+            CancellationToken token)
+            : this(requestData, discoveryResultCache, testPlatformEventSource, new AssemblyProperties(), token)
         {
-            this.discoveryResultCache = discoveryResultCache;
-            this.testPlatformEventSource = testPlatformEventSource;
-            this.requestData = requestData;
-            this.assemblyProperties = new AssemblyProperties();
         }
 
-        // Added this to make class testable, needed a PEHeader mocked Instance
-        internal DiscovererEnumerator(IRequestData requestData,
+        /// <summary>
+        /// Initializes a new instance of the <see cref="DiscovererEnumerator"/> class.
+        /// </summary>
+        /// <param name="requestData">The request data for providing discovery services and data.</param>
+        /// <param name="discoveryResultCache"> The discovery result cache. </param>
+        /// <param name="testPlatformEventSource">Telemetry events receiver</param>
+        /// <param name="assemblyProperties">Information on the assemblies being discovered</param>
+        /// <param name="token">Cancellation Token to abort discovery</param>
+        public DiscovererEnumerator(IRequestData requestData,
             DiscoveryResultCache discoveryResultCache,
             ITestPlatformEventSource testPlatformEventSource,
-            IAssemblyProperties assemblyProperties)
+            IAssemblyProperties assemblyProperties,
+            CancellationToken token)
         {
+            // Added this to make class testable, needed a PEHeader mocked Instance
             this.discoveryResultCache = discoveryResultCache;
             this.testPlatformEventSource = testPlatformEventSource;
             this.requestData = requestData;
             this.assemblyProperties = assemblyProperties;
+            this.cancellationToken = token;
         }
 
-
         /// <summary>
         /// Discovers tests from the sources.
         /// </summary>
         /// <param name="testExtensionSourceMap"> The test extension source map. </param>
         /// <param name="settings"> The settings. </param>
         /// <param name="testCaseFilter"> The test case filter. </param>
         /// <param name="logger"> The logger. </param>
-        internal void LoadTests(IDictionary<string, IEnumerable<string>> testExtensionSourceMap, IRunSettings settings, string testCaseFilter, IMessageLogger logger)
+        public void LoadTests(IDictionary<string, IEnumerable<string>> testExtensionSourceMap, IRunSettings settings, string testCaseFilter, IMessageLogger logger)
         {
             this.testPlatformEventSource.DiscoveryStart();
-            foreach (var kvp in testExtensionSourceMap)
+            try
             {
-                this.LoadTestsFromAnExtension(kvp.Key, kvp.Value, settings, testCaseFilter, logger);
+                foreach (var kvp in testExtensionSourceMap)
+                {
+                    this.LoadTestsFromAnExtension(kvp.Key, kvp.Value, settings, testCaseFilter, logger);
+                }
+            }
+            finally
+            {
+                this.testPlatformEventSource.DiscoveryStop(this.discoveryResultCache.TotalDiscoveredTests);
             }
-            this.testPlatformEventSource.DiscoveryStop(this.discoveryResultCache.TotalDiscoveredTests);
         }
 
         /// <summary>
@@ -105,7 +126,7 @@ private void LoadTestsFromAnExtension(string extensionAssembly, IEnumerable<stri
             // Stopwatch to collect metrics
             var timeStart = DateTime.UtcNow;
 
-            var discovererToSourcesMap = DiscovererEnumerator.GetDiscovererToSourcesMap(extensionAssembly, sources, logger, this.assemblyProperties);
+            var discovererToSourcesMap = GetDiscovererToSourcesMap(extensionAssembly, sources, logger, this.assemblyProperties);
             var totalAdapterLoadTIme = DateTime.UtcNow - timeStart;
 
             // Collecting Data Point for TimeTaken to Load Adapters
@@ -117,30 +138,46 @@ private void LoadTestsFromAnExtension(string extensionAssembly, IEnumerable<stri
                 return;
             }
 
-            // Collecting Total Number of Adapters Discovered in Machine
-            this.requestData.MetricsCollection.Add(TelemetryDataConstants.NumberOfAdapterDiscoveredDuringDiscovery, discovererToSourcesMap.Keys.Count());
+            double totalTimeTakenByAdapters = 0;
+            double totalAdaptersUsed = 0;
+            try
+            {
+                // Collecting Total Number of Adapters Discovered in Machine
+                this.requestData.MetricsCollection.Add(TelemetryDataConstants.NumberOfAdapterDiscoveredDuringDiscovery, discovererToSourcesMap.Keys.Count());
 
-            var context = new DiscoveryContext { RunSettings = settings };
-            context.FilterExpressionWrapper = !string.IsNullOrEmpty(testCaseFilter) ? new FilterExpressionWrapper(testCaseFilter) : null;
+                var context = new DiscoveryContext { RunSettings = settings };
+                context.FilterExpressionWrapper = !string.IsNullOrEmpty(testCaseFilter) ? new FilterExpressionWrapper(testCaseFilter) : null;
 
-            // Set on the logger the TreatAdapterErrorAsWarning setting from runsettings.
-            this.SetAdapterLoggingSettings(logger, settings);
+                // Set on the logger the TreatAdapterErrorAsWarning setting from runsettings.
+                this.SetAdapterLoggingSettings(logger, settings);
 
-            var discoverySink = new TestCaseDiscoverySink(this.discoveryResultCache);
-            double totalTimeTakenByAdapters = 0;
-            double totalAdaptersUsed = 0;
+                var discoverySink = new TestCaseDiscoverySink(this.discoveryResultCache);
+                foreach (var discoverer in discovererToSourcesMap.Keys)
+                {
+                    if (this.cancellationToken.IsCancellationRequested)
+                    {
+                        EqtTrace.Info("DiscovererEnumerator.LoadTestsFromAnExtension: Cancellation Requested. Aborting the discovery");
+                        LogTestsDiscoveryCancellation(logger);
+                        return;
+                    }
+
+                    this.DiscoverTestsFromSingleDiscoverer(discoverer, discovererToSourcesMap, context, discoverySink, logger, ref totalAdaptersUsed, ref totalTimeTakenByAdapters);
+                }
 
-            foreach (var discoverer in discovererToSourcesMap.Keys)
-            {
-                this.DiscoverTestsFromSingleDiscoverer(discoverer, discovererToSourcesMap, context, discoverySink, logger, ref totalAdaptersUsed, ref totalTimeTakenByAdapters);
+                if (this.discoveryResultCache.TotalDiscoveredTests == 0)
+                {
+                    LogWarningOnNoTestsDiscovered(sources, testCaseFilter, logger);
+                }
             }
-
-            if (this.discoveryResultCache.TotalDiscoveredTests == 0)
+            finally
             {
-                DiscovererEnumerator.LogWarningOnNoTestsDiscovered(sources, testCaseFilter, logger);
+                this.CollectTelemetryAtEnd(totalTimeTakenByAdapters, totalAdaptersUsed);
             }
+        }
 
-            this.CollectTelemetryAtEnd(totalTimeTakenByAdapters, totalAdaptersUsed);
+        private void LogTestsDiscoveryCancellation(IMessageLogger logger)
+        {
+            logger.SendMessage(TestMessageLevel.Warning, CrossPlatEngineResources.TestDiscoveryCancelled);
         }
 
         private void CollectTelemetryAtEnd(double totalTimeTakenByAdapters, double totalAdaptersUsed)
@@ -427,16 +464,16 @@ private static bool IsAssembly(string filePath)
                 ".exe".Equals(fileExtension, StringComparison.OrdinalIgnoreCase);
         }
 
-        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes")]
-        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design",
+        [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes")]
+        [SuppressMessage("Microsoft.Design",
             "CA1006:DoNotNestGenericTypesInMemberSignatures")]
         private static IEnumerable<LazyExtension<ITestDiscoverer, ITestDiscovererCapabilities>> GetDiscoverers(
             string extensionAssembly,
             bool throwOnError)
         {
             try
             {
-                if (string.IsNullOrEmpty(extensionAssembly) || string.Equals(extensionAssembly, ObjectModel.Constants.UnspecifiedAdapterPath))
+                if (string.IsNullOrEmpty(extensionAssembly) || string.Equals(extensionAssembly, Constants.UnspecifiedAdapterPath))
                 {
                     // full discovery.
                     return TestDiscoveryExtensionManager.Create().Discoverers;
@@ -448,10 +485,7 @@ private static IEnumerable<LazyExtension<ITestDiscoverer, ITestDiscovererCapabil
             }
             catch (Exception ex)
             {
-                EqtTrace.Error(
-                    "TestDiscoveryManager: LoadExtensions: Exception occured while loading extensions {0}",
-                    ex);
-
+                EqtTrace.Error($"TestDiscoveryManager: LoadExtensions: Exception occured while loading extensions {ex}");
                 if (throwOnError)
                 {
                     throw;

src/Microsoft.TestPlatform.CrossPlatEngine/Discovery/DiscoveryManager.cs

@@ -9,7 +9,7 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Discovery
     using System.Globalization;
     using System.IO;
     using System.Linq;
-
+    using System.Threading;
     using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework;
     using Microsoft.VisualStudio.TestPlatform.Common.Logging;
     using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
@@ -28,11 +28,12 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Discovery
     /// </summary>
     public class DiscoveryManager : IDiscoveryManager
     {
-        private TestSessionMessageLogger sessionMessageLogger;
-        private ITestPlatformEventSource testPlatformEventSource;
-        private IRequestData requestData;
+        private readonly TestSessionMessageLogger sessionMessageLogger;
+        private readonly ITestPlatformEventSource testPlatformEventSource;
+        private readonly IRequestData requestData;
         private ITestDiscoveryEventsHandler2 testDiscoveryEventsHandler;
         private DiscoveryCriteria discoveryCriteria;
+        private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
 
         /// <summary>
         /// Initializes a new instance of the <see cref="DiscoveryManager"/> class.
@@ -109,7 +110,7 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
                 // If there are sources to discover
                 if (verifiedExtensionSourceMap.Any())
                 {
-                    new DiscovererEnumerator(this.requestData, discoveryResultCache).LoadTests(
+                    new DiscovererEnumerator(this.requestData, discoveryResultCache, cancellationTokenSource.Token).LoadTests(
                         verifiedExtensionSourceMap,
                         RunSettingsUtilities.CreateAndInitializeRunSettings(discoveryCriteria.RunSettings),
                         discoveryCriteria.TestCaseFilter,
@@ -138,9 +139,10 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
 
                     // Collecting Total Tests Discovered
                     this.requestData.MetricsCollection.Add(TelemetryDataConstants.TotalTestsDiscovered, totalDiscoveredTestCount);
-
-                    var discoveryCompleteEventsArgs = new DiscoveryCompleteEventArgs(totalDiscoveredTestCount, false);
-                    discoveryCompleteEventsArgs.Metrics = this.requestData.MetricsCollection.Metrics;
+                    var discoveryCompleteEventsArgs = new DiscoveryCompleteEventArgs(totalDiscoveredTestCount, false)
+                    {
+                        Metrics = this.requestData.MetricsCollection.Metrics
+                    };
 
                     eventHandler.HandleDiscoveryComplete(discoveryCompleteEventsArgs, lastChunk);
                 }
@@ -161,7 +163,7 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
         /// </summary>
         public void Abort()
         {
-            // do nothing for now.
+            this.cancellationTokenSource.Cancel();
         }
 
         private void OnReportTestCases(IEnumerable<TestCase> testCases)

src/Microsoft.TestPlatform.CrossPlatEngine/Resources/Resources.resx

@@ -195,4 +195,7 @@
   <data name="NoTestsAvailableForGivenTestCaseFilter" xml:space="preserve">
     <value>No test matches the given testcase filter `{0}` in {1}</value>
   </data>
+  <data name="TestDiscoveryCancelled" xml:space="preserve">
+    <value>Discovery of tests cancelled.</value>
+  </data>
 </root>