Enable parallel discovery

scripts/build/TestPlatform.Dependencies.props

@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <Project ToolsVersion="4.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <PropertyGroup>
     <VSSdkBuildToolsVersion>15.8.3247</VSSdkBuildToolsVersion>

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

@@ -226,7 +226,16 @@ private void ProcessRequests(ITestRequestManager testRequestManager)
 
                         case MessageType.CancelDiscovery:
                             {
-                                testRequestManager.CancelDiscovery();
+                                // If testhost has old version, we should use old cancel logic
+                                // to be consistent and not create regression issues
+                                if (this.protocolConfig.Version < ObjectModel.Constants.MinimumProtocolVersionWithCancelDiscoveryEventHandlerSupport)
+                                {
+                                    testRequestManager.CancelDiscovery();
+                                }
+                                else
+                                {
+                                    testRequestManager.CancelDiscoveryWithEventHandler();
+                                }
                                 break;
                             }
 

src/Microsoft.TestPlatform.Client/Discovery/DiscoveryRequest.cs

@@ -143,6 +143,42 @@ public void Abort()
             }
         }
 
+        /// <inheritdoc/>
+        public void AbortWithEventHandler()
+        {
+            if (EqtTrace.IsVerboseEnabled)
+            {
+                EqtTrace.Verbose("DiscoveryRequest.AbortWithEventHandler: Aborting.");
+            }
+
+            lock (this.syncObject)
+            {
+                if (this.disposed)
+                {
+                    throw new ObjectDisposedException("DiscoveryRequest");
+                }
+
+                if (this.discoveryInProgress)
+                {
+                    this.DiscoveryManager.Abort(this);
+                }
+                else
+                {
+                    if (EqtTrace.IsInfoEnabled)
+                    {
+                        EqtTrace.Info("DiscoveryRequest.AbortWithEventHandler: No operation to abort.");
+                    }
+
+                    return;
+                }
+            }
+
+            if (EqtTrace.IsInfoEnabled)
+            {
+                EqtTrace.Info("DiscoveryRequest.AbortWithEventHandler: Aborted.");
+            }
+        }
+
         /// <summary>
         /// Wait for discovery completion
         /// </summary>

src/Microsoft.TestPlatform.Client/PublicAPI/PublicAPI.Unshipped.txt

@@ -0,0 +1,2 @@
+Microsoft.VisualStudio.TestPlatform.Client.Discovery.DiscoveryRequest.AbortWithEventHandler() -> void
+Microsoft.VisualStudio.TestPlatform.Client.RequestHelper.ITestRequestManager.CancelDiscoveryWithEventHandler() -> void

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

@@ -102,6 +102,11 @@ void StartTestSession(
         /// </summary>
         void CancelDiscovery();
 
+        /// <summary>
+        /// Cancels the current discovery request with discovery complete event handler
+        /// </summary>
+        void CancelDiscoveryWithEventHandler();
+
         /// <summary>
         /// Cancels the current test run attachments processing request.
         /// </summary>

src/Microsoft.TestPlatform.Common/Interfaces/Engine/ClientProtocol/IParallelProxyDiscoveryManager.cs

@@ -10,6 +10,11 @@ namespace Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine
     /// </summary>
     public interface IParallelProxyDiscoveryManager : IParallelOperationManager, IProxyDiscoveryManager
     {
+        /// <summary>
+        /// Indicates if user requested an abortion
+        /// </summary>
+        bool IsAbortRequested { get; set; }
+
         /// <summary>
         /// Handles Partial Discovery Complete event coming from a specific concurrent proxy discovery manager
         /// Each concurrent proxy discovery manager will signal the parallel discovery manager when its complete
@@ -25,4 +30,14 @@ bool HandlePartialDiscoveryComplete(
             IEnumerable<TestCase> lastChunk,
             bool isAborted);
     }
+
+    /// <summary>
+    /// Enums for indicating discovery status of source
+    /// </summary>
+    public enum DiscoveryStatus
+    {
+        FullyDiscovered, // Indicates that source was fully discovered
+        PartiallyDiscovered, // Indicates that we started discovery of the source but something happened (cancel/abort) and we stopped processing it
+        NotDiscovered // Indicates the sources which were not touched during discovery
+    }
 }

src/Microsoft.TestPlatform.Common/Interfaces/Engine/ClientProtocol/IProxyDiscoveryManager.cs

@@ -28,6 +28,12 @@ public interface IProxyDiscoveryManager
         /// </summary>
         void Abort();
 
+        /// <summary>
+        /// Aborts discovery operation with EventHandler.
+        /// </summary>
+        /// <param name="eventHandler">EventHandler for handling discovery events from Engine</param>
+        void Abort(ITestDiscoveryEventsHandler2 eventHandler);
+
         /// <summary>
         /// Closes the current test operation.
         /// Send a EndSession message to close the test host and channel gracefully.

src/Microsoft.TestPlatform.Common/Interfaces/Engine/TesthostProtocol/IDiscoveryManager.cs

