Fix LD_PRELOAD (#184)

- __libc_start_main in libomnitrace-dl wasn't wrapping bc of -fvisibility=hidden - fix OMNITRACE_STRIP_TARGET - omnitrace_reset_preload function in main library - defer removing libomnitrace from LD_PRELOAD
ROCm · Oct 21, 2022 · a41a5c1 · a41a5c1
1 parent f2a0485
commit a41a5c1
Show file tree

Hide file tree

Showing 8 changed files with 111 additions and 20 deletions.
diff --git a/cmake/MacroUtilities.cmake b/cmake/MacroUtilities.cmake
@@ -122,11 +122,11 @@ function(OMNITRACE_STRIP_TARGET _TARGET)
                 --keep-symbol="omnitrace_finalize" --keep-symbol="omnitrace_push_trace"
                 --keep-symbol="omnitrace_pop_trace" --keep-symbol="omnitrace_push_region"
                 --keep-symbol="omnitrace_pop_region" --keep-symbol="omnitrace_set_env"
-                --keep-symbol="omnitrace_set_mpi" --keep-symbol="omnitrace_user_*"
-                --keep-symbol="ompt_start_tool" --keep-symbol="kokkosp_*"
-                --keep-symbol="OnLoad" --keep-symbol="OnUnload"
-                --keep-symbol="OnLoadToolProp" --keep-symbol="OnUnloadTool" ${ARGN}
-                $<TARGET_FILE:${_TARGET}>
+                --keep-symbol="omnitrace_set_mpi" --keep-symbol="omnitrace_reset_preload"
+                --keep-symbol="omnitrace_user_*" --keep-symbol="ompt_start_tool"
+                --keep-symbol="kokkosp_*" --keep-symbol="OnLoad" --keep-symbol="OnUnload"
+                --keep-symbol="OnLoadToolProp" --keep-symbol="OnUnloadTool"
+                --keep-symbol="__libc_start_main" ${ARGN} $<TARGET_FILE:${_TARGET}>
             WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
             COMMENT "Stripping ${_TARGET}...")
     endif()

diff --git a/source/bin/omnitrace-avail/CMakeLists.txt b/source/bin/omnitrace-avail/CMakeLists.txt
@@ -31,6 +31,8 @@ set_target_properties(
     omnitrace-avail PROPERTIES BUILD_RPATH "\$ORIGIN:\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}"
                                INSTALL_RPATH "${OMNITRACE_EXE_INSTALL_RPATH}")
 
