Added telemetry data point for extensions loaded during test discovery/run

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

@@ -6,6 +6,8 @@
 using System.Linq;
 using System.Threading;
 
+using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework;
+using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework.Utilities;
 using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces;
@@ -248,6 +250,22 @@ public void HandleDiscoveryComplete(DiscoveryCompleteEventArgs discoveryComplete
                     OnDiscoveredTests.SafeInvoke(this, discoveredTestsEvent, "DiscoveryRequest.DiscoveryComplete");
                 }
 
+                // Add extensions discovered by vstest.console.
+                //
+                // TODO(copoiena): Writing telemetry twice is less than ideal.
+                // We first write telemetry data in the _requestData variable in the ParallelRunEventsHandler
+                // and then we write again here. We should refactor this code and write only once.
+                discoveryCompleteEventArgs.DiscoveredExtensions = TestExtensions.CreateMergedDictionary(
+                    discoveryCompleteEventArgs.DiscoveredExtensions,
+                    TestPluginCache.Instance.TestExtensions.GetCachedExtensions());
+
+                if (RequestData.IsTelemetryOptedIn)
+                {
+                    TestExtensions.AddExtensionTelemetry(
+                        discoveryCompleteEventArgs.Metrics,
+                        discoveryCompleteEventArgs.DiscoveredExtensions);
+                }
+
                 LoggerManager.HandleDiscoveryComplete(discoveryCompleteEventArgs);
                 OnDiscoveryComplete.SafeInvoke(this, discoveryCompleteEventArgs, "DiscoveryRequest.DiscoveryComplete");
             }
@@ -406,6 +424,25 @@ private string UpdateRawMessageWithTelemetryInfo(DiscoveryCompletePayload discov
 
                 // Collecting Total Time Taken
                 discoveryCompletePayload.Metrics[TelemetryDataConstants.TimeTakenInSecForDiscovery] = discoveryFinalTimeTakenForDesignMode.TotalSeconds;
+
+                // Add extensions discovered by vstest.console.
+                //
+                // TODO(copoiena):
+                // Doing extension merging here is incorrect because we can end up not merging the
+                // cached extensions for the current process (i.e. vstest.console) and hence have
+                // an incomplete list of discovered extensions. This can happen because this method
+                // is called only if telemetry is opted in (see: HandleRawMessage). We should handle
+                // this merge a level above in order to be consistent, but that means we'd have to
+                // deserialize all raw messages no matter if telemetry is opted in or not and that
+                // would probably mean a performance hit.
+                discoveryCompletePayload.DiscoveredExtensions = TestExtensions.CreateMergedDictionary(
+                    discoveryCompletePayload.DiscoveredExtensions,
+                    TestPluginCache.Instance.TestExtensions.GetCachedExtensions());
+
+                // Write extensions to telemetry data.
+                TestExtensions.AddExtensionTelemetry(
+                    discoveryCompletePayload.Metrics,
+                    discoveryCompletePayload.DiscoveredExtensions);
             }
 
             if (message is VersionedMessage message1)

src/Microsoft.TestPlatform.Client/Execution/TestRunRequest.cs

@@ -8,6 +8,8 @@
 using System.Linq;
 using System.Threading;
 
+using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework;
+using Microsoft.VisualStudio.TestPlatform.Common.ExtensionFramework.Utilities;
 using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces;
@@ -400,6 +402,22 @@ public void HandleTestRunComplete(TestRunCompleteEventArgs runCompleteArgs!!, Te
                         runCompleteArgs.InvokedDataCollectors,
                         _runRequestTimeTracker.Elapsed);
 
+                // Add extensions discovered by vstest.console.
+                //
+                // TODO(copoiena): Writing telemetry twice is less than ideal.
+                // We first write telemetry data in the _requestData variable in the ParallelRunEventsHandler
+                // and then we write again here. We should refactor this code and write only once.
+                runCompleteArgs.DiscoveredExtensions = TestExtensions.CreateMergedDictionary(
+                    runCompleteArgs.DiscoveredExtensions,
+                    TestPluginCache.Instance.TestExtensions.GetCachedExtensions());
+
+                if (_requestData.IsTelemetryOptedIn)
+                {
+                    TestExtensions.AddExtensionTelemetry(
+                        runCompleteArgs.Metrics,
+                        runCompleteArgs.DiscoveredExtensions);
+                }
+
                 // Ignore the time sent (runCompleteArgs.ElapsedTimeInRunningTests)
                 // by either engines - as both calculate at different points
                 // If we use them, it would be an incorrect comparison between TAEF and Rocksteady
