Refactor the NGEN tracking fix to be more performant (#6598)

## Summary of changes

This PR refactor the NGEN tracking fix to remove the `lock` and the
`unordered_set` and reuse the `module_ids` lock and a vector + 3
mdMethodDef field.

## Reason for change

After #6588 was merged, @andrewlock saw some performance impact in the
Execution Benchmarks.

## Implementation details

Remove a lock and an unordered_set for just a field comparison + a
vector Contains if a field is found.

This should improve everything on the missed hit because is just a
comparison over 3 fields.

## Test coverage

## Other details
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

Author: tonyredondo

tracer/src/Datadog.Tracer.Native/clr_helpers.h

@@ -628,28 +628,6 @@ struct GenericTypeProps
     HRESULT TryParse();
 };
 
-struct ModuleIDMethodDef {
-    ModuleID ModuleId;
-    mdMethodDef MethodDef;
-
-    // Hash function
-    std::size_t operator()(const ModuleIDMethodDef& i) const {
-        return static_cast<std::size_t>(i.ModuleId) ^ (static_cast<std::size_t>(i.MethodDef) << 1);
-    }
-    // Comparison for equality
-    bool operator ()(const ModuleIDMethodDef& lhs, const ModuleIDMethodDef& rhs) const  {
-        return lhs.ModuleId == rhs.ModuleId && lhs.MethodDef == rhs.MethodDef;
-    }
-
-    bool operator ==(const ModuleIDMethodDef& other) const  {
-        return this->ModuleId == other.ModuleId && this->MethodDef == other.MethodDef;
-    }
-
-    bool operator !=(const ModuleIDMethodDef& other) const  {
-        return this->ModuleId != other.ModuleId || this->MethodDef != other.MethodDef;
-    }
-};
-
 RuntimeInformation GetRuntimeInformation(ICorProfilerInfo4* info);
 
 AssemblyInfo GetAssemblyInfo(ICorProfilerInfo4* info, const AssemblyID& assembly_id);

tracer/src/Datadog.Tracer.Native/cor_profiler.cpp

@@ -343,6 +343,7 @@ HRESULT STDMETHODCALLTYPE CorProfiler::Initialize(IUnknown* cor_profiler_info_un
 
     //
     managed_profiler_assembly_reference = AssemblyReference::GetFromCache(managed_profiler_full_assembly_version);
+    managedInternalModules_.reserve(10);
 
     const auto currentModuleFileName = shared::GetCurrentModuleFileName();
     if (currentModuleFileName == shared::EmptyWStr)
@@ -504,8 +505,7 @@ HRESULT STDMETHODCALLTYPE CorProfiler::AssemblyLoadFinished(AssemblyID assembly_
     return S_OK;
 }
 
-void CorProfiler::RewritingPInvokeMaps(const ModuleID module_id,
-                                       const ModuleMetadata& module_metadata,
+void CorProfiler::RewritingPInvokeMaps(const ModuleMetadata& module_metadata,
                                        const shared::WSTRING& rewrite_reason,
                                        const shared::WSTRING& nativemethods_type_name,
                                        const shared::WSTRING& library_path)
@@ -570,10 +570,6 @@ void CorProfiler::RewritingPInvokeMaps(const ModuleID module_id,
                                                              shared::WSTRING(importName).c_str(),
                                                              profiler_ref);
 
-                        // Store this methodDef token in the internal tokens list
-                        auto intTokens = internal_rewrite_tokens.Get();
-                        intTokens->insert({ module_id, methodDef });
-
                         if (FAILED(hr))
                         {
                             Logger::Warn("ModuleLoadFinished: DefinePinvokeMap to the actual profiler file path "
@@ -799,6 +795,7 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
                                   module_info.assembly.name == system_private_corelib_assemblyName))
     {
         corlib_module_loaded = true;
+        corlib_module_id = module_id;
         corlib_app_domain_id = app_domain_id;
 
         ComPtr<IUnknown> metadata_interfaces;
@@ -909,9 +906,6 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
         }
     }
 
-    // Add the module_id to the module metadata vector (to allow NGEN analysis later)
-    modules.push_back(module_id);
-
     if (module_info.assembly.name == managed_profiler_name)
     {
         // Fix PInvoke Rewriting
@@ -940,17 +934,17 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
         const auto& assemblyVersion = assemblyImport.version.str();
 
         Logger::Info("ModuleLoadFinished: ", managed_profiler_name, " v", assemblyVersion, " - Fix PInvoke maps");
-        managedProfilerModuleId_ = module_id;
+        managedInternalModules_.push_back(module_id);
 #ifdef _WIN32
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("windows"), windows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("ASM"), appsec_windows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("debugger"), debugger_windows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("fault_tolerant"), fault_tolerant_windows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("windows"), windows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("ASM"), appsec_windows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("debugger"), debugger_windows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("fault_tolerant"), fault_tolerant_windows_nativemethods_type);
 #else
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("non-windows"), nonwindows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("ASM"), appsec_nonwindows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("debugger"), debugger_nonwindows_nativemethods_type);
-        RewritingPInvokeMaps(module_id, module_metadata, WStr("fault_tolerant"), fault_tolerant_nonwindows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("non-windows"), nonwindows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("ASM"), appsec_nonwindows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("debugger"), debugger_nonwindows_nativemethods_type);
+        RewritingPInvokeMaps(module_metadata, WStr("fault_tolerant"), fault_tolerant_nonwindows_nativemethods_type);
 #endif // _WIN32
 
         mdTypeDef bubbleUpTypeDef;