@@ -29,5 +29,11 @@ public interface IDiscoveryManager
         /// Aborts the test discovery.
         /// </summary>
         void Abort();
+
+        /// <summary>
+        /// Aborts the test discovery with eventHandler
+        /// </summary>
+        /// <param name="eventHandler">EventHandler for handling discovery events from Engine</param>
+        void Abort(ITestDiscoveryEventsHandler2 eventHandler);
     }
 }

src/Microsoft.TestPlatform.Common/PublicAPI/PublicAPI.Unshipped.txt

@@ -0,0 +1,8 @@
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus.FullyDiscovered = 0 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus.NotDiscovered = 2 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus.PartiallyDiscovered = 1 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.IsAbortRequested.get -> bool
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.IsAbortRequested.set -> void
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IProxyDiscoveryManager.Abort(Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.ITestDiscoveryEventsHandler2 eventHandler) -> void
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.TesthostProtocol.IDiscoveryManager.Abort(Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.ITestDiscoveryEventsHandler2 eventHandler) -> void

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

@@ -87,6 +87,11 @@ public interface ITestRequestSender : IDisposable
         /// </summary>
         void SendTestRunAbort();
 
+        /// <summary>
+        /// Send the request to abort the discovery
+        /// </summary>
+        void SendDiscoveryAbort();
+
         /// <summary>
         /// Handle client process exit
         /// </summary>

src/Microsoft.TestPlatform.CommunicationUtilities/JsonDataSerializer.cs

@@ -230,6 +230,7 @@ private JsonSerializer GetPayloadSerializer(int? version)
                 case 2:
                 case 4:
                 case 5:
+                case 6:
                     return payloadSerializer2;
                 default:
                     throw new NotSupportedException($"Protocol version {version} is not supported. " +

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

@@ -30,5 +30,20 @@ public class DiscoveryCompletePayload
         /// Gets or sets the Metrics
         /// </summary>
         public IDictionary<string, object> Metrics { get; set; }
+
+        /// <summary>
+        /// Gets or sets list of sources which were fully discovered
+        /// </summary>
+        public IReadOnlyCollection<string> FullyDiscoveredSources { get; set; } = new List<string>();
+
+        /// <summary>
+        /// Gets or sets list of sources which were partially discovered (started discover tests, but then discovery aborted)
+        /// </summary>
+        public IReadOnlyCollection<string> PartiallyDiscoveredSources { get; set; } = new List<string>();
+
+        /// <summary>
+        /// Gets or sets list of sources which were not discovered at all
+        /// </summary>
+        public IReadOnlyCollection<string> NotDiscoveredSources { get; set; } = new List<string>();
     }
 }

src/Microsoft.TestPlatform.CommunicationUtilities/PublicAPI/PublicAPI.Unshipped.txt

@@ -0,0 +1,8 @@
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces.ITestRequestSender.SendDiscoveryAbort() -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.FullyDiscoveredSources.get -> System.Collections.Generic.IReadOnlyCollection<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.FullyDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.NotDiscoveredSources.get -> System.Collections.Generic.IReadOnlyCollection<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.NotDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.PartiallyDiscoveredSources.get -> System.Collections.Generic.IReadOnlyCollection<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.PartiallyDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.TestRequestSender.SendDiscoveryAbort() -> void

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -14,8 +14,8 @@ namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Client;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Host;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
-    using CommonResources = Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Resources.Resources;
-    using ObjectModelConstants = Microsoft.VisualStudio.TestPlatform.ObjectModel.Constants;
+    using CommonResources = Resources.Resources;
+    using ObjectModelConstants = TestPlatform.ObjectModel.Constants;
 
     /// <summary>
     /// Test request sender implementation.
@@ -54,7 +54,7 @@ public class TestRequestSender : ITestRequestSender
 
         // Must be in sync with the highest supported version in
         // src/Microsoft.TestPlatform.CrossPlatEngine/EventHandlers/TestRequestHandler.cs file.
-        private int highestSupportedVersion = 5;
+        private int highestSupportedVersion = 6;
 
         private TestHostConnectionInfo connectionInfo;
 
@@ -290,6 +290,24 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
 
             this.channel.Send(message);
         }
+
+        /// <inheritdoc/>
+        public void SendDiscoveryAbort()
+        {
+            if (this.IsOperationComplete())
+            {
+                EqtTrace.Verbose("TestRequestSender.SendDiscoveryAbort: Operation is already complete. Skip error message.");
+                return;
+            }
+
+            if (EqtTrace.IsVerboseEnabled)
+            {
+                EqtTrace.Verbose("TestRequestSender.SendDiscoveryAbort: Sending discovery abort.");
+            }
+
+            this.channel?.Send(this.dataSerializer.SerializeMessage(MessageType.CancelDiscovery));
+        }
+
         #endregion
 
         #region Execution Protocol