@@ -420,7 +438,6 @@ public void HandleTestRunComplete(TestRunCompleteEventArgs runCompleteArgs!!, Te
                 // Notify the waiting handle that run is complete
                 _runCompletionEvent.Set();
 
-
                 var executionTotalTimeTaken = DateTime.UtcNow - _executionStartTime;
 
                 // Fill in the time taken to complete the run
@@ -583,6 +600,25 @@ private string UpdateRawMessageWithTelemetryInfo(TestRunCompletePayload testRunC
                 // Fill in the time taken to complete the run
                 var executionTotalTimeTakenForDesignMode = DateTime.UtcNow - _executionStartTime;
                 testRunCompletePayload.TestRunCompleteArgs.Metrics[TelemetryDataConstants.TimeTakenInSecForRun] = executionTotalTimeTakenForDesignMode.TotalSeconds;
+
+                // Add extensions discovered by vstest.console.
+                //
+                // TODO(copoiena):
+                // Doing extension merging here is incorrect because we can end up not merging the
+                // cached extensions for the current process (i.e. vstest.console) and hence have
+                // an incomplete list of discovered extensions. This can happen because this method
+                // is called only if telemetry is opted in (see: HandleRawMessage). We should handle
+                // this merge a level above in order to be consistent, but that means we'd have to
+                // deserialize all raw messages no matter if telemetry is opted in or not and that
+                // would probably mean a performance hit.
+                testRunCompletePayload.TestRunCompleteArgs.DiscoveredExtensions = TestExtensions.CreateMergedDictionary(
+                    testRunCompletePayload.TestRunCompleteArgs.DiscoveredExtensions,
+                    TestPluginCache.Instance.TestExtensions.GetCachedExtensions());
+
+                // Write extensions to telemetry data.
+                TestExtensions.AddExtensionTelemetry(
+                    testRunCompletePayload.TestRunCompleteArgs.Metrics,
+                    testRunCompletePayload.TestRunCompleteArgs.DiscoveredExtensions);
             }
 
             if (message is VersionedMessage message1)

src/Microsoft.TestPlatform.Common/ExtensionFramework/TestPluginCache.cs

@@ -117,7 +117,8 @@ public List<string> GetExtensionPaths(string endsWithPattern, bool skipDefaultEx
     /// <returns>
     /// The <see cref="Dictionary"/>. of test plugin info.
     /// </returns>
-    public Dictionary<string, TPluginInfo> DiscoverTestExtensions<TPluginInfo, TExtension>(string endsWithPattern)
+    public Dictionary<string, TPluginInfo> DiscoverTestExtensions<TPluginInfo, TExtension>(
+        string endsWithPattern)
         where TPluginInfo : TestPluginInformation
     {
         EqtTrace.Verbose("TestPluginCache.DiscoverTestExtensions: finding test extensions in assemblies ends with: {0} TPluginInfo: {1} TExtension: {2}", endsWithPattern, typeof(TPluginInfo), typeof(TExtension));
@@ -309,35 +310,37 @@ internal IEnumerable<string> DefaultExtensionPaths
     /// <returns>
     /// The <see cref="Dictionary"/>.
     /// </returns>
-    internal Dictionary<string, TPluginInfo> GetTestExtensions<TPluginInfo, TExtension>(string extensionAssembly, bool skipCache = false) where TPluginInfo : TestPluginInformation
+    internal Dictionary<string, TPluginInfo> GetTestExtensions<TPluginInfo, TExtension>(
+        string extensionAssembly,
+        bool skipCache = false)
+        where TPluginInfo : TestPluginInformation
     {
         if (skipCache)
         {
             return GetTestExtensions<TPluginInfo, TExtension>(new List<string>() { extensionAssembly });
         }
-        else
-        {
-            // Check if extensions from this assembly have already been discovered.
-            var extensions = TestExtensions?.GetExtensionsDiscoveredFromAssembly(
-                TestExtensions.GetTestExtensionCache<TPluginInfo>(),
-                extensionAssembly);
 
-            if (extensions != null && extensions.Count > 0)
-            {
-                return extensions;
-            }
+        // Check if extensions from this assembly have already been discovered.
+        var extensions = TestExtensions?.GetExtensionsDiscoveredFromAssembly(
+            TestExtensions.GetTestExtensionCache<TPluginInfo>(),
+            extensionAssembly);
 
-            var pluginInfos = GetTestExtensions<TPluginInfo, TExtension>(new List<string>() { extensionAssembly });
+        if (extensions?.Count > 0)
+        {
+            return extensions;
+        }
 
-            // Add extensions discovered to the cache.
-            if (TestExtensions == null)
-            {
-                TestExtensions = new TestExtensions();
-            }
+        var pluginInfos = GetTestExtensions<TPluginInfo, TExtension>(new List<string>() { extensionAssembly });
 
-            TestExtensions.AddExtension(pluginInfos);
-            return pluginInfos;
+        // Add extensions discovered to the cache.
+        if (TestExtensions == null)
+        {
+            TestExtensions = new TestExtensions();
         }
+
+        TestExtensions.AddExtension(pluginInfos);
+
+        return pluginInfos;
     }
 
     /// <summary>

src/Microsoft.TestPlatform.Common/ExtensionFramework/Utilities/TestExtensions.cs

@@ -5,11 +5,11 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
+using System.Text;
 
 using Microsoft.VisualStudio.TestPlatform.Common.DataCollector;
-
+using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel;
-
 using Microsoft.VisualStudio.TestPlatform.PlatformAbstractions;
 
 #nullable disable
@@ -91,6 +91,80 @@ public class TestExtensions
     /// </summary>
     internal bool AreDataCollectorsCached { get; set; }
 
+    /// <summary>
+    /// Merge two extension dictionaries.
+    /// </summary>
+    /// 
+    /// <param name="first">First extension dictionary.</param>
+    /// <param name="second">Second extension dictionary.</param>
+    /// 
+    /// <returns>
+    /// A dictionary representing the merger between the two input dictionaries.
+    /// </returns>
+    internal static Dictionary<string, HashSet<string>> CreateMergedDictionary(
+        Dictionary<string, HashSet<string>> first,
+        Dictionary<string, HashSet<string>> second)
+    {
+        var isFirstNullOrEmpty = first == null || first.Count == 0;
+        var isSecondNullOrEmpty = second == null || second.Count == 0;
+
+        // Sanity checks.
+        if (isFirstNullOrEmpty && isSecondNullOrEmpty)
+        {
+            return new Dictionary<string, HashSet<string>>();
+        }
+        if (isFirstNullOrEmpty)
+        {
+            return new Dictionary<string, HashSet<string>>(second);
+        }
+        if (isSecondNullOrEmpty)
+        {
+            return new Dictionary<string, HashSet<string>>(first);
+        }
+
+        // Copy all the keys in the first dictionary into the resulting dictionary.
+        var result = new Dictionary<string, HashSet<string>>(first);
+
+        foreach (var kvp in second)
+        {
+            // If the "source" set is empty there's no reason to continue merging for this key.
+            if (kvp.Value == null || kvp.Value.Count == 0)
+            {
+                continue;
+            }
+
+            // If there's no key-value pair entry in the "destination" dictionary for the current
+            // key in the "source" dictionary, we copy the "source" set wholesale.
+            if (!result.ContainsKey(kvp.Key))
+            {
+                result.Add(kvp.Key, kvp.Value);
+                continue;
+            }
+
+            // Getting here means there's already an entry for the "source" key in the "destination"
+            // dictionary which means we need to copy individual set elements from the "source" set
+            // to the "destination" set.
+            result[kvp.Key] = MergeSets(result[kvp.Key], kvp.Value);
+        }
+
+        return result;
+    }
+
+    /// <summary>
+    /// Add extension-related telemetry.
+    /// </summary>
+    /// 
+    /// <param name="metrics">A collection representing the telemetry data.</param>
+    /// <param name="extensions">The input extension collection.</param>
+    internal static void AddExtensionTelemetry(
+        IDictionary<string, object> metrics,
+        Dictionary<string, HashSet<string>> extensions)
+    {
+        metrics.Add(
+            TelemetryDataConstants.DiscoveredExtensions,
+            SerializeExtensionDictionary(extensions));
+    }
+
     /// <summary>
     /// Adds the extensions specified to the current set of extensions.
     /// </summary>
@@ -307,6 +381,27 @@ internal void SetTestExtensionsCacheStatusToTrue<TPluginInfo>() where TPluginInf
         }
     }
 
+    /// <summary>
+    /// Gets the cached extensions for the current process.
+    /// </summary>
+    /// 
+    /// <returns>A dictionary representing the cached extensions for the current process.</returns>
+    internal Dictionary<string, HashSet<string>> GetCachedExtensions()
+    {
+        var extensions = new Dictionary<string, HashSet<string>>();
+
+        // Write all "known" cached extension.
+        AddCachedExtensionToDictionary(extensions, "TestDiscoverers", TestDiscoverers?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "TestExecutors", TestExecutors?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "TestExecutors2", TestExecutors2?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "TestSettingsProviders", TestSettingsProviders?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "TestLoggers", TestLoggers?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "TestHosts", TestHosts?.Values.ToList());
+        AddCachedExtensionToDictionary(extensions, "DataCollectors", DataCollectors?.Values.ToList());
+
+        return extensions;
+    }
+
     /// <summary>
     /// The invalidate cache of plugin infos.
     /// </summary>
@@ -390,4 +485,45 @@ private void SetTestExtensionCache<TPluginInfo>(Dictionary<string, TPluginInfo>
         }
     }
 
+    private void AddCachedExtensionToDictionary<T>(
+        Dictionary<string, HashSet<string>> extensionDict,
+        string extensionType,
+        IEnumerable<T> extensions)
+        where T : TestPluginInformation
+    {
+        if (extensions == null)
+        {
+            return;
+        }
+
+        extensionDict.Add(extensionType, new HashSet<string>(extensions.Select(e => e.IdentifierData)));
+    }
+
+    private static string SerializeExtensionDictionary(IDictionary<string, HashSet<string>> extensions)
+    {
+        StringBuilder sb = new();
+
+        foreach (var kvp in extensions)
+        {
+            if (kvp.Value?.Count > 0)
+            {
+                sb.AppendFormat("{0}=[{1}];", kvp.Key, string.Join(",", kvp.Value));
+            }
+        }
+
+        return sb.ToString();
+    }
+
+    private static HashSet<string> MergeSets(HashSet<string> firstSet, HashSet<string> secondSet)
+    {
+        var mergedSet = new HashSet<string>(firstSet);
+
+        // No need to worry about duplicates as the set implementation handles this already.
+        foreach (var key in secondSet)
+        {
+            mergedSet.Add(key);
+        }
+
+        return mergedSet;
+    }
 }

