Reject method NGEN image if there's a rejit request for that method. (#…

…6184)

tracer/build/_build/Build.Steps.cs

@@ -1488,19 +1488,41 @@ _ when exclude.Contains(project.Path) => false,
                         .SetProjectFile(project)));
 
                 var projectsToPublish = includeIntegration
-                   .Select(x => Solution.GetProject(x))
-                   .Where(x => x.Name switch
+                   .Select(x =>
                     {
-                        "Samples.Trimming" => Framework.IsGreaterThanOrEqualTo(TargetFramework.NET6_0),
-                        _ => false,
+                        var project = Solution.GetProject(x);
+                        return project?.Name switch
+                        {
+                            "Samples.Trimming" => (project, include: Framework.IsGreaterThanOrEqualTo(TargetFramework.NET6_0), r2r: false),
+                            "Samples.ManualInstrumentation" => (project, include: Framework.IsGreaterThanOrEqualTo(TargetFramework.NETCOREAPP2_1), r2r: true),
+                            _ => (project, include: false, r2r: false),
+                        };
+                    })
+                   .Where(x => (x, x.project.TryGetTargetFrameworks(), x.project.RequiresDockerDependency()) switch
+                    {
+                        ({include: false }, _, _) => false,
+                        _ when exclude.Contains(x.project.Path) => false,
+                        _ when !string.IsNullOrWhiteSpace(SampleName) => x.project.Path.ToString().Contains(SampleName, StringComparison.OrdinalIgnoreCase),
+                        (_, _, DockerDependencyType.All) => false, // can't use docker on Windows
+                        (_, { } targets, _) => targets.Contains(Framework),
+                        _ => true,
                     });
 
-                var rid = IsArm64 ? "win-arm64" : "win-x64";
+                var rid = TargetPlatform.ToString() switch
+                {
+                    "x64" => "win-x64",
+                    "x86" => "win-x86",
+                    "ARM64" or "ARM64EC" => "win-arm64",
+                    _ => throw new InvalidOperationException("Unsupported architecture " + RuntimeInformation.ProcessArchitecture),
+                };
+
                 DotNetPublish(config => config
                    .SetConfiguration(BuildConfiguration)
                    .SetFramework(Framework)
                    .SetRuntime(rid)
-                   .CombineWith(projectsToPublish, (s, project) => s.SetProject(project)));
+                   .CombineWith(projectsToPublish, (s, project) => s
+                      .SetProject(project.project)
+                      .When(project.r2r, x => x.SetPublishReadyToRun(true))));
             }
         });
 
@@ -1912,20 +1934,25 @@ var name when multiPackageProjects.Contains(name) => false,
 
             // We have to explicitly publish the trimming sample separately (written so we can add to this later if needs be)
             var projectsToPublish = sampleProjects
-               .Select(x => Solution.GetProject(x))
-               .Where(x => x?.Name switch
+               .Select(x =>
                 {
-                    "Samples.Trimming" => x.TryGetTargetFrameworks().Contains(Framework),
-                    _ => false,
+                    var project = Solution.GetProject(x);
+                    return project?.Name switch
+                    {
+                        "Samples.Trimming" => (project, include: Framework.IsGreaterThanOrEqualTo(TargetFramework.NET6_0), r2r: false),
+                        "Samples.ManualInstrumentation" => (project, include: Framework.IsGreaterThanOrEqualTo(TargetFramework.NETCOREAPP2_1), r2r: true),
+                        _ => (project, include: false, r2r: false),
+                    };
                 })
                .Where(x => (IncludeTestsRequiringDocker, x) switch
                 {
                     // filter out or to integration tests that have docker dependencies
+                    (_, {include: false }) => false,
                     (null, _) => true,
-                    (_, null) => true,
-                    (_, { } p) when !string.IsNullOrWhiteSpace(SampleName) => p.Name.Contains(SampleName, StringComparison.OrdinalIgnoreCase),
-                    (false, { } p) => p.RequiresDockerDependency() == DockerDependencyType.None,
-                    (true, { } p) => p.RequiresDockerDependency() != DockerDependencyType.None,
+                    (_, { project: null}) => true,
+                    (_, { } p) when !string.IsNullOrWhiteSpace(SampleName) => p.project.Name.Contains(SampleName, StringComparison.OrdinalIgnoreCase),
+                    (false, { } p) => p.project.RequiresDockerDependency() == DockerDependencyType.None,
+                    (true, { } p) => p.project.RequiresDockerDependency() != DockerDependencyType.None,
                 });
 
             var rid = (IsLinux, IsArm64) switch
@@ -1939,7 +1966,9 @@ var name when multiPackageProjects.Contains(name) => false,
                .SetConfiguration(BuildConfiguration)
                .SetFramework(Framework)
                .SetRuntime(rid)
-               .CombineWith(projectsToPublish, (s, project) => s.SetProject(project)));
+               .CombineWith(projectsToPublish, (s, project) => s
+                  .SetProject(project.project)
+                  .When(project.r2r, x => x.SetPublishReadyToRun(true))));
         });
 
     Target CompileMultiApiPackageVersionSamples => _ => _

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