@@ -964,15 +958,15 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
         if (fs::exists(native_loader_library_path))
         {
             auto native_loader_file_path = shared::ToWSTRING(native_loader_library_path);
-            RewritingPInvokeMaps(module_id, module_metadata, WStr("native loader"), native_loader_nativemethods_type, native_loader_file_path);
+            RewritingPInvokeMaps(module_metadata, WStr("native loader"), native_loader_nativemethods_type, native_loader_file_path);
         }
 
         if (ShouldRewriteProfilerMaps())
         {
             auto profiler_library_path = shared::GetEnvironmentValue(WStr("DD_INTERNAL_PROFILING_NATIVE_ENGINE_PATH"));
             if (!profiler_library_path.empty() && fs::exists(profiler_library_path))
             {
-                RewritingPInvokeMaps(module_id, module_metadata, WStr("continuous profiler"), profiler_nativemethods_type, profiler_library_path);
+                RewritingPInvokeMaps(module_metadata, WStr("continuous profiler"), profiler_nativemethods_type, profiler_library_path);
             }
         }
 
@@ -1042,11 +1036,14 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
             const auto& assemblyVersion = assemblyImport.version.str();
 
             Logger::Info("ModuleLoadFinished: ", manual_instrumentation_name, " v", assemblyVersion, " - RewriteIsManualInstrumentationOnly");
+            managedInternalModules_.push_back(module_id);
 
             // Rewrite Instrumentation.IsManualInstrumentationOnly()
             RewriteIsManualInstrumentationOnly(module_metadata, module_id);
         }
 
+        modules.push_back(module_id);
+
         bool searchForTraceAttribute = trace_annotations_enabled;
         if (searchForTraceAttribute)
         {
@@ -1761,7 +1758,6 @@ HRESULT STDMETHODCALLTYPE CorProfiler::AppDomainShutdownFinished(AppDomainID app
     const auto& count = first_jit_compilation_app_domains.erase(appDomainId);
 
     Logger::Debug("AppDomainShutdownFinished: AppDomain: ", appDomainId, ", removed ", count, " elements");
-
     return S_OK;
 }
 
@@ -2721,10 +2717,8 @@ HRESULT CorProfiler::RewriteForDistributedTracing(const ModuleMetadata& module_m
                                 module_metadata.metadata_import));
     }
 
-    // Store this methodDef token in the internal tokens list
-    auto intTokens = internal_rewrite_tokens.Get();
-    intTokens->insert({ module_id, getDistributedTraceMethodDef });
-
+    Logger::Debug("MethodDef was added as an internal rewrite: ", getDistributedTraceMethodDef);
+    getDistributedTraceMethodDef_ = getDistributedTraceMethodDef;
     return hr;
 }
 
@@ -2799,8 +2793,7 @@ HRESULT CorProfiler::RewriteForTelemetry(const ModuleMetadata& module_metadata,
 
     // Store this methodDef token in the internal tokens list
     Logger::Debug("MethodDef was added as an internal rewrite: ", getNativeTracerVersionMethodDef);
-    auto intTokens = internal_rewrite_tokens.Get();
-    intTokens->insert({ module_id, getNativeTracerVersionMethodDef });
+    getNativeTracerVersionMethodDef_ = getNativeTracerVersionMethodDef;
 
     return hr;
 }
@@ -2866,8 +2859,7 @@ HRESULT CorProfiler::RewriteIsManualInstrumentationOnly(const ModuleMetadata& mo
 
     // Store this methodDef token in the internal tokens list
     Logger::Debug("MethodDef was added as an internal rewrite: ", isAutoEnabledMethodDef);
-    auto intTokens = internal_rewrite_tokens.Get();
-    intTokens->insert({ module_id, isAutoEnabledMethodDef });
+    isManualInstrumentationOnlyMethodDef_ = isAutoEnabledMethodDef;
 
     return hr;
 }
@@ -3256,10 +3248,6 @@ HRESULT CorProfiler::RunILStartupHook(const ComPtr<IMetaDataEmit2>& metadata_emi
         return hr;
     }
 