src/Microsoft.TestPlatform.Common/Telemetry/TelemetryDataConstants.cs

@@ -10,6 +10,9 @@ namespace Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
 /// </summary>
 internal static class TelemetryDataConstants
 {
+    // ******************** General ***********************
+    public static readonly string DiscoveredExtensions = "VS.TestPlatform.DiscoveredExtensions";
+
     // ******************** Execution ***********************
     public static readonly string ParallelEnabledDuringExecution = "VS.TestRun.ParallelEnabled";
 

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

@@ -48,4 +48,9 @@ public class DiscoveryCompletePayload
     /// Gets or sets list of sources which were not discovered at all.
     /// </summary>
     public IList<string> NotDiscoveredSources { get; set; } = new List<string>();
+
+    /// <summary>
+    /// Gets or sets the collection of discovered extensions.
+    /// </summary>
+    public Dictionary<string, HashSet<string>> DiscoveredExtensions { get; set; } = new();
 }

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

@@ -6,3 +6,5 @@ Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.Discovery
 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
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.DiscoveredExtensions.get -> System.Collections.Generic.Dictionary<string, System.Collections.Generic.HashSet<string>>
+Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.ObjectModel.DiscoveryCompletePayload.DiscoveredExtensions.set -> void

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -576,7 +576,9 @@ private void OnDiscoveryMessageReceived(ITestDiscoveryEventsHandler2 discoveryEv
                         discoveryCompletePayload.IsAborted,
                         discoveryCompletePayload.FullyDiscoveredSources,
                         discoveryCompletePayload.PartiallyDiscoveredSources,
-                        discoveryCompletePayload.NotDiscoveredSources);
+                        discoveryCompletePayload.NotDiscoveredSources,
+                        discoveryCompletePayload.DiscoveredExtensions);
+
                     discoveryCompleteEventArgs.Metrics = discoveryCompletePayload.Metrics;
                     discoveryEventsHandler.HandleDiscoveryComplete(
                         discoveryCompleteEventArgs,