@@ -255,6 +255,8 @@ HRESULT STDMETHODCALLTYPE CorProfiler::Initialize(IUnknown* cor_profiler_info_un
     rejit_handler = info10 != nullptr
                         ? std::make_shared<RejitHandler>(info10, work_offloader)
                         : std::make_shared<RejitHandler>(this->info_, work_offloader);
+    rejit_handler->SetRejitTracking(runtime_information_.is_core());
+
     tracer_integration_preprocessor = std::make_unique<TracerRejitPreprocessor>(this, rejit_handler, work_offloader);
 
     fault_tolerant_method_duplicator = std::make_shared<fault_tolerant::FaultTolerantMethodDuplicator>(this, rejit_handler, work_offloader);
@@ -4333,7 +4335,44 @@ HRESULT STDMETHODCALLTYPE CorProfiler::JITCachedFunctionSearchStarted(FunctionID
         return S_OK;
     }
 
+    // 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
+    // are not resolved (is not clear why, maybe a bug). So we end up missing spans.
+    // Also in .NET Framework  we don't need to reject the ngen image, the rejit will always do the right job.
+    // In .NET Core if we don't reject the image, the image will be used and the rejit callback will never get called
+    // (this was confirmed on issue-6124).
+    // The following code handle both scenarios.
     *pbUseCachedFunction = true;
+    if (runtime_information_.is_core())
+    {
+        // Check if this method has been rejitted, if that's the case we don't accept the image
+        bool hasBeenRejitted = this->rejit_handler->HasBeenRejitted(module_id, function_token);
+        *pbUseCachedFunction = !hasBeenRejitted;
+
+        // 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)
+        {
+            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,
+                              ", 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,
+                              ", Function: ", HexStr(function_token),
+                              " previous value = ", *pbUseCachedFunction ? "true" : "false");
+            }          
+        }
+    }
+
     return S_OK;
 }
 

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

@@ -280,6 +280,15 @@ void RejitHandler::RequestRejit(std::vector<ModuleID>& modulesVector, std::vecto
         if (SUCCEEDED(hr))
         {
             Logger::Info("Request ReJIT done for ", modulesVector.size(), " methods");
+
+            if (enable_rejit_tracking)
+            {
+                WriteLock wlock(m_rejit_history_lock);
+                for (auto i = 0; i < modulesVector.size(); i++)
+                {
+                    m_rejit_history.push_back({modulesVector[i], modulesMethodDef[i]});
+                }
+            }
         }
         else
         {
@@ -493,4 +502,37 @@ void RejitHandler::AddNGenInlinerModule(ModuleID moduleId)
     }
 }
 
+void RejitHandler::SetRejitTracking(bool enabled) {
+    if (IsShutdownRequested())
+    {
+        return;
+    }
+
+    enable_rejit_tracking = enabled;
+}
+
+bool RejitHandler::HasBeenRejitted(ModuleID moduleId, mdMethodDef methodDef) {
+    if (IsShutdownRequested())
+    {
+        return false;
+    }
+
+    if (!enable_rejit_tracking)
+    {
+        return false;
+    }
+
+    ReadLock rlock(m_rejit_history_lock);
+    for (auto i = 0; i < m_rejit_history.size(); i++)
+    {
+        const auto mod_met_pair = m_rejit_history[i];
+        if (get<0>(mod_met_pair) == moduleId && get<1>(mod_met_pair) == methodDef)
+        {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 } // namespace trace

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

@@ -114,6 +114,9 @@ class RejitHandler
 
     std::vector<Rejitter*> m_rejitters;
 
+    Lock m_rejit_history_lock;
+    std::vector<std::tuple<ModuleID, mdMethodDef>> m_rejit_history;
+    bool enable_rejit_tracking = false;
 public:
     RejitHandler(ICorProfilerInfo7* pInfo, std::shared_ptr<RejitWorkOffloader> work_offloader);
     RejitHandler(ICorProfilerInfo10* pInfo, std::shared_ptr<RejitWorkOffloader> work_offloader);
@@ -143,6 +146,9 @@ class RejitHandler
     bool HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef);
     void RemoveModule(ModuleID moduleId);
     void AddNGenInlinerModule(ModuleID moduleId);
+
+    void SetRejitTracking(bool enabled);
+    bool HasBeenRejitted(ModuleID moduleId, mdMethodDef methodDef);
 };
 
 } // namespace trace

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/ManualInstrumentationTests.cs

@@ -3,7 +3,9 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
+using System;
 using System.Threading.Tasks;
+using Datadog.Trace.Configuration;
 using Datadog.Trace.TestHelpers;
 using FluentAssertions;
 using FluentAssertions.Execution;
