DataDog · gleocadie · Sep 7, 2026 · Sep 7, 2026 · Sep 9, 2026
@@ -552,6 +552,8 @@ HRESULT STDMETHODCALLTYPE CorProfiler::ModuleLoadFinished(ModuleID module_id, HR
         return S_OK;
     }
 
+    rejit_handler->NotifyModuleLoaded(module_id);
+
     auto hr = TryRejitModule(module_id, modules.Ref());
 
     // Push integration definitions from past modules that were unable to be added
@@ -601,6 +603,13 @@ HRESULT STDMETHODCALLTYPE CorProfiler::ModuleLoadFinished(ModuleID module_id, HR
                 Logger::Warn("Timeout while waiting for the rejit requests to be processed. Rejit will continue asynchronously, but some initial calls may not be instrumented");
             }
         }
+
+        // The enqueued request holds its own token now, so drop the ones taken when the modules
+        // were queued above.
+        for (size_t i = 0; i < rejitModuleIds.size(); i++)
+        {
+            rejit_handler->ReleaseInFlightRequest();
+        }
     }
 
     if (debugger_instrumentation_requester != nullptr)
@@ -1163,6 +1172,9 @@ HRESULT CorProfiler::TryRejitModule(ModuleID module_id, std::vector<ModuleID>& m
 
                 if (methodReferences.size() > 0)
                 {
+                    // The module can unload while it waits here, so keep its unload recorded until
+                    // we dequeue it.
+                    rejit_handler->AcquireInFlightRequest();
                     rejit_module_method_pairs.push_back(std::make_pair(module_id, methodReferences));
                 }
             }

@@ -821,6 +821,15 @@ bool Dataflow::HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef)
 {
     return false;
 }
+void Dataflow::NotifyModuleLoaded(ModuleID moduleId)
+{
+}
+void Dataflow::AcquireInFlightRequest()
+{
+}
+void Dataflow::ReleaseInFlightRequest()
+{
+}
 void Dataflow::RemoveModule(ModuleID moduleId)
 {
 }

@@ -116,6 +116,9 @@ namespace iast
         void Shutdown() override;
         RejitHandlerModule* GetOrAddModule(ModuleID moduleId) override;
         bool HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef) override;
+        void NotifyModuleLoaded(ModuleID moduleId) override;
+        void AcquireInFlightRequest() override;
+        void ReleaseInFlightRequest() override;
         void RemoveModule(ModuleID moduleId) override;
         void AddNGenInlinerModule(ModuleID moduleId) override;
 

@@ -495,6 +495,52 @@ bool RejitHandler::HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef)
     return false;
 }
 
+void RejitHandler::NotifyModuleLoaded(ModuleID moduleId)
+{
+    if (IsShutdownRequested())
+    {
+        return;
+    }
+
+    Rejitter* prev = nullptr;
+    for (size_t x = 0; x < m_rejittersCount; x++)
+    {
+        const auto current = m_rejitters[x];
+        if (current != prev)
+        {
+            current->NotifyModuleLoaded(moduleId);
+        }
+    }
+}
+
+// No shutdown check here: it must stay paired with ReleaseInFlightRequest, otherwise a release
+// would decrement a count that was never taken.
+void RejitHandler::AcquireInFlightRequest()
+{
+    Rejitter* prev = nullptr;
+    for (size_t x = 0; x < m_rejittersCount; x++)
+    {
+        const auto current = m_rejitters[x];
+        if (current != prev)
+        {
+            current->AcquireInFlightRequest();
+        }
+    }
+}
+
+void RejitHandler::ReleaseInFlightRequest()
+{
+    Rejitter* prev = nullptr;
+    for (size_t x = 0; x < m_rejittersCount; x++)
+    {
+        const auto current = m_rejitters[x];
+        if (current != prev)
+        {
+            current->ReleaseInFlightRequest();
+        }
+    }
+}
+
 void RejitHandler::RemoveModule(ModuleID moduleId)
 {
     if (IsShutdownRequested())

@@ -146,6 +146,9 @@ class RejitHandler
     AssemblyProperty* GetCorAssemblyProperty();
 
     bool HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef);
