Enable parallel discovery

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

@@ -227,7 +227,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 (_protocolConfig.Version < ObjectModel.Constants.MinimumProtocolVersionWithCancelDiscoveryEventHandlerSupport)
+                            {
+                                testRequestManager.CancelDiscovery();
+                            }
+                            else
+                            {
+                                testRequestManager.CancelDiscoveryWithEventHandler();
+                            }
                             break;
                         }
 

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

@@ -164,6 +164,42 @@ bool IRequest.WaitForCompletion(int timeout)
         return _discoveryCompleted == null || _discoveryCompleted.WaitOne(timeout);
     }
 
+    /// <inheritdoc/>
+    public void AbortWithEventHandler()
+    {
+        if (EqtTrace.IsVerboseEnabled)
+        {
+            EqtTrace.Verbose("DiscoveryRequest.AbortWithEventHandler: Aborting.");
+        }
+
+        lock (_syncObject)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException("DiscoveryRequest");
+            }
+
+            if (DiscoveryInProgress)
+            {
+                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>
     /// Raised when the test discovery starts.
     /// </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

@@ -20,7 +20,7 @@ public interface ITestRequestManager : IDisposable
     /// <summary>
     /// Initializes the extensions while probing additional paths.
     /// </summary>
-    /// 
+    ///
     /// <param name="pathToAdditionalExtensions">Paths to additional extensions.</param>
     /// <param name="skipExtensionFilters">Skip extension filtering by name if true.</param>
     void InitializeExtensions(
@@ -35,7 +35,7 @@ void InitializeExtensions(
     /// <summary>
     /// Discovers tests given a list of sources and some run settings.
     /// </summary>
-    /// 
+    ///
     /// <param name="discoveryPayload">Discovery payload.</param>
     /// <param name="disoveryEventsRegistrar">Discovery events registrar.</param>
     /// <param name="protocolConfig">Protocol related information.</param>
@@ -47,7 +47,7 @@ void DiscoverTests(
     /// <summary>
     /// Runs tests given a list of sources and some run settings.
     /// </summary>
-    /// 
+    ///
     /// <param name="testRunRequestPayLoad">Test run request payload.</param>
     /// <param name="customTestHostLauncher">Custom test host launcher for the run.</param>
     /// <param name="testRunEventsRegistrar">Run events registrar.</param>
@@ -61,7 +61,7 @@ void RunTests(
     /// <summary>
     /// Processes test run attachments.
     /// </summary>
-    /// 
+    ///
     /// <param name="testRunAttachmentsProcessingPayload">
     /// Test run attachments processing payload.
     /// </param>
@@ -77,7 +77,7 @@ void ProcessTestRunAttachments(
     /// <summary>
     /// Starts a test session.
     /// </summary>
-    /// 
+    ///
     /// <param name="payload">The start test session payload.</param>
     /// <param name="testHostLauncher">The custom test host launcher.</param>
     /// <param name="eventsHandler">The events handler.</param>
@@ -103,6 +103,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
@@ -24,4 +29,14 @@ bool HandlePartialDiscoveryComplete(
         long totalTests,
         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,4 +29,10 @@ 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.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
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus.FullyDiscovered = 0 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus.NotDiscovered = 2 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus.PartiallyDiscovered = 1 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.DiscoveryStatus

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

@@ -214,7 +214,7 @@ private JsonSerializer GetPayloadSerializer(int? version)
             // unless this is disabled by VSTEST_DISABLE_PROTOCOL_3_VERSION_DOWNGRADE
             // env variable.
             0 or 1 or 3 => s_payloadSerializer,
-            2 or 4 or 5 => s_payloadSerializer2,
+            2 or 4 or 5 or 6 => s_payloadSerializer2,
             _ => throw new NotSupportedException($"Protocol version {version} is not supported. " +
                 "Ensure it is compatible with the latest serializer or add a new one."),
         };

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

@@ -31,4 +31,19 @@ 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

@@ -58,7 +58,7 @@ public class TestRequestSender : ITestRequestSender
 
     // Must be in sync with the highest supported version in
     // src/Microsoft.TestPlatform.CrossPlatEngine/EventHandlers/TestRequestHandler.cs file.
-    private readonly int _highestSupportedVersion = 5;
+    private readonly int _highestSupportedVersion = 6;
 
     private readonly TestHostConnectionInfo _connectionInfo;
 
@@ -295,6 +295,23 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
 
         _channel.Send(message);
     }
+
+    /// <inheritdoc/>
+    public void SendDiscoveryAbort()
+    {
+        if (IsOperationComplete())
+        {
+            EqtTrace.Verbose("TestRequestSender.SendDiscoveryAbort: Operation is already complete. Skip error message.");
+            return;
+        }
+
+        if (EqtTrace.IsVerboseEnabled)
+        {
+            EqtTrace.Verbose("TestRequestSender.SendDiscoveryAbort: Sending discovery abort.");
+        }
+
+        _channel?.Send(_dataSerializer.SerializeMessage(MessageType.CancelDiscovery));
+    }
     #endregion
 
     #region Execution Protocol
@@ -594,7 +611,12 @@ private void OnDiscoveryMessageReceived(ITestDiscoveryEventsHandler2 discoveryEv
                 case MessageType.DiscoveryComplete:
                     var discoveryCompletePayload =
                         _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(() => _testHostManagerFactory.GetDiscoveryManager().Abort());
     }
 
+    /// <inheritdoc/>
+    public void Abort(ITestDiscoveryEventsHandler2 eventHandler)
+    {
+        Task.Run(() => _testHostManagerFactory.GetDiscoveryManager().Abort(eventHandler));
+    }
+
     private void InitializeExtensions(IEnumerable<string> sources)
     {
         var extensionsFromSource = _testHostManager.GetTestPlatformExtensions(sources, Enumerable.Empty<string>());

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

@@ -5,11 +5,15 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Client.Parallel;
 
 using Common.Telemetry;
 
+using Microsoft.VisualStudio.TestPlatform.ObjectModel;
+
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 
+using static Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager;
+
 /// <summary>
 /// ParallelDiscoveryDataAggregator aggregates discovery data from parallel discovery managers
 /// </summary>
@@ -43,6 +47,20 @@ public ParallelDiscoveryDataAggregator()
 
     #endregion
 
+    #region Internal Properties
+
+    /// <summary>
+    /// Dictionary which stores source with corresponding discoveryStatus
+    /// </summary>
+    internal ConcurrentDictionary<string, DiscoveryStatus> SourcesWithDiscoveryStatus = new();
+
+    /// <summary>
+    /// Indicates if discovery complete payload already sent back to IDE
+    /// </summary>
+    internal bool IsMessageSent { get; set; }
+
+    #endregion
+
     #region Public Methods
 
     /// <summary>
@@ -128,5 +146,57 @@ 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)
+    {
+        IsMessageSent = 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
 }