@@ -23,22 +25,34 @@ public ManualInstrumentationTests(ITestOutputHelper output)
 
     [SkippableFact]
     [Trait("RunOnWindows", "True")]
-    public async Task ManualAndAutomatic()
-    {
-        SetEnvironmentVariable("AUTO_INSTRUMENT_ENABLED", "1");
-        const int expectedSpans = 36;
-        using var telemetry = this.ConfigureTelemetry();
-        using var agent = EnvironmentHelper.GetMockAgent();
-        using var assert = new AssertionScope();
-        using var process = await RunSampleAndWaitForExit(agent);
-
-        var spans = agent.WaitForSpans(expectedSpans);
-        spans.Should().HaveCount(expectedSpans);
+    public async Task ManualAndAutomatic() => await RunTest(usePublishWithRID: false);
 
-        var settings = VerifyHelper.GetSpanVerifierSettings();
-
-        await VerifyHelper.VerifySpans(spans, settings);
+#if NETFRAMEWORK
+    [SkippableFact]
+    [Trait("RunOnWindows", "True")]
+    public async Task NGenRunManualAndAutomatic()
+    {
+        SetEnvironmentVariable("READY2RUN_ENABLED", "1");
+        var sampleAppPath = EnvironmentHelper.GetSampleApplicationPath();
+        NgenHelper.InstallToNativeImageCache(Output, sampleAppPath);
+        try
+        {
+            await RunTest(usePublishWithRID: false);
+        }
+        finally
+        {
+            NgenHelper.UninstallFromNativeImageCache(Output, sampleAppPath);
+        }
     }
+#else
+    [SkippableFact]
+    [Trait("RunOnWindows", "True")]
+    public async Task ReadyToRunManualAndAutomatic()
+    {
+        SetEnvironmentVariable("READY2RUN_ENABLED", "1");
+        await RunTest(usePublishWithRID: true);
+    }
+#endif
 
     [SkippableFact]
     [Trait("RunOnWindows", "True")]
@@ -55,4 +69,23 @@ public async Task ManualOnly()
         var spans = agent.WaitForSpans(0, timeoutInMilliseconds: 500);
         spans.Should().BeEmpty();
     }
+
+    private async Task RunTest(bool usePublishWithRID = false)
+    {
+        SetEnvironmentVariable("AUTO_INSTRUMENT_ENABLED", "1");
+        const int expectedSpans = 37;
+        using var telemetry = this.ConfigureTelemetry();
+        using var agent = EnvironmentHelper.GetMockAgent();
+        using var assert = new AssertionScope();
+        using var process = await RunSampleAndWaitForExit(agent, usePublishWithRID: usePublishWithRID);
+
+        var spans = agent.WaitForSpans(expectedSpans);
+        spans.Should().HaveCount(expectedSpans);
+
+        var settings = VerifyHelper.GetSpanVerifierSettings();
+        settings.UseMethodName(nameof(ManualAndAutomatic)); // they should be identical, so share
+        settings.DisableRequireUniquePrefix();
+
+        await VerifyHelper.VerifySpans(spans, settings);
+    }
 }

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/SourceCodeIntegrationGitMetadataTests.cs

@@ -24,7 +24,7 @@ public SourceCodeIntegrationGitMetadataTests(ITestOutputHelper output)
     [Trait("RunOnWindows", "True")]
     public async Task ManualAndAutomatic()
     {
-        const int expectedSpans = 36;
+        const int expectedSpans = 37;
         using var telemetry = this.ConfigureTelemetry();
         using var agent = EnvironmentHelper.GetMockAgent();
         using var assert = new AssertionScope();

tracer/test/Datadog.Trace.TestHelpers/EnvironmentHelper.cs

@@ -487,12 +487,10 @@ public string GetSampleApplicationOutputDirectory(string packageVersion = "", st
 
             string GetPivot()
             {
-                var rid = (usePublishWithRID, IsAlpine()) switch
-                {
-                    (false, _) => string.Empty,
-                    (true, false) => $"_{EnvironmentTools.GetOS()}-{(EnvironmentTools.GetPlatform() == "Arm64" ? "arm64" : "x64")}",
-                    (true, true) => $"_{EnvironmentTools.GetOS()}-musl-{(EnvironmentTools.GetPlatform() == "Arm64" ? "arm64" : "x64")}",
-                };
+                var rid = usePublishWithRID
+                              ? $"_{EnvironmentTools.GetOS()}{(IsAlpine() ? "-musl" : string.Empty)}-{EnvironmentTools.GetPlatform().ToLowerInvariant()}"
+                              : string.Empty;
+
                 var config = EnvironmentTools.GetBuildConfiguration().ToLowerInvariant();
                 var packageVersionPivot = string.IsNullOrEmpty(packageVersion) ? string.Empty : $"_{packageVersion}";
                 return $"{config}_{targetFramework}{packageVersionPivot}{rid}";