+    void NotifyModuleLoaded(ModuleID moduleId);
+    void AcquireInFlightRequest();
+    void ReleaseInFlightRequest();
     void RemoveModule(ModuleID moduleId);
     void AddNGenInlinerModule(ModuleID moduleId);
 

@@ -8,6 +8,10 @@
 
 namespace trace
 {
+// Upper bound on RejitPreprocessor::m_unloaded_modules.
+// Handle cases where modules keep on being unloaded and loaded again.
+static constexpr size_t kMaxTrackedUnloadedModules = 4096;
+
 // Rejitter
 Rejitter::Rejitter(std::shared_ptr<RejitHandler> handler, RejitterPriority priority, bool registerRejitter) :
     m_rejitHandler(handler), m_priority(priority)
@@ -42,9 +46,14 @@ void RejitPreprocessor<RejitRequestDefinition>::Shutdown()
 
     std::lock_guard<std::mutex> moduleGuard(m_modules_lock);
     std::lock_guard<std::mutex> ngenModuleGuard(m_ngenInlinersModules_lock);
+    std::lock_guard<std::mutex> unloadedModuleGuard(m_unloaded_modules_lock);
 
     m_modules.clear();
     m_ngenInlinersModules.clear();
+
+    // Nothing reads m_unloaded_modules after shutdown, but tokens can still be outstanding, so
+    // clear the set and leave m_in_flight_requests for them to unwind.
+    m_unloaded_modules.clear();
 }
 
 template <class RejitRequestDefinition>
@@ -62,6 +71,9 @@ RejitHandlerModule* RejitPreprocessor<RejitRequestDefinition>::GetOrAddModule(Mo
         return find_res->second.get();
     }
 
+    // Deliberately does not clear m_unloaded_modules: this runs for unloaded modules too (callers
+    // such as RemoveProbes pass stored ModuleIDs), and clearing here would unprotect in-flight
+    // requests. NotifyModuleLoaded handles ModuleID reuse.
     RejitHandlerModule* moduleHandler = new RejitHandlerModule(moduleId, m_rejit_handler.get());
     m_modules[moduleId] = std::unique_ptr<RejitHandlerModule>(moduleHandler);
     return moduleHandler;
@@ -86,6 +98,47 @@ bool RejitPreprocessor<RejitRequestDefinition>::HasModuleAndMethod(ModuleID modu
     return false;
 }
 
+template <class RejitRequestDefinition>
+void RejitPreprocessor<RejitRequestDefinition>::NotifyModuleLoaded(ModuleID moduleId)
+{
+    if (m_rejit_handler->IsShutdownRequested())
+    {
+        return;
+    }
+
+    // The CLR can reuse a ModuleID, so drop it from the unloaded set.
+    std::lock_guard<std::mutex> unloadedGuard(m_unloaded_modules_lock);
+    m_unloaded_modules.erase(moduleId);
+}
+
+template <class RejitRequestDefinition>
+void RejitPreprocessor<RejitRequestDefinition>::AcquireInFlightRequest()
+{
+    std::lock_guard<std::mutex> unloadedGuard(m_unloaded_modules_lock);
+    ++m_in_flight_requests;
+}
+
+template <class RejitRequestDefinition>
+void RejitPreprocessor<RejitRequestDefinition>::ReleaseInFlightRequest()
+{
+    std::lock_guard<std::mutex> unloadedGuard(m_unloaded_modules_lock);
+    if (--m_in_flight_requests == 0)
+    {
+        // Nothing can name an unloaded module any more.
+        m_unloaded_modules.clear();
+    }
+}
+
+template <class RejitRequestDefinition>
+std::shared_ptr<void> RejitPreprocessor<RejitRequestDefinition>::AcquireInFlightRequestToken()
+{
+    AcquireInFlightRequest();
+
+    // The deleter runs when the last owner goes away, which also covers a work item that is
+    // discarded without running.
+    return std::shared_ptr<void>(nullptr, [this](void*) { ReleaseInFlightRequest(); });
+}
+
 template <class RejitRequestDefinition>
 void RejitPreprocessor<RejitRequestDefinition>::RemoveModule(ModuleID moduleId)
 {
@@ -95,13 +148,28 @@ void RejitPreprocessor<RejitRequestDefinition>::RemoveModule(ModuleID moduleId)
     }
 
     // Removes the RejitHandlerModule instance
