Rejit - Switch to shared pointer and apply same behavior for initial …

…call (#3877) Co-authored-by: Paulo Janotti <[email protected]>
open-telemetry · Dec 13, 2024 · 9a4cd41 · 9a4cd41
1 parent 1f39185
commit 9a4cd41
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 13 deletions.
diff --git a/src/OpenTelemetry.AutoInstrumentation.Native/cor_profiler.cpp b/src/OpenTelemetry.AutoInstrumentation.Native/cor_profiler.cpp
@@ -36,8 +36,6 @@
 #include <mach-o/getsect.h>
 #endif
 
-using namespace std::chrono_literals;
-
 #ifdef _WIN32
 #include "netfx_assembly_redirection.h"
 #endif
@@ -53,6 +51,8 @@ using namespace std::chrono_literals;
         return E_FAIL;                                                                                                 \
     }
 
+using namespace std::chrono_literals;
+
 namespace trace
 {
 
@@ -788,10 +788,10 @@ HRESULT STDMETHODCALLTYPE CorProfiler::ModuleLoadFinished(ModuleID module_id, HR
         // We call the function to analyze the module and request the ReJIT of integrations defined in this module.
         if (tracer_integration_preprocessor != nullptr && !integration_definitions_.empty())
         {
-            std::promise<ULONG> promise;
-            std::future<ULONG>  future = promise.get_future();
+            auto               promise = std::make_shared<std::promise<ULONG>>();
+            std::future<ULONG> future  = promise->get_future();
             tracer_integration_preprocessor->EnqueueRequestRejitForLoadedModules(std::vector<ModuleID>{module_id},
-                                                                                 integration_definitions_, &promise);
+                                                                                 integration_definitions_, promise);
 
             // wait and get the value from the future<ULONG>
             const auto status = future.wait_for(100ms);
@@ -1109,14 +1109,24 @@ void CorProfiler::InternalAddInstrumentation(WCHAR* id, CallTargetDefinition* it
         {
             if (rejit_handler != nullptr)
             {
-                std::promise<ULONG> promise;
-                std::future<ULONG>  future = promise.get_future();
+                auto               promise = std::make_shared<std::promise<ULONG>>();
+                std::future<ULONG> future  = promise->get_future();
                 tracer_integration_preprocessor->EnqueueRequestRejitForLoadedModules(module_ids_,
-                                                                                     integrationDefinitions, &promise);
+                                                                                     integrationDefinitions, promise);
 
-                // wait and get the value from the future<int>
-                const auto& numReJITs = future.get();
-                Logger::Debug("Total number of ReJIT Requested: ", numReJITs);
+                // wait and get the value from the future<ULONG>
+                const auto status = future.wait_for(100ms);
+
+                if (status != std::future_status::timeout)
+                {
+                    const auto& numReJITs = future.get();
+                    Logger::Debug("Total number of ReJIT Requested: ", numReJITs);
+                }
+                else
+                {
+                    Logger::Warn("Timeout while waiting for the rejit requests to be processed. Rejit will continue "
+                                 "asynchronously, but some initial calls may not be instrumented");
+                }
             }
 
             integration_definitions_.reserve(integration_definitions_.size() + integrationDefinitions.size());

diff --git a/src/OpenTelemetry.AutoInstrumentation.Native/rejit_preprocessor.cpp b/src/OpenTelemetry.AutoInstrumentation.Native/rejit_preprocessor.cpp
@@ -395,7 +395,7 @@ template <class RejitRequestDefinition>
 void RejitPreprocessor<RejitRequestDefinition>::EnqueueRequestRejitForLoadedModules(
     const std::vector<ModuleID>&               modulesVector,
     const std::vector<RejitRequestDefinition>& definitions,
-    std::promise<ULONG>*                       promise)
+    std::shared_ptr<std::promise<ULONG>>       promise)
 {
     if (m_rejit_handler->IsShutdownRequested())
     {

diff --git a/src/OpenTelemetry.AutoInstrumentation.Native/rejit_preprocessor.h b/src/OpenTelemetry.AutoInstrumentation.Native/rejit_preprocessor.h
@@ -53,7 +53,7 @@ class RejitPreprocessor
 
     void EnqueueRequestRejitForLoadedModules(const std::vector<ModuleID>& modulesVector,
                                              const std::vector<RejitRequestDefinition>& requests,
-                                             std::promise<ULONG>* promise);
+                                             std::shared_ptr<std::promise<ULONG>> promise);
 };
 
 /// <summary>