@@ -595,7 +613,12 @@ private void OnDiscoveryMessageReceived(ITestDiscoveryEventsHandler2 discoveryEv
                     case MessageType.DiscoveryComplete:
                         var discoveryCompletePayload =
                             this.dataSerializer.DeserializePayload<DiscoveryCompletePayload>(data);
-                        var discoveryCompleteEventArgs = new DiscoveryCompleteEventArgs(discoveryCompletePayload.TotalTests, discoveryCompletePayload.IsAborted);
+                        var discoveryCompleteEventArgs = new DiscoveryCompleteEventArgs(
+                                                                                        discoveryCompletePayload.TotalTests,
+                                                                                        discoveryCompletePayload.IsAborted,
+                                                                                        discoveryCompletePayload.FullyDiscoveredSources,
+                                                                                        discoveryCompletePayload.PartiallyDiscoveredSources,
+                                                                                        discoveryCompletePayload.NotDiscoveredSources);
                         discoveryCompleteEventArgs.Metrics = discoveryCompletePayload.Metrics;
                         discoveryEventsHandler.HandleDiscoveryComplete(
                             discoveryCompleteEventArgs,

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

@@ -95,6 +95,12 @@ public void Abort()
             Task.Run(() => this.testHostManagerFactory.GetDiscoveryManager().Abort());
         }
 
+        /// <inheritdoc/>
+        public void Abort(ITestDiscoveryEventsHandler2 eventHandler)
+        {
+            Task.Run(() => this.testHostManagerFactory.GetDiscoveryManager().Abort(eventHandler));
+        }
+
         private void InitializeExtensions(IEnumerable<string> sources)
         {
             var extensionsFromSource = this.testHostManager.GetTestPlatformExtensions(sources, Enumerable.Empty<string>());

src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryDataAggregator.cs

@@ -4,6 +4,8 @@
 namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Client.Parallel
 {
     using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
+    using Microsoft.VisualStudio.TestPlatform.ObjectModel;
+    using Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine;
     using System;
     using System.Collections.Concurrent;
     using System.Collections.Generic;
@@ -42,6 +44,20 @@ public ParallelDiscoveryDataAggregator()
 
         #endregion
 
+        #region Internal Properties
+
+        /// <summary>
+        /// Dictionary which stores source with corresponding discoveryStatus
+        /// </summary>
+        internal ConcurrentDictionary<string, DiscoveryStatus> SourcesWithDiscoveryStatus = new ConcurrentDictionary<string, DiscoveryStatus>();
+
+        /// <summary>
+        /// Indicates if discovery complete payload already sent back to IDE
+        /// </summary>
+        internal bool IsMessageSent { get; set; }
+
+        #endregion
+
         #region Public Methods
 
         /// <summary>
@@ -128,6 +144,58 @@ public void AggregateDiscoveryDataMetrics(IDictionary<string, object> metrics)
             }
         }
 
+        /// <summary>
+        /// Aggregate the source as fully discovered
+        /// </summary>
+        /// <param name="sorce">Fully discovered source</param>
+        internal void AggregateTheSourcesWithDiscoveryStatus(IEnumerable<string> sources, DiscoveryStatus status)
+        {
+            if (sources == null || sources.Count() == 0) return;
+
+            foreach (var source in sources)
+            {
+                if (status == DiscoveryStatus.NotDiscovered) SourcesWithDiscoveryStatus[source] = status;
+
+                if (!SourcesWithDiscoveryStatus.ContainsKey(source))
+                {
+                    EqtTrace.Warning($"ParallelDiscoveryDataAggregator.AggregateTheSourcesWithDiscoveryStatus : " +
+                                     $"{source} is not present in SourcesWithDiscoveryStatus dictionary");
+                }
+                else
+                {
+                    SourcesWithDiscoveryStatus[source] = status;
+
+                    if (EqtTrace.IsInfoEnabled)
+                    {
+                        EqtTrace.Info($"ParallelDiscoveryDataAggregator.AggregateTheSourcesWithDiscoveryStatus : " +
+                                         $"{source} is marked with {status} status");
+                    }
+                }
+            }
+        }
+
+        /// <summary>
+        /// Aggregate the value indicating if we already sent message to IDE
+        /// </summary>
+        /// <param name="isMessageSent">Boolean value if we already sent message to IDE</param>
+        internal void AggregateIsMessageSent(bool isMessageSent)
+        {
+            this.IsMessageSent = this.IsMessageSent || isMessageSent;
+        }
+
+        /// <summary>
+        /// Returning sources with particular discovery status
+        /// </summary>
+        /// <param name="status">Status to filter</param>
+        /// <returns></returns>
+        internal ICollection<string> GetSourcesWithStatus(DiscoveryStatus status)
+        {
+            if (SourcesWithDiscoveryStatus == null || SourcesWithDiscoveryStatus.IsEmpty) return new List<string>();
+
+            return SourcesWithDiscoveryStatus.Where(source => source.Value == status)
+                                             .Select(source => source.Key).ToList();
+        }
+
         #endregion
     }
 }