+omnitrace_strip_target(omnitrace-avail)
+
 install(
     TARGETS omnitrace-avail
     DESTINATION ${CMAKE_INSTALL_BINDIR}

diff --git a/source/bin/omnitrace-sample/CMakeLists.txt b/source/bin/omnitrace-sample/CMakeLists.txt
@@ -17,6 +17,8 @@ set_target_properties(
     omnitrace-sample PROPERTIES BUILD_RPATH "\$ORIGIN:\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}"
                                 INSTALL_RPATH "${OMNITRACE_EXE_INSTALL_RPATH}")
 
+omnitrace_strip_target(omnitrace-sample)
+
 install(
     TARGETS omnitrace-sample
     DESTINATION ${CMAKE_INSTALL_BINDIR}

diff --git a/source/lib/omnitrace-dl/dl.cpp b/source/lib/omnitrace-dl/dl.cpp
@@ -963,7 +963,7 @@ omnitrace_preload()
 
     if(_preload)
     {
-        reset_omnitrace_preload();
+        // reset_omnitrace_preload();
         omnitrace_preinit_library();
         OMNITRACE_DL_LOG(1, "[%s] invoking %s(%s)\n", __FUNCTION__, "omnitrace_init",
                          ::omnitrace::join(::omnitrace::QuoteStrings{}, ", ", "sampling",

diff --git a/source/lib/omnitrace-dl/main.c b/source/lib/omnitrace-dl/main.c
@@ -22,13 +22,46 @@
 
 #define _GNU_SOURCE
 
+#define OMNITRACE_PUBLIC_API __attribute__((visibility("default")));
+#define OMNITRACE_HIDDEN_API __attribute__((visibility("hidden")));
+
 #include <dlfcn.h>
 #include <stdio.h>
 #include <unistd.h>
 #include <stdlib.h>
 #include <string.h>
 #include <stdbool.h>
 
+//
+// local type definitions
+//
+typedef int (*start_main_t)(int (*)(int, char**, char**), int, char**,
+                            int (*)(int, char**, char**), void (*)(void), void (*)(void),
+                            void*);
+
+//
+// local variables
+//
+static int (*main_real)(int, char**, char**);  // Trampoline for the real main()
+
+//
+// local function declarations
+//
+int
+omnitrace_main(int, char**, char**) OMNITRACE_HIDDEN_API;
+
+int
+omnitrace_libc_start_main(int (*)(int, char**, char**), int, char**,
+                     int (*)(int, char**, char**), void (*)(void), void (*)(void),
+                     void*) OMNITRACE_HIDDEN_API;
+
+int
+__libc_start_main(int (*)(int, char**, char**), int, char**, int (*)(int, char**, char**),
+                  void (*)(void), void (*)(void), void*) OMNITRACE_PUBLIC_API;
+
+//
+// external function declarations
+//
 extern int
 omnitrace_preload_library(void);
 
@@ -41,18 +74,15 @@ omnitrace_push_trace(const char* name);
 extern void
 omnitrace_pop_trace(const char* name);
 
-// extern void
-// omnitrace_update_env(char*** envp);
-
 extern void
 omnitrace_init_tooling(void);
 
 extern void
 omnitrace_init(const char*, bool, const char*);
 
-// Trampoline for the real main()
-static int (*main_real)(int, char**, char**);
-
+//
+//  local function definitions
+//
 int
 omnitrace_main(int argc, char** argv, char** envp)
 {
@@ -77,13 +107,8 @@ omnitrace_main(int argc, char** argv, char** envp)
     return ret;
 }
 
-typedef int
-(*omnitrace_libc_start_main)(int (*)(int, char**, char**), int, char**,
-                  int (*)(int, char**, char**), void (*)(void),
-                  void (*)(void), void*);
-
 int
-__libc_start_main(int (*_main)(int, char**, char**), int _argc, char** _argv,
+omnitrace_libc_start_main(int (*_main)(int, char**, char**), int _argc, char** _argv,
                   int (*_init)(int, char**, char**), void (*_fini)(void),
                   void (*_rtld_fini)(void), void* _stack_end)
 {
@@ -101,7 +126,7 @@ __libc_start_main(int (*_main)(int, char**, char**), int _argc, char** _argv,
     main_real = _main;
 
     // Find the real __libc_start_main()
-    omnitrace_libc_start_main user_main = dlsym(RTLD_NEXT, "__libc_start_main");
+    start_main_t user_main = dlsym(RTLD_NEXT, "__libc_start_main");
 
     // disable future LD_PRELOADs
     setenv("OMNITRACE_PRELOAD", "0", 1);
@@ -127,3 +152,12 @@ __libc_start_main(int (*_main)(int, char**, char**), int _argc, char** _argv,
         return -1;
     }
 }
+
+int
+__libc_start_main(int (*_main)(int, char**, char**), int _argc, char** _argv,
+                  int (*_init)(int, char**, char**), void (*_fini)(void),
+                  void (*_rtld_fini)(void), void* _stack_end)
+{
+    return omnitrace_libc_start_main(_main, _argc, _argv, _init, _fini, _rtld_fini,
+                                _stack_end);
+}
diff --git a/source/lib/omnitrace/api.cpp b/source/lib/omnitrace/api.cpp
@@ -90,6 +90,12 @@ omnitrace_finalize(void)
     omnitrace_finalize_hidden();
 }
 
+extern "C" void
+omnitrace_reset_preload(void)
+{
+    omnitrace_reset_preload_hidden();
+}
+
 extern "C" void
 omnitrace_set_env(const char* env_name, const char* env_val)
 {

diff --git a/source/lib/omnitrace/api.hpp b/source/lib/omnitrace/api.hpp
@@ -43,6 +43,9 @@ extern "C"
     /// shuts down all tooling and generates output
     void omnitrace_finalize(void) OMNITRACE_PUBLIC_API;
 
+    /// remove libomnitrace from LD_PRELOAD
+    void omnitrace_reset_preload(void) OMNITRACE_PUBLIC_API;
+
     /// sets an environment variable
     void omnitrace_set_env(const char* env_name,
                            const char* env_val) OMNITRACE_PUBLIC_API;
@@ -76,6 +79,7 @@ extern "C"
     bool omnitrace_init_tooling_hidden(void) OMNITRACE_HIDDEN_API;
     void omnitrace_init_hidden(const char*, bool, const char*) OMNITRACE_HIDDEN_API;
     void omnitrace_finalize_hidden(void) OMNITRACE_HIDDEN_API;
+    void omnitrace_reset_preload_hidden(void) OMNITRACE_HIDDEN_API;
     void omnitrace_set_env_hidden(const char* env_name,
                                   const char* env_val) OMNITRACE_HIDDEN_API;
     void omnitrace_set_mpi_hidden(bool use, bool attached) OMNITRACE_HIDDEN_API;

diff --git a/source/lib/omnitrace/library.cpp b/source/lib/omnitrace/library.cpp
@@ -424,7 +424,7 @@ omnitrace_init_tooling_hidden()
     omnitrace_preinit_hidden();
 
     // start these gotchas once settings have been initialized
-    get_init_bundle()->start();
+    if(get_init_bundle()) get_init_bundle()->start();
 
     if(get_use_sampling()) sampling::block_signals();
 
@@ -564,6 +564,28 @@ omnitrace_init_hidden(const char* _mode, bool _is_binary_rewrite, const char* _a
 
 //======================================================================================//
 
+extern "C" void
+omnitrace_reset_preload_hidden(void)
+{
+    tim::set_env("OMNITRACE_PRELOAD", "0", 1);
+    auto&& _preload_libs = get_env("LD_PRELOAD", std::string{});
+    if(_preload_libs.find("libomnitrace") != std::string::npos)
+    {
+        auto _modified_preload = std::string{};
+        for(const auto& itr : delimit(_preload_libs, ":"))
+        {
+            if(itr.find("libomnitrace") != std::string::npos) continue;
+            _modified_preload += common::join("", ":", itr);
+        }
+        if(!_modified_preload.empty() && _modified_preload.find(':') == 0)
+            _modified_preload = _modified_preload.substr(1);
+
+        tim::set_env("LD_PRELOAD", _modified_preload, 1);
+    }
+}
+
+//======================================================================================//
+
 extern "C" void
 omnitrace_finalize_hidden(void)
 {
@@ -590,6 +612,8 @@ omnitrace_finalize_hidden(void)
     tim::signals::block_signals(get_sampling_signals(),
                                 tim::signals::sigmask_scope::process);
 
+    omnitrace_reset_preload_hidden();
+
     // some functions called during finalization may alter the push/pop count so we need
     // to save them here
     auto _push_count = tracing::push_count().load();
@@ -635,6 +659,13 @@ omnitrace_finalize_hidden(void)
     tim::signals::enable_signal_detection({ tim::signals::sys_signal::Interrupt },
                                           [](int) {});
 
+    std::string              _bundle_name = OMNITRACE_FUNCTION;
+    comp::user_global_bundle _bundle{ _bundle_name.c_str() };
+    _bundle.clear();
+    _bundle.insert<comp::wall_clock, comp::cpu_clock, comp::peak_rss, comp::page_rss,
+                   comp::cpu_util>();
+    _bundle.start();
+
     OMNITRACE_DEBUG_F("Copying over all timemory hash information to main thread...\n");
     // copy these over so that all hashes are known
     auto& _hzero = tracing::get_timemory_hash_ids(0);
@@ -929,6 +960,18 @@ omnitrace_finalize_hidden(void)
         }
     }
 
+    _bundle.stop();
+    auto _get_metric = [](auto* _v, std::string_view _tail) -> std::string {
+        return (_v) ? JOIN("", *_v, _tail) : std::string{};
+    };
+
+    OMNITRACE_VERBOSE_F(0, "Finalization metrics: %s%s%s%s%s\n",
+                        _get_metric(_bundle.get<comp::wall_clock>(), ", ").c_str(),
+                        _get_metric(_bundle.get<comp::peak_rss>(), ", ").c_str(),
+                        _get_metric(_bundle.get<comp::page_rss>(), ", ").c_str(),
+                        _get_metric(_bundle.get<comp::cpu_clock>(), ", ").c_str(),
+                        _get_metric(_bundle.get<comp::cpu_util>(), "").c_str());
+
     if(_timemory_manager && _timemory_manager != nullptr)
     {
         _timemory_manager->add_metadata([](auto& ar) {