Lower reflection overhead (#2839)

Co-authored-by: Amaury Levé <amauryleve@microsoft.com>

src/Adapter/MSTest.TestAdapter/BannedSymbols.txt

@@ -0,0 +1 @@
+M:Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers.ReflectHelper.#ctor; This is allowed only for tests.

src/Adapter/MSTest.TestAdapter/Discovery/AssemblyEnumerator.cs

@@ -189,14 +189,15 @@ internal static string GetLoadExceptionDetails(ReflectionTypeLoadException ex)
     /// <param name="assemblyFileName">The reflected assembly name.</param>
     /// <param name="discoverInternals">True to discover test classes which are declared internal in
     /// addition to test classes which are declared public.</param>
+    /// <param name="discoveryOption"><see cref="TestDataSourceDiscoveryOption"/> to use when generating tests.</param>
     /// <param name="testIdGenerationStrategy"><see cref="TestIdGenerationStrategy"/> to use when generating TestId.</param>
     /// <returns>a TypeEnumerator instance.</returns>
-    internal virtual TypeEnumerator GetTypeEnumerator(Type type, string assemblyFileName, bool discoverInternals, TestIdGenerationStrategy testIdGenerationStrategy)
+    internal virtual TypeEnumerator GetTypeEnumerator(Type type, string assemblyFileName, bool discoverInternals, TestDataSourceDiscoveryOption discoveryOption, TestIdGenerationStrategy testIdGenerationStrategy)
     {
         var typeValidator = new TypeValidator(ReflectHelper, discoverInternals);
         var testMethodValidator = new TestMethodValidator(ReflectHelper, discoverInternals);
 
-        return new TypeEnumerator(type, assemblyFileName, ReflectHelper, typeValidator, testMethodValidator, testIdGenerationStrategy);
+        return new TypeEnumerator(type, assemblyFileName, ReflectHelper, typeValidator, testMethodValidator, discoveryOption, testIdGenerationStrategy);
     }
 
     private List<UnitTestElement> DiscoverTestsInType(
@@ -220,9 +221,8 @@ internal virtual TypeEnumerator GetTypeEnumerator(Type type, string assemblyFile
         try
         {
             typeFullName = type.FullName;
-            TypeEnumerator testTypeEnumerator = GetTypeEnumerator(type, assemblyFileName, discoverInternals, testIdGenerationStrategy);
+            TypeEnumerator testTypeEnumerator = GetTypeEnumerator(type, assemblyFileName, discoverInternals, discoveryOption, testIdGenerationStrategy);
             ICollection<UnitTestElement>? unitTestCases = testTypeEnumerator.Enumerate(out ICollection<string> warningsFromTypeEnumerator);
-
             warningMessages.AddRange(warningsFromTypeEnumerator);
 
             if (unitTestCases != null)
@@ -262,6 +262,18 @@ private bool DynamicDataAttached(IDictionary<string, object?> sourceLevelParamet
             return false;
         }
 
+        if (test.TestMethod.DataType == DynamicDataType.None)
+        {
+            return false;
+        }
+
+        // PERF: For perf we started setting DataType in TypeEnumerator, so when it is None we will not reach this line.
+        // But if we do run this code, we still reset it to None, because the code that determines if this is data drive test expects the value to be None
+        // and only sets it when needed.
+        //
+        // If you remove this line and acceptance tests still pass you are okay.
+        test.TestMethod.DataType = DynamicDataType.None;
+
         // NOTE: From this place we don't have any path that would let the user write a message on the TestContext and we don't do
         // anything with what would be printed anyway so we can simply use a simple StringWriter.
         using var writer = new StringWriter();
@@ -274,11 +286,11 @@ private bool DynamicDataAttached(IDictionary<string, object?> sourceLevelParamet
     private static bool TryProcessTestDataSourceTests(UnitTestElement test, TestMethodInfo testMethodInfo, List<UnitTestElement> tests)
     {
         MethodInfo methodInfo = testMethodInfo.MethodInfo;
-        IEnumerable<FrameworkITestDataSource>? testDataSources = ReflectHelper.GetAttributes<Attribute>(methodInfo, false)?.OfType<FrameworkITestDataSource>();
-        if (testDataSources == null || !testDataSources.Any())
-        {
-            return false;
-        }
+
+        // We don't have a special method to filter attributes that are not derived from Attribute, so we take all
+        // attributes and filter them. We don't have to care if there is one, because this method is only entered when
+        // there is at least one (we determine this in TypeEnumerator.GetTestFromMethod.
+        IEnumerable<FrameworkITestDataSource>? testDataSources = ReflectHelper.Instance.GetDerivedAttributes<Attribute>(methodInfo, inherit: false).OfType<FrameworkITestDataSource>();
 
         try
         {
@@ -310,9 +322,11 @@ private static bool TryProcessTestDataSourceTests(UnitTestElement test, TestMeth
                     throw new ArgumentException(string.Format(CultureInfo.CurrentCulture, FrameworkMessages.DynamicDataIEnumerableEmpty, "GetData", dataSource.GetType().Name));
                 }
             }
-            catch (Exception ex) when (ex is ArgumentException && MSTestSettings.CurrentSettings.ConsiderEmptyDataSourceAsInconclusive)
+            catch (ArgumentException) when (MSTestSettings.CurrentSettings.ConsiderEmptyDataSourceAsInconclusive)
             {
                 UnitTestElement discoveredTest = test.Clone();
+                // Make the test not data driven, because it had no data.
+                discoveredTest.TestMethod.DataType = DynamicDataType.None;
                 discoveredTest.DisplayName = dataSource.GetDisplayName(methodInfo, null) ?? discoveredTest.DisplayName;
                 tests.Add(discoveredTest);
                 continue;

src/Adapter/MSTest.TestAdapter/Discovery/TestMethodValidator.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Globalization;
@@ -48,8 +48,14 @@ internal TestMethodValidator(ReflectHelper reflectHelper, bool discoverInternals
     /// <returns> Return true if a method is a valid test method. </returns>
     internal virtual bool IsValidTestMethod(MethodInfo testMethodInfo, Type type, ICollection<string> warnings)
     {
-        if (!_reflectHelper.IsAttributeDefined<TestMethodAttribute>(testMethodInfo, false)
-            && !_reflectHelper.HasAttributeDerivedFrom<TestMethodAttribute>(testMethodInfo, false))
+        // PERF: We are doing caching reflection here, meaning we will cache every method info in the
+        // assembly, this is because when we discover and run we will repeatedly inspect all the methods in the assembly, and this
+        // gives us a better performance.
+        // It would be possible to use non-caching reflection here if we knew that we are only doing discovery that won't be followed by run,
+        // but the difference is quite small, and we don't expect a huge amount of non-test methods in the assembly.
+        //
+        // Also skip all methods coming from object, because they cannot be tests.
+        if (testMethodInfo.DeclaringType == typeof(object) || !_reflectHelper.IsDerivedAttributeDefined<TestMethodAttribute>(testMethodInfo, inherit: false))
         {
             return false;
         }

src/Adapter/MSTest.TestAdapter/Discovery/TypeEnumerator.cs

@@ -22,6 +22,7 @@ internal class TypeEnumerator
     private readonly TypeValidator _typeValidator;
     private readonly TestMethodValidator _testMethodValidator;
     private readonly TestIdGenerationStrategy _testIdGenerationStrategy;
+    private readonly TestDataSourceDiscoveryOption _discoveryOption;
     private readonly ReflectHelper _reflectHelper;
 
     /// <summary>
@@ -33,14 +34,15 @@ internal class TypeEnumerator
     /// <param name="typeValidator"> The validator for test classes. </param>
     /// <param name="testMethodValidator"> The validator for test methods. </param>
     /// <param name="testIdGenerationStrategy"><see cref="TestIdGenerationStrategy"/> to use when generating TestId.</param>
-    internal TypeEnumerator(Type type, string assemblyFilePath, ReflectHelper reflectHelper, TypeValidator typeValidator, TestMethodValidator testMethodValidator, TestIdGenerationStrategy testIdGenerationStrategy)
+    internal TypeEnumerator(Type type, string assemblyFilePath, ReflectHelper reflectHelper, TypeValidator typeValidator, TestMethodValidator testMethodValidator, TestDataSourceDiscoveryOption discoveryOption, TestIdGenerationStrategy testIdGenerationStrategy)
     {
         _type = type;
         _assemblyFilePath = assemblyFilePath;
         _reflectHelper = reflectHelper;
         _typeValidator = typeValidator;
         _testMethodValidator = testMethodValidator;
         _testIdGenerationStrategy = testIdGenerationStrategy;
+        _discoveryOption = discoveryOption;
     }
 
     /// <summary>
@@ -73,6 +75,8 @@ internal Collection<UnitTestElement> GetTests(ICollection<string> warnings)
         var tests = new Collection<UnitTestElement>();
 
         // Test class is already valid. Verify methods.
+        // PERF: GetRuntimeMethods is used here to get all methods, including non-public, and static methods.
+        // if we rely on analyzers to identify all invalid methods on build, we can change this to fit the current settings.
         foreach (MethodInfo method in _type.GetRuntimeMethods())
         {
             bool isMethodDeclaredInTestTypeAssembly = _reflectHelper.IsMethodDeclaredInSameAssemblyAsType(method, _type);
@@ -147,14 +151,30 @@ internal UnitTestElement GetTestFromMethod(MethodInfo method, bool isDeclaredInT
                     method.DeclaringType.Assembly);
         }
 
+        // PERF: When discovery option is set to DuringDiscovery, we will expand data on tests to one test case
+        // per data item. This will happen in AssemblyEnumerator. But AssemblyEnumerator does not have direct access to
+        // the method info or method attributes, so it would create a TestMethodInfo to see if the test is data driven.
+        // Creating TestMethodInfo is expensive and should be done only for a test that we know is data driven.
+        //
+        // So to optimize this we check if we have some data source attribute. Because here we have access to all attributes
+        // and we store that info in DataType. AssemblyEnumerator will pick this up and will get the real test data in the expensive way
+        // or it will skip over getting the data cheaply, when DataType = DynamicDataType.None.
+        //
+        // This needs to be done only when DuringDiscovery is set, because otherwise we would populate the DataType, but we would not populate
+        // and execution would not try to re-populate the data, because DataType is already set to data driven, so it would just throw error about empty data.
+        if (_discoveryOption == TestDataSourceDiscoveryOption.DuringDiscovery)
+        {
+            testMethod.DataType = GetDynamicDataType(method);
+        }
+
         var testElement = new UnitTestElement(testMethod)
         {
             // Get compiler generated type name for async test method (either void returning or task returning).
             AsyncTypeName = method.GetAsyncTypeName(),
-            TestCategory = _reflectHelper.GetCategories(method, _type),
+            TestCategory = _reflectHelper.GetTestCategories(method, _type),
             DoNotParallelize = _reflectHelper.IsDoNotParallelizeSet(method, _type),
             Priority = _reflectHelper.GetPriority(method),
-            Ignored = _reflectHelper.IsAttributeDefined<IgnoreAttribute>(method, false),
+            Ignored = _reflectHelper.IsNonDerivedAttributeDefined<IgnoreAttribute>(method, inherit: false),
             DeploymentItems = PlatformServiceProvider.Instance.TestDeployment.GetDeploymentItems(method, _type, warnings),
         };
 
@@ -174,31 +194,49 @@ internal UnitTestElement GetTestFromMethod(MethodInfo method, bool isDeclaredInT
 
         testElement.Traits = traits.ToArray();
 
-        if (_reflectHelper.GetCustomAttribute<CssIterationAttribute>(method) is CssIterationAttribute cssIteration)
+        if (_reflectHelper.GetFirstDerivedAttributeOrDefault<CssIterationAttribute>(method, inherit: true) is CssIterationAttribute cssIteration)
         {
             testElement.CssIteration = cssIteration.CssIteration;
         }
 
-        if (_reflectHelper.GetCustomAttribute<CssProjectStructureAttribute>(method) is CssProjectStructureAttribute cssProjectStructure)
+        if (_reflectHelper.GetFirstDerivedAttributeOrDefault<CssProjectStructureAttribute>(method, inherit: true) is CssProjectStructureAttribute cssProjectStructure)
         {
             testElement.CssProjectStructure = cssProjectStructure.CssProjectStructure;
         }
 
-        if (_reflectHelper.GetCustomAttribute<DescriptionAttribute>(method) is DescriptionAttribute descriptionAttribute)
+        if (_reflectHelper.GetFirstDerivedAttributeOrDefault<DescriptionAttribute>(method, inherit: true) is DescriptionAttribute descriptionAttribute)
         {
             testElement.Description = descriptionAttribute.Description;
         }
 
-        WorkItemAttribute[] workItemAttributes = _reflectHelper.GetCustomAttributes<WorkItemAttribute>(method);
+        WorkItemAttribute[] workItemAttributes = _reflectHelper.GetDerivedAttributes<WorkItemAttribute>(method, inherit: true).ToArray();
         if (workItemAttributes.Length != 0)
         {
             testElement.WorkItemIds = workItemAttributes.Select(x => x.Id.ToString(CultureInfo.InvariantCulture)).ToArray();
         }
 
         // get DisplayName from TestMethodAttribute (or any inherited attribute)
-        TestMethodAttribute? testMethodAttribute = _reflectHelper.GetCustomAttribute<TestMethodAttribute>(method);
+        TestMethodAttribute? testMethodAttribute = _reflectHelper.GetFirstDerivedAttributeOrDefault<TestMethodAttribute>(method, inherit: true);
         testElement.DisplayName = testMethodAttribute?.DisplayName ?? method.Name;
 
         return testElement;
     }
+
+    private DynamicDataType GetDynamicDataType(MethodInfo method)
+    {
+        foreach (Attribute attribute in _reflectHelper.GetDerivedAttributes<Attribute>(method, inherit: true))
+        {
+            if (AttributeComparer.IsDerived<ITestDataSource>(attribute))
+            {
+                return DynamicDataType.ITestDataSource;
+            }
+
+            if (AttributeComparer.IsDerived<DataSourceAttribute>(attribute))
+            {
+                return DynamicDataType.DataSourceAttribute;
+            }
+        }
+
+        return DynamicDataType.None;
+    }
 }

src/Adapter/MSTest.TestAdapter/Discovery/TypeValidator.cs

@@ -51,9 +51,12 @@ internal virtual bool IsValidTestClass(Type type, ICollection<string> warnings)
     {
         TypeInfo typeInfo = type.GetTypeInfo();
 
-        if (!typeInfo.IsClass
-            || (!_reflectHelper.IsAttributeDefined<TestClassAttribute>(typeInfo, false)
-            && !_reflectHelper.HasAttributeDerivedFrom<TestClassAttribute>(typeInfo, false)))
+        // PERF: We are doing caching reflection here, meaning we will cache every class info in the
+        // assembly, this is because when we discover and run we will repeatedly inspect all the types in the assembly, and this
+        // gives us a better performance.
+        // It would be possible to use non-caching reflection here if we knew that we are only doing discovery that won't be followed by run,
+        // but the difference is quite small, and we don't expect a huge amount of non-test classes in the assembly.
+        if (!typeInfo.IsClass || !_reflectHelper.IsDerivedAttributeDefined<TestClassAttribute>(typeInfo, inherit: false))
         {
             return false;
         }

src/Adapter/MSTest.TestAdapter/Execution/TestMethodInfo.cs

@@ -83,12 +83,11 @@ public class TestMethodInfo : ITestMethod
     /// </summary>
     internal TestMethodOptions TestMethodOptions { get; }
 
-    public Attribute[]? GetAllAttributes(bool inherit) => ReflectHelper.GetCustomAttributes(TestMethod, inherit) as Attribute[];
+    public Attribute[]? GetAllAttributes(bool inherit) => ReflectHelper.Instance.GetDerivedAttributes<Attribute>(TestMethod, inherit).ToArray();
 
     public TAttributeType[] GetAttributes<TAttributeType>(bool inherit)
         where TAttributeType : Attribute
-        => ReflectHelper.GetAttributes<TAttributeType>(TestMethod, inherit)
-        ?? [];
+        => ReflectHelper.Instance.GetDerivedAttributes<TAttributeType>(TestMethod, inherit).ToArray();
 
     /// <summary>
     /// Execute test method. Capture failures, handle async and return result.