-    // Store this methodDef token in the internal tokens list
-    auto intTokens = internal_rewrite_tokens.Get();
-    intTokens->insert({ module_id, function_token });
-
     return S_OK;
 }
 
@@ -4364,11 +4352,31 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
     }
 
     auto _ = trace::Stats::Instance()->JITCachedFunctionSearchStartedMeasure();
-    if (!pbUseCachedFunction)
+    if (pbUseCachedFunction == nullptr || !*pbUseCachedFunction)
     {
         return S_OK;
     }
 
+    // Extract Module metadata
+    ModuleID module_id;
+    mdToken function_token = mdTokenNil;
+
+    HRESULT hr = this->info_->GetFunctionInfo(functionId, nullptr, &module_id, &function_token);
+    if (FAILED(hr))
+    {
+        Logger::Warn("JITCachedFunctionSearchStarted: Call to ICorProfilerInfo.GetFunctionInfo() failed for ",
+                        functionId);
+        return S_OK;
+    }
+
+    // Verify if is the COR module
+    if (module_id == corlib_module_id)
+    {
+        // we don't rewrite the COR module, so we accept all the images from there.
+        *pbUseCachedFunction = true;
+        return S_OK;
+    }
+
     // keep this lock until we are done using the module,
     // to prevent it from unloading while in use
     auto modulesOpt = module_ids.TryGet();
@@ -4382,18 +4390,6 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
 
     auto& modules = modulesOpt.value();
 
-    // Extract Module metadata
-    ModuleID module_id;
-    mdToken function_token = mdTokenNil;
-
-    HRESULT hr = this->info_->GetFunctionInfo(functionId, nullptr, &module_id, &function_token);
-    if (FAILED(hr))
-    {
-        Logger::Warn("JITCachedFunctionSearchStarted: Call to ICorProfilerInfo4.GetFunctionInfo() failed for ",
-                        functionId);
-        return S_OK;
-    }
-
     // Call RequestRejitOrRevert for register inliners and current NGEN module.
     if (rejit_handler != nullptr)
     {
@@ -4410,30 +4406,50 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
         return S_OK;
     }
 
+    bool isAnInternalModule = false;
+
     // Verify that we have the metadata for this module
     if (!shared::Contains(modules.Ref(), module_id))
     {
-        // we haven't stored a ModuleMetadata for this module,
-        // so there's nothing to do here, we accept the NGEN image.
-        *pbUseCachedFunction = true;
-        return S_OK;
+        isAnInternalModule = shared::Contains(managedInternalModules_, module_id);
+        if (!isAnInternalModule)
+        {
+            // we haven't stored a ModuleMetadata for this module,
+            // so there's nothing to do here, we accept the NGEN image.
+            *pbUseCachedFunction = true;
+            return S_OK;
+        }
     }
 
-    const auto& module_info = GetModuleInfo(this->info_, module_id);
-    if (!module_info.IsValid())
+    // let's get the AssemblyID
+    DWORD module_path_len = 0;
+    AssemblyID assembly_id = 0;
+    hr = this->info_->GetModuleInfo(module_id, nullptr, 0, &module_path_len,
+                                            nullptr, &assembly_id);
+    if (FAILED(hr) || module_path_len == 0)
     {
-        Logger::Debug("Disabling NGEN. ModuleInfo is not valid. ModuleId=", module_id);
+        Logger::Warn("JITCachedFunctionSearchStarted: Call to ICorProfilerInfo.GetModuleInfo() failed for ",
+                        functionId);
         return S_OK;
     }
 
-    const auto& appDomainId = module_info.assembly.app_domain_id;
+    // now the assembly info
+    DWORD assembly_name_len = 0;
+    AppDomainID app_domain_id;
+    hr = this->info_->GetAssemblyInfo(assembly_id, 0, &assembly_name_len, nullptr, &app_domain_id, nullptr);
+    if (FAILED(hr) || assembly_name_len == 0)
+    {
+        Logger::Warn("JITCachedFunctionSearchStarted: Call to ICorProfilerInfo.GetAssemblyInfo() failed for ",
+                        functionId);
+        return S_OK;
+    }
 
     const bool has_loader_injected_in_appdomain =
