Enable parallel discovery

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

@@ -38,7 +38,7 @@ public class DesignModeClient : IDesignModeClient
     private readonly ICommunicationManager _communicationManager;
     private readonly IDataSerializer _dataSerializer;
 
-    private readonly ProtocolConfig _protocolConfig = ObjectModel.Constants.DefaultProtocolConfig;
+    private readonly ProtocolConfig _protocolConfig = Constants.DefaultProtocolConfig;
     private readonly IEnvironment _platformEnvironment;
     private readonly TestSessionMessageLogger _testSessionMessageLogger;
     private readonly object _lockObject = new();

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

@@ -26,6 +26,7 @@ namespace Microsoft.VisualStudio.TestPlatform.Client.Discovery;
 public sealed class DiscoveryRequest : IDiscoveryRequest, ITestDiscoveryEventsHandler2
 {
     private readonly IDataSerializer _dataSerializer;
+    internal /* for testing purposes */ readonly ProtocolConfig ProtocolConfig = Constants.DefaultProtocolConfig;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="DiscoveryRequest"/> class.
@@ -70,7 +71,7 @@ public void DiscoverAsync()
         {
             if (_disposed)
             {
-                throw new ObjectDisposedException("DiscoveryRequest");
+                throw new ObjectDisposedException(nameof(DiscoveryRequest));
             }
 
             // Reset the discovery completion event
@@ -112,12 +113,21 @@ public void Abort()
         {
             if (_disposed)
             {
-                throw new ObjectDisposedException("DiscoveryRequest");
+                throw new ObjectDisposedException(nameof(DiscoveryRequest));
             }
 
             if (DiscoveryInProgress)
             {
-                DiscoveryManager.Abort();
+                // If testhost has old version, we should use old cancel logic
+                // to be consistent and not create regression issues
+                if (ProtocolConfig.Version < Constants.MinimumProtocolVersionWithCancelDiscoveryEventHandlerSupport)
+                {
+                    DiscoveryManager.Abort();
+                }
+                else
+                {
+                    DiscoveryManager.Abort(this);
+                }
             }
             else
             {

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

@@ -22,7 +22,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(
@@ -37,7 +37,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>
@@ -49,7 +49,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>
@@ -63,7 +63,7 @@ void RunTests(
     /// <summary>
     /// Processes test run attachments.
     /// </summary>
-    /// 
+    ///
     /// <param name="testRunAttachmentsProcessingPayload">
     /// Test run attachments processing payload.
     /// </param>
@@ -79,7 +79,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>

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

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine;
+
+/// <summary>
+/// Enums for indicating discovery status of source
+/// </summary>
+public enum DiscoveryStatus
+{
+    /// <summary>
+    /// Indicates the sources which were not touched during discovery.
+    /// </summary>
+    NotDiscovered,
+
+    /// <summary>
+    /// Indicates that we started discovery of the source but something happened (cancel/abort) and we stopped processing it.
+    /// </summary>
+    PartiallyDiscovered,
+
+    /// <summary>
+    /// Indicates that source was fully discovered.
+    /// </summary>
+    FullyDiscovered,
+}

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

@@ -12,6 +12,11 @@ namespace Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine;
 /// </summary>
 public interface IParallelProxyDiscoveryManager : IParallelOperationManager, IProxyDiscoveryManager
 {
+    /// <summary>
+    /// Indicates if user requested an abortion
+    /// </summary>
+    bool IsAbortRequested { get; }
+
     /// <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

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

@@ -30,6 +30,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

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

@@ -6,3 +6,10 @@ static Microsoft.VisualStudio.TestPlatform.Common.Telemetry.TelemetryDataConstan
 static Microsoft.VisualStudio.TestPlatform.Common.Telemetry.TelemetryDataConstants.StopTestSessionCompleteEvent -> string
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IProxyTestSessionManager.StartSession(Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.ITestSessionEventsHandler eventsHandler, Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.IRequestData requestData) -> bool
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IProxyTestSessionManager.StopSession(Microsoft.VisualStudio.TestPlatform.ObjectModel.Client.IRequestData requestData) -> bool
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.IParallelProxyDiscoveryManager.IsAbortRequested.get -> bool
+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.DiscoveryStatus
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus.NotDiscovered = 0 -> 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.DiscoveryStatus.FullyDiscovered = 2 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine.DiscoveryStatus

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

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

src/Microsoft.TestPlatform.CommunicationUtilities/JsonDataSerializer.cs

@@ -208,6 +208,15 @@ private JsonSerializer GetPayloadSerializer(int? version)
             version = 1;
         }
 
+        // 0: the original protocol with no versioning (Message). It is used during negotiation.
+        // 1: new protocol with versioning (VersionedMessage).
+        // 2: changed serialization because the serialization of properties in bag was too verbose,
+        //    so common properties are considered built-in and serialized without type info.
+        // 3: introduced because of changes to allow attaching debugger to external process.
+        // 4: introduced because 3 did not update this table and ended up using the serializer for protocol v1,
+        //    which is extremely slow. We negotiate 2 or 4, but never 3 unless the flag above is set.
+        // 5: ???
+        // 6: accepts abort and cancel with handlers that report the status.
         return version switch
         {
             // 0 is used during negotiation.
@@ -216,9 +225,9 @@ 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,
-            _ => throw new NotSupportedException($"Protocol version {version} is not supported. " +
-                "Ensure it is compatible with the latest serializer or add a new one."),
+            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

@@ -33,4 +33,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 IList<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 IList<string> PartiallyDiscoveredSources { get; set; } = new List<string>();
+
+    /// <summary>
+    /// Gets or sets list of sources which were not discovered at all.
+    /// </summary>
+    public IList<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.IList<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.FullyDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.NotDiscoveredSources.get -> System.Collections.Generic.IList<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.NotDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.PartiallyDiscoveredSources.get -> System.Collections.Generic.IList<string>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.PartiallyDiscoveredSources.set -> void
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.TestRequestSender.SendDiscoveryAbort() -> void

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -60,7 +60,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;
 
@@ -276,6 +276,19 @@ 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;
+        }
+
+        EqtTrace.Verbose("TestRequestSender.SendDiscoveryAbort: Sending discovery abort.");
+        _channel?.Send(_dataSerializer.SerializeMessage(MessageType.CancelDiscovery));
+    }
     #endregion
 
     #region Execution Protocol
@@ -543,7 +556,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

@@ -97,6 +97,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

@@ -7,6 +7,9 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Client.Parallel;
 
 using Common.Telemetry;
 
+using Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Discovery;
+using Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine;
+
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
@@ -37,6 +40,17 @@ public ParallelDiscoveryDataAggregator()
     /// </summary>
     public long TotalTests { get; private set; }
 
+
+    /// <summary>
+    /// Dictionary which stores source with corresponding discoveryStatus
+    /// </summary>
+    private readonly ConcurrentDictionary<string, DiscoveryStatus> _sourcesWithDiscoveryStatus = new();
+
+    /// <summary>
+    /// Indicates if discovery complete payload already sent back to IDE
+    /// </summary>
+    internal bool IsMessageSent { get; private set; }
+
     /// <summary>
     /// Returns the Aggregated Metrics.
     /// </summary>
@@ -120,4 +134,27 @@ public void AggregateDiscoveryDataMetrics(IDictionary<string, object> metrics)
         }
     }
 
+    /// <summary>
+    /// Aggregate the source as fully discovered
+    /// </summary>
+    /// <param name="sorce">Fully discovered source</param>
+    public void MarkSourcesWithStatus(ICollection<string> sources, DiscoveryStatus status)
+        => DiscoveryManager.MarkSourcesWithStatus(sources, status, _sourcesWithDiscoveryStatus);
+
+    /// <summary>
+    /// Aggregates the value indicating if we already sent message to IDE.
+    /// </summary>
+    /// <param name="isMessageSent">Boolean value if we already sent message to IDE</param>
+    public void AggregateIsMessageSent(bool isMessageSent)
+    {
+        IsMessageSent = IsMessageSent || isMessageSent;
+    }
+
+    /// <summary>
+    /// Returns sources with particular discovery status.
+    /// </summary>
+    /// <param name="status">Status to filter</param>
+    /// <returns></returns>
+    public List<string> GetSourcesWithStatus(DiscoveryStatus status)
+        => DiscoveryManager.GetSourcesWithStatus(status, _sourcesWithDiscoveryStatus);
 }