+    //
+    // TOFIX (gleocadie): this destroys the handler (and its methods and metadata) while other threads can still
+    // be using it: GetOrAddModule hands out a raw pointer and drops m_modules_lock, so callers such
+    // as ProcessTypeDefForRejit and RejitMethod dereference it unsynchronized with this erase.
+    // The unloaded-module tracking below only stops us picking up a module that unloaded *before*
+    // the request ran; it does not protect a module that unloads mid-request. Pre-existing, and
+    // fixable by holding the handlers in shared_ptr so in-flight users keep them alive.
     std::lock_guard<std::mutex> modulesGuard(m_modules_lock);
     m_modules.erase(moduleId);
 
     // Removes the moduleID from the inliners vector
     std::lock_guard<std::mutex> inlinersGuard(m_ngenInlinersModules_lock);
     m_ngenInlinersModules.erase(std::remove(m_ngenInlinersModules.begin(), m_ngenInlinersModules.end(), moduleId),
                                 m_ngenInlinersModules.end());
+
+    // Record unloaded modules to avoid re-processing them. With no request in flight there is
+    // nothing holding a module list that could name this module, so there is nothing to track.
+    std::lock_guard<std::mutex> unloadedGuard(m_unloaded_modules_lock);
+    if (m_in_flight_requests > 0)
+    {
+        m_unloaded_modules.insert(moduleId);
+    }
 }
 
 template <class RejitRequestDefinition>
