Enable parallel discovery (#3349)

Co-authored-by: Medeni Baykal <433724+Haplois@users.noreply.github.com>
Co-authored-by: Sanan Yuzbashiyev <sayuzbas@microsoft.com>
Co-authored-by: Amaury Levé <amauryleve@microsoft.com>

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

@@ -70,7 +70,7 @@ public void DiscoverAsync()
  {
  if (_disposed)
  {
- throw new ObjectDisposedException("DiscoveryRequest");
+ throw new ObjectDisposedException(nameof(DiscoveryRequest));
  }
 
  // Reset the discovery completion event
@@ -112,12 +112,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 (Constants.DefaultProtocolConfig.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);
 }