From c0b008e93b8332fe5acd9bfd10218707aa6dae1f Mon Sep 17 00:00:00 2001
From: Kevin Gosse <kevin.gosse@datadoghq.com>
Date: Mon, 1 Jul 2024 13:05:46 +0200
Subject: [PATCH] Lock access to rejitters

---
 tracer/src/Datadog.Tracer.Native/rejit_handler.cpp | 12 ++++++++++++
 tracer/src/Datadog.Tracer.Native/rejit_handler.h   |  1 +
 2 files changed, 13 insertions(+)

diff --git a/tracer/src/Datadog.Tracer.Native/rejit_handler.cpp b/tracer/src/Datadog.Tracer.Native/rejit_handler.cpp
index 864e6cba8c6a..2594db9e6115 100644
--- a/tracer/src/Datadog.Tracer.Native/rejit_handler.cpp
+++ b/tracer/src/Datadog.Tracer.Native/rejit_handler.cpp
@@ -356,6 +356,8 @@ void RejitHandler::Shutdown()
     WriteLock w_lock(m_shutdown_lock);
     m_shutdown.store(true);
 
+    ReadLock r_lock(m_rejitters_lock);
+
     for (auto rejitter : m_rejitters)
     {
         rejitter->Shutdown();
@@ -373,6 +375,8 @@ bool RejitHandler::IsShutdownRequested()
 
 void RejitHandler::RegisterRejitter(Rejitter* rejitter)
 {
+    WriteLock w_lock(m_rejitters_lock);
+
     if (m_rejitters.size() == 0)
     {
         m_rejitters.push_back(rejitter);
@@ -405,6 +409,8 @@ HRESULT RejitHandler::NotifyReJITParameters(ModuleID moduleId, mdMethodDef metho
     // Create the FunctionControlWrapper
     FunctionControlWrapper functionControl((ICorProfilerInfo*)m_profilerInfo, moduleId, methodId);
 
+    ReadLock r_lock(m_rejitters_lock);
+
     // Call all rejitters sequentially
     for (auto rejitter : m_rejitters)
     {
@@ -461,6 +467,8 @@ bool RejitHandler::HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef)
         return false;
     }
 
+    ReadLock r_lock(m_rejitters_lock);
+
     for (auto rejitter : m_rejitters)
     {
         if (rejitter->HasModuleAndMethod(moduleId, methodDef))
@@ -479,6 +487,8 @@ void RejitHandler::RemoveModule(ModuleID moduleId)
         return;
     }
 
+    ReadLock r_lock(m_rejitters_lock);
+
     for (auto rejitter : m_rejitters)
     {
         rejitter->RemoveModule(moduleId);
@@ -492,6 +502,8 @@ void RejitHandler::AddNGenInlinerModule(ModuleID moduleId)
         return;
     }
 
+    ReadLock r_lock(m_rejitters_lock);
+
     for (auto rejitter : m_rejitters)
     {
         rejitter->AddNGenInlinerModule(moduleId);
diff --git a/tracer/src/Datadog.Tracer.Native/rejit_handler.h b/tracer/src/Datadog.Tracer.Native/rejit_handler.h
index 95073c7069d2..72d1dab49f83 100644
--- a/tracer/src/Datadog.Tracer.Native/rejit_handler.h
+++ b/tracer/src/Datadog.Tracer.Native/rejit_handler.h
@@ -101,6 +101,7 @@ class RejitHandler
 private:
     std::atomic_bool m_shutdown = {false};
     Lock m_shutdown_lock;
+    Lock m_rejitters_lock;
 
     AssemblyProperty* m_pCorAssemblyProperty = nullptr;