Lock access to rejitters

DataDog · Jul 1, 2024 · c0b008e · c0b008e
1 parent 3d66935
commit c0b008e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git 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
@@ -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;