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/rocprofiler-systems commit: a41a5c155e]
2022-10-21 17:39:08 -05:00
@@ -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()
@@ -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}
@@ -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}
@@ -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",
@@ -22,6 +22,9 @@

 #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>
@@ -29,6 +32,36 @@
 #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);
+}
@@ -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)
 {
@@ -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;
@@ -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) {