-        first_jit_compilation_app_domains.find(appDomainId) != first_jit_compilation_app_domains.end();
+        first_jit_compilation_app_domains.find(app_domain_id) != first_jit_compilation_app_domains.end();
 
     if (!has_loader_injected_in_appdomain)
     {
-        Logger::Debug("Disabling NGEN due to missing loader.");
+        Logger::Debug("JITCachedFunctionSearchStarted: Disabling NGEN due to missing loader.");
         // The loader is missing in this AppDomain, we skip the NGEN image to allow the JITCompilationStart inject
         // it.
         *pbUseCachedFunction = false;
@@ -4442,37 +4458,32 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
 
     // Let's check if the method has been rewritten internally
     // if that's the case we don't accept the image
-    auto internalTokensOpt = internal_rewrite_tokens.TryGet();
-    if (!internalTokensOpt.has_value())
-    {
-        Logger::Error(
-            "JITCachedFunctionSearchStarted: Failed on exception while tried to acquire the lock for the internal_rewrite_tokens collection for functionId ",
-            functionId);
-        return S_OK;
-    }
-
-    auto& internalTokens = internalTokensOpt.value();
-    bool hasBeenRewritten = internalTokens->find({ module_id, function_token }) != internalTokens->end();
+    bool isKnownMethodDef =
+        function_token == getDistributedTraceMethodDef_ ||
+            function_token == getNativeTracerVersionMethodDef_ ||
+                function_token == isManualInstrumentationOnlyMethodDef_;
+    bool hasBeenRewritten = isKnownMethodDef && isAnInternalModule;
     if (hasBeenRewritten)
     {
         // If we are in debug mode and the image is rejected because has been rewritten then let's write a couple of logs
         if (Logger::IsDebugEnabled())
         {
+            const auto& module_info = GetModuleInfo(this->info_, module_id);
             ComPtr<IUnknown> metadata_interfaces;
             if (this->info_->GetModuleMetaData(module_id, ofRead, IID_IMetaDataImport2,
                                                metadata_interfaces.GetAddressOf()) == S_OK)
             {
                 const auto& metadata_import = metadata_interfaces.As<IMetaDataImport2>(IID_IMetaDataImport);
                 auto functionInfo = GetFunctionInfo(metadata_import, function_token);
 
-                Logger::Debug("NGEN Image: Rejected (because rewritten) for Module: ", module_info.assembly.name,
+                Logger::Debug("JITCachedFunctionSearchStarted: Rejected (because rewritten) for Module: ", module_info.assembly.name,
                               ", Method:", functionInfo.type.name, ".", functionInfo.name,
                               "() previous value =  ", *pbUseCachedFunction ? "true" : "false", "[moduleId=", module_id,
                               ", methodDef=", HexStr(function_token), "]");
             }
             else
             {
-                Logger::Debug("NGEN Image: Rejected (because rewritten) for Module: ", module_info.assembly.name,
+                Logger::Debug("JITCachedFunctionSearchStarted: Rejected (because rewritten) for Module: ", module_info.assembly.name,
                               ", Function: ", HexStr(function_token),
                               " previous value = ", *pbUseCachedFunction ? "true" : "false");
             }
@@ -4482,10 +4493,6 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
         *pbUseCachedFunction = false;
         return S_OK;
     }
-    else
-    {
-        Logger::Debug("JITCachedFunctionSearchStarted: non rejected by internal token on token: ", function_token, " Looking into a hashset with ", internalTokens->size(), " elements");
-    }
 
     // JITCachedFunctionSearchStarted has a different behaviour between .NET Framework and .NET Core
     // On .NET Framework when we reject bcl images the rejit calls of the integrations for those bcl assemblies
@@ -4504,21 +4511,22 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
         // If we are in debug mode and the image is rejected because has been rejitted then let's write a couple of logs
         if (Logger::IsDebugEnabled() && hasBeenRejitted)
         {
+            const auto& module_info = GetModuleInfo(this->info_, module_id);
             ComPtr<IUnknown> metadata_interfaces;
             if (this->info_->GetModuleMetaData(module_id, ofRead, IID_IMetaDataImport2,
                                                metadata_interfaces.GetAddressOf()) == S_OK)
             {
                 const auto& metadata_import = metadata_interfaces.As<IMetaDataImport2>(IID_IMetaDataImport);
                 auto functionInfo = GetFunctionInfo(metadata_import, function_token);
 
-                Logger::Debug("NGEN Image: Rejected (because rejitted) for Module: ", module_info.assembly.name,
+                Logger::Debug("JITCachedFunctionSearchStarted: Rejected (because rejitted) for Module: ", module_info.assembly.name,
                               ", Method:", functionInfo.type.name, ".", functionInfo.name,
                               "() previous value =  ", *pbUseCachedFunction ? "true" : "false", "[moduleId=", module_id,
                               ", methodDef=", HexStr(function_token), "]");
             }
             else
             {
-                Logger::Debug("NGEN Image: Rejected (because rejitted) for Module: ", module_info.assembly.name,
+                Logger::Debug("JITCachedFunctionSearchStarted: Rejected (because rejitted) for Module: ", module_info.assembly.name,
                               ", Function: ", HexStr(function_token),
                               " previous value = ", *pbUseCachedFunction ? "true" : "false");
             }