@@ -432,6 +500,8 @@ ULONG RejitPreprocessor<RejitRequestDefinition>::RequestRejitForLoadedModules(
     const std::vector<ModuleID>& modules, const std::vector<RejitRequestDefinition>& definitions,
     bool enqueueInSameThread)
 {
+    // No token taken here: every caller either already holds one for this module list, or holds
+    // CorProfiler::module_ids for the duration of the call, which blocks the unload.
     std::vector<MethodIdentifier> rejitRequests{};
     const auto rejitCount = PreprocessRejitRequests(modules, definitions, rejitRequests);
     RequestRejit(rejitRequests, enqueueInSameThread);
@@ -529,8 +599,12 @@ void RejitPreprocessor<RejitRequestDefinition>::EnqueueRequestRejitForLoadedModu
     DBG("RejitHandler::EnqueueRequestRejitForLoadedModules");
     auto enqueueMeasure = trace::Stats::Instance()->EnqueueRequestRejitForLoadedModulesMeasure();
 
+    // Taken before enqueuing: the modules can unload while the work item sits in the queue.
+    auto requestToken = AcquireInFlightRequestToken();
+
     std::function<void()> action = [=, modules = std::move(modulesVector), definitions = std::move(definitions),
-                                    localPromise = promise, enqueueMeasure = std::move(enqueueMeasure)]() mutable {
+                                    localPromise = promise, enqueueMeasure = std::move(enqueueMeasure),
+                                    requestToken = std::move(requestToken)]() mutable {
         // Process modules for rejit
         const auto rejitCount = RequestRejitForLoadedModules(modules, definitions, true);
 
@@ -562,6 +636,17 @@ ULONG RejitPreprocessor<RejitRequestDefinition>::PreprocessRejitRequests(
     for (const auto& module : modules)
     {
         auto _ = trace::Stats::Instance()->CallTargetRequestRejitMeasure();
+
+        {
+            // Skip modules that have been unloaded.
+            std::lock_guard<std::mutex> unloadedGuard(m_unloaded_modules_lock);
+            if (m_unloaded_modules.find(module) != m_unloaded_modules.end())
+            {
+                DBG("PreprocessRejitRequests: skipping module already unloaded: ", module);
+                continue;
+            }
+        }
+
         const ModuleInfo& moduleInfo = GetModuleInfo(corProfilerInfo, module);
         if (!moduleInfo.IsValid())
         {
@@ -876,8 +961,12 @@ void RejitPreprocessor<RejitRequestDefinition>::EnqueuePreprocessRejitRequests(
 
     DBG("RejitHandler::EnqueuePreprocessRejitRequests");
 
+    // Taken before enqueuing: the modules can unload while the work item sits in the queue.
+    auto requestToken = AcquireInFlightRequestToken();
+
     std::function<void()> action = [=, modules = std::move(modulesVector), definitions = std::move(definitions),
-                                    localRejitRequests = rejitRequests, localPromise = promise]() mutable {
+                                    localRejitRequests = rejitRequests, localPromise = promise,
+                                    requestToken = std::move(requestToken)]() mutable {
         // Process modules for rejit
         const auto rejitCount = PreprocessRejitRequests(modules, definitions, localRejitRequests);
 

@@ -3,6 +3,8 @@
 
 #include "integration.h"
 #include <future>
+#include <mutex>
+#include <unordered_set>
 #include "cor.h"
 #include "corprof.h"
 #include "module_metadata.h"
@@ -39,6 +41,9 @@ class Rejitter
     virtual void Shutdown() = 0;
     virtual RejitHandlerModule* GetOrAddModule(ModuleID moduleId) = 0;
     virtual bool HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef) = 0;
+    virtual void NotifyModuleLoaded(ModuleID moduleId) = 0;
+    virtual void AcquireInFlightRequest() = 0;
+    virtual void ReleaseInFlightRequest() = 0;
     virtual void RemoveModule(ModuleID moduleId) = 0;
     virtual void AddNGenInlinerModule(ModuleID moduleId) = 0;
 
@@ -95,19 +100,31 @@ class RejitPreprocessor : public Rejitter
                                   const std::vector<RejitRequestDefinition>& definitions,
                                   std::vector<MethodIdentifier>& rejitRequests);
 
+    // Counts one request as in flight until the returned token is destroyed. Must be held for as
+    // long as a module list taken from CorProfiler::module_ids is used.
+    std::shared_ptr<void> AcquireInFlightRequestToken();
+
 protected:
     std::mutex m_modules_lock;
     std::unordered_map<ModuleID, std::unique_ptr<RejitHandlerModule>> m_modules;
     std::mutex m_ngenInlinersModules_lock;
     std::vector<ModuleID> m_ngenInlinersModules;
 
+    // Unloaded modules handling: prevent from calling into the CLR for an unloaded module.
+    std::mutex m_unloaded_modules_lock;
+    std::unordered_set<ModuleID> m_unloaded_modules;
+    size_t m_in_flight_requests = 0;
+
 public:
     RejitPreprocessor(CorProfiler* corProfiler, std::shared_ptr<RejitHandler> rejit_handler,
                       std::shared_ptr<RejitWorkOffloader> work_offloader, RejitterPriority priority);
 
     void Shutdown() override;
     RejitHandlerModule* GetOrAddModule(ModuleID moduleId) override;
     bool HasModuleAndMethod(ModuleID moduleId, mdMethodDef methodDef) override;
+    void NotifyModuleLoaded(ModuleID moduleId) override;
+    void AcquireInFlightRequest() override;
+    void ReleaseInFlightRequest() override;
     void RemoveModule(ModuleID moduleId) override;
     void AddNGenInlinerModule(ModuleID moduleId) override;
     HRESULT RejitMethod(FunctionControlWrapper& functionControl) override;