Pass workloads to proxy managers

playground/MSTest1/MSTest1.csproj

@@ -7,6 +7,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Microsoft.NET.Test.SDK" Version="17.5.0" />
     <PackageReference Include="MSTest.TestAdapter" Version="2.1.0" />
     <PackageReference Include="MSTest.TestFramework" Version="2.1.0" />
     <PackageReference Include="coverlet.collector" Version="1.2.0" />

playground/MSTest2/MSTest2.csproj

@@ -7,6 +7,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Microsoft.NET.Test.SDK" Version="17.5.0" />
     <PackageReference Include="MSTest.TestAdapter" Version="2.1.0" />
     <PackageReference Include="MSTest.TestFramework" Version="2.1.0" />
     <PackageReference Include="coverlet.collector" Version="1.2.0" />

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

@@ -24,7 +24,7 @@ internal sealed class ParallelOperationManager<TManager, TEventHandler, TWorkloa
                 out int num)
         ? num
         : PreStart;
-    private readonly Func<TestRuntimeProviderInfo, TManager> _createNewManager;
+    private readonly Func<TestRuntimeProviderInfo, TWorkload, TManager> _createNewManager;
 
     /// <summary>
     /// Default number of Processes
@@ -50,7 +50,7 @@ internal sealed class ParallelOperationManager<TManager, TEventHandler, TWorkloa
     /// <param name="createNewManager">Creates a new manager that is responsible for running a single part of the overall workload.
     /// A manager is typically a testhost, and the part of workload is discovering or running a single test dll.</param>
     /// <param name="parallelLevel">Determines the maximum amount of parallel managers that can be active at the same time.</param>
-    public ParallelOperationManager(Func<TestRuntimeProviderInfo, TManager> createNewManager, int parallelLevel)
+    public ParallelOperationManager(Func<TestRuntimeProviderInfo, TWorkload, TManager> createNewManager, int parallelLevel)
     {
         _createNewManager = createNewManager;
         MaxParallelLevel = parallelLevel;
@@ -144,7 +144,7 @@ private bool RunWorkInParallel()
                 var workload = workloadsToAdd[i];
                 slot.ShouldPreStart = occupiedSlots + i + 1 > MaxParallelLevel;
 
-                var manager = _createNewManager(workload.Provider);
+                var manager = _createNewManager(workload.Provider, workload.Work);
                 var eventHandler = _getEventHandler(_eventHandler, manager);
                 slot.EventHandler = eventHandler;
                 slot.Manager = manager;

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

@@ -43,7 +43,7 @@ internal sealed class ParallelProxyDiscoveryManager : IParallelProxyDiscoveryMan
 
     public ParallelProxyDiscoveryManager(
         IRequestData requestData,
-        Func<TestRuntimeProviderInfo, IProxyDiscoveryManager> actualProxyManagerCreator,
+        Func<TestRuntimeProviderInfo, DiscoveryCriteria, IProxyDiscoveryManager> actualProxyManagerCreator,
         DiscoveryDataAggregator dataAggregator,
         int parallelLevel,
         List<TestRuntimeProviderInfo> testHostProviders)
@@ -53,7 +53,7 @@ public ParallelProxyDiscoveryManager(
 
     internal ParallelProxyDiscoveryManager(
         IRequestData requestData,
-        Func<TestRuntimeProviderInfo, IProxyDiscoveryManager> actualProxyManagerCreator,
+        Func<TestRuntimeProviderInfo, DiscoveryCriteria, IProxyDiscoveryManager> actualProxyManagerCreator,
         DiscoveryDataAggregator dataAggregator,
         IDataSerializer dataSerializer,
         int parallelLevel,

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

@@ -64,7 +64,7 @@ internal sealed class ParallelProxyExecutionManager : IParallelProxyExecutionMan
 
     public ParallelProxyExecutionManager(
         IRequestData requestData,
-        Func<TestRuntimeProviderInfo, IProxyExecutionManager> actualProxyManagerCreator,
+        Func<TestRuntimeProviderInfo, TestRunCriteria, IProxyExecutionManager> actualProxyManagerCreator,
         int parallelLevel,
          List<TestRuntimeProviderInfo> testHostProviders)
         : this(requestData, actualProxyManagerCreator, JsonDataSerializer.Instance, parallelLevel, testHostProviders)
@@ -73,7 +73,7 @@ public ParallelProxyExecutionManager(
 
     internal ParallelProxyExecutionManager(
         IRequestData requestData,
-        Func<TestRuntimeProviderInfo, IProxyExecutionManager> actualProxyManagerCreator,
+        Func<TestRuntimeProviderInfo, TestRunCriteria, IProxyExecutionManager> actualProxyManagerCreator,
         IDataSerializer dataSerializer,
         int parallelLevel,
         List<TestRuntimeProviderInfo> testHostProviders)

src/Microsoft.TestPlatform.CrossPlatEngine/TestEngine.cs

@@ -118,9 +118,9 @@ public IProxyDiscoveryManager GetDiscoveryManager(
         // discovery manager to publish its current state. But doing so we are losing the collected state of all the
         // other managers.
         var discoveryDataAggregator = new DiscoveryDataAggregator();
-        Func<TestRuntimeProviderInfo, IProxyDiscoveryManager> proxyDiscoveryManagerCreator = runtimeProviderInfo =>
+        Func<TestRuntimeProviderInfo, DiscoveryCriteria, IProxyDiscoveryManager> proxyDiscoveryManagerCreator = (runtimeProviderInfo, discoveryCriteria) =>
         {
-            var sources = runtimeProviderInfo.SourceDetails.Select(r => r.Source!).ToList();
+            var sources = discoveryCriteria.Sources.ToList();
             var hostManager = _testHostProviderManager.GetTestHostManagerByRunConfiguration(runtimeProviderInfo.RunSettings, sources);
             hostManager?.Initialize(TestSessionMessageLogger.Instance, runtimeProviderInfo.RunSettings!);
 
@@ -241,15 +241,16 @@ public IProxyExecutionManager GetExecutionManager(
         }
 
         // This creates a single non-parallel execution manager, based requestData, isDataCollectorEnabled and the
-        // overall testRunCriteria. The overall testRunCriteria are split to smaller pieces (e.g. each source from the overall
-        // testRunCriteria) so we can run them in parallel, and those are then passed to those non-parallel execution managers.
+        // split testRunCriteria. The overall testRunCriteria are split to smaller pieces (e.g. each source from the overall
+        // testRunCriteria) so we can run them in parallel.
         //
         // The function below grabs most of the parameter via closure from the local context,
-        // but gets the runtime provider later, because that is specific info to the source (or sources) it will be running.
+        // but gets the runtime provider later, as well as the discovery request, because that is specific info to the source (or sources)
+        // it will be running.
         // This creator does not get those smaller pieces of testRunCriteria, those come later when we call a method on
         // the non-parallel execution manager we create here. E.g. StartTests(<single piece of testRunCriteria>).
-        Func<TestRuntimeProviderInfo, IProxyExecutionManager> proxyExecutionManagerCreator = runtimeProviderInfo =>
-            CreateNonParallelExecutionManager(requestData, testRunCriteria, isDataCollectorEnabled, runtimeProviderInfo);
+        Func<TestRuntimeProviderInfo, TestRunCriteria, IProxyExecutionManager> proxyExecutionManagerCreator = (runtimeProviderInfo, runCriteria) =>
+            CreateNonParallelExecutionManager(requestData, runCriteria, isDataCollectorEnabled, runtimeProviderInfo);
 
         var executionManager = new ParallelProxyExecutionManager(requestData, proxyExecutionManagerCreator, parallelLevel, testHostProviders);