From 7075a0a87da6955810583f212bc2865293172e6e Mon Sep 17 00:00:00 2001
From: Igor Chorazewicz <igor.chorazewicz@intel.com>
Date: Wed, 21 Aug 2024 23:27:35 +0200
Subject: [PATCH] [Loader] Fix xpti finalization order

Global xptiContext was beeing (sometimes) destroyed before the
tracing context (as both were global variables, their order of
destruction was not guaranteed). This resulted in lost
notifications as xptiFinalize was being called after xptiFrameworkFinalize.

Fix this by creating a global shared_ptr to the xptiContext.
Bump the refcount of that pointer every time a context is created
and decrease when it is destroyed to make sure that xptiFrameworkInitialize()
is not beeing called before the last context is destroyed.
---
 .../layers/tracing/ur_tracing_layer.cpp       | 51 +++----------------
 .../layers/tracing/ur_tracing_layer.hpp       |  4 ++
 2 files changed, 12 insertions(+), 43 deletions(-)

diff --git a/source/loader/layers/tracing/ur_tracing_layer.cpp b/source/loader/layers/tracing/ur_tracing_layer.cpp
index 140f653ba5..722ee77faa 100644
--- a/source/loader/layers/tracing/ur_tracing_layer.cpp
+++ b/source/loader/layers/tracing/ur_tracing_layer.cpp
@@ -29,36 +29,19 @@ constexpr auto STREAM_VER_MINOR = UR_MINOR_VERSION(UR_API_VERSION_CURRENT);
 // Unfortunately this doesn't match the semantics of XPTI, which can be initialized
 // and finalized exactly once. To workaround this, XPTI is globally initialized on
 // first use and finalized in the destructor.
-class XptiContext {
-    XptiContext() {
-        xptiFrameworkInitialize();
-        inited = true;
-    }
-
-    ~XptiContext() {
-        xptiFrameworkFinalize();
-        inited = false;
-    }
-
-    // Accessing this after destruction is technically UB, but if we get there,
-    // it means something is calling UR after it has been destroyed at program
-    // exit.
-    std::atomic_bool inited;
-
-  public:
-    static bool running() {
-        static XptiContext context;
-        return context.inited;
-    }
+struct XptiContextManager {
+    XptiContextManager() { xptiFrameworkInitialize(); }
+    ~XptiContextManager() { xptiFrameworkFinalize(); }
 };
 
+static std::shared_ptr<XptiContextManager> xptiContextManagerGlobal = [] {
+    return std::make_shared<XptiContextManager>();
+}();
 static thread_local xpti_td *activeEvent;
 
 ///////////////////////////////////////////////////////////////////////////////
 context_t::context_t() : logger(logger::create_logger("tracing", true, true)) {
-    if (!XptiContext::running()) {
-        return;
-    }
+    this->xptiContextManager = xptiContextManagerGlobal;
 
     call_stream_id = xptiRegisterStream(CALL_STREAM_NAME);
     std::ostringstream streamv;
@@ -69,20 +52,12 @@ context_t::context_t() : logger(logger::create_logger("tracing", true, true)) {
 
 void context_t::notify(uint16_t trace_type, uint32_t id, const char *name,
                        void *args, ur_result_t *resultp, uint64_t instance) {
-    if (!XptiContext::running()) {
-        return;
-    }
-
     xpti::function_with_args_t payload{id, name, args, resultp, nullptr};
     xptiNotifySubscribers(call_stream_id, trace_type, nullptr, activeEvent,
                           instance, &payload);
 }
 
 uint64_t context_t::notify_begin(uint32_t id, const char *name, void *args) {
-    if (!XptiContext::running()) {
-        return 0;
-    }
-
     if (auto loc = codelocData.get_codeloc()) {
         xpti::payload_t payload =
             xpti::payload_t(loc->functionName, loc->sourceFile, loc->lineNumber,
@@ -101,20 +76,10 @@ uint64_t context_t::notify_begin(uint32_t id, const char *name, void *args) {
 
 void context_t::notify_end(uint32_t id, const char *name, void *args,
                            ur_result_t *resultp, uint64_t instance) {
-    if (!XptiContext::running()) {
-        return;
-    }
-
     notify((uint16_t)xpti::trace_point_type_t::function_with_args_end, id, name,
            args, resultp, instance);
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-context_t::~context_t() {
-    if (!XptiContext::running()) {
-        return;
-    }
-
-    xptiFinalize(CALL_STREAM_NAME);
-}
+context_t::~context_t() { xptiFinalize(CALL_STREAM_NAME); }
 } // namespace ur_tracing_layer
diff --git a/source/loader/layers/tracing/ur_tracing_layer.hpp b/source/loader/layers/tracing/ur_tracing_layer.hpp
index 1a5c542ee6..b7e3fc0314 100644
--- a/source/loader/layers/tracing/ur_tracing_layer.hpp
+++ b/source/loader/layers/tracing/ur_tracing_layer.hpp
@@ -21,6 +21,8 @@
 #define TRACING_COMP_NAME "tracing layer"
 
 namespace ur_tracing_layer {
+struct XptiContextManager;
+
 ///////////////////////////////////////////////////////////////////////////////
 class __urdlllocal context_t : public proxy_layer_context_t,
                                public AtomicSingleton<context_t> {
@@ -47,6 +49,8 @@ class __urdlllocal context_t : public proxy_layer_context_t,
     uint8_t call_stream_id;
 
     inline static const std::string name = "UR_LAYER_TRACING";
+
+    std::shared_ptr<XptiContextManager> xptiContextManager;
 };
 
 context_t *getContext();