diff --git a/source/bin/omnitrace-run/impl.cpp b/source/bin/omnitrace-run/impl.cpp index 58e0e8f94..ae3f486c2 100644 --- a/source/bin/omnitrace-run/impl.cpp +++ b/source/bin/omnitrace-run/impl.cpp @@ -191,6 +191,7 @@ prepare_environment_for_run(parser_data_t& _data) if(_data.launcher.empty()) { omnitrace::argparse::add_ld_preload(_data); + omnitrace::argparse::add_ld_library_path(_data); } } diff --git a/source/bin/omnitrace-sample/impl.cpp b/source/bin/omnitrace-sample/impl.cpp index ce5abc848..4833f9b56 100644 --- a/source/bin/omnitrace-sample/impl.cpp +++ b/source/bin/omnitrace-sample/impl.cpp @@ -31,6 +31,7 @@ #include #include #include +#include #include #include @@ -129,13 +130,11 @@ get_initial_environment() } } - update_env(_env, "LD_PRELOAD", - get_realpath(get_internal_libpath("libomnitrace-dl.so")), true); + auto _dl_libpath = get_realpath(get_internal_libpath("libomnitrace-dl.so")); + auto _omni_libpath = get_realpath(get_internal_libpath("libomnitrace.so")); - auto* _dl_libpath = - realpath(get_internal_libpath("libomnitrace-dl.so").c_str(), nullptr); - auto* _omni_libpath = - realpath(get_internal_libpath("libomnitrace.so").c_str(), nullptr); + update_env(_env, "LD_PRELOAD", _dl_libpath, UPD_APPEND); + update_env(_env, "LD_LIBRARY_PATH", tim::filepath::dirname(_dl_libpath), UPD_APPEND); auto _mode = get_env("OMNITRACE_MODE", "sampling", false); @@ -154,12 +153,9 @@ get_initial_environment() #if defined(OMNITRACE_USE_OMPT) if(!getenv("OMP_TOOL_LIBRARIES")) - update_env(_env, "OMP_TOOL_LIBRARIES", _dl_libpath, true); + update_env(_env, "OMP_TOOL_LIBRARIES", _dl_libpath, UPD_APPEND); #endif - free(_dl_libpath); - free(_omni_libpath); - return _env; } @@ -223,29 +219,42 @@ print_updated_environment(std::vector _env) template void update_env(std::vector& _environ, std::string_view _env_var, Tp&& _env_val, - bool _append) + update_mode&& _mode, std::string_view _join_delim) { updated_envs.emplace(_env_var); + auto _prepend = (_mode & UPD_PREPEND) == UPD_PREPEND; + auto _append = (_mode & UPD_APPEND) == UPD_APPEND; + auto _weak_upd = (_mode & UPD_WEAK) == UPD_WEAK; + auto _key = join("", _env_var, "="); for(auto& itr : _environ) { if(!itr) continue; if(std::string_view{ itr }.find(_key) == 0) { - if(_append) + if(_weak_upd) + { + // if the value has changed, do not update but allow overridding the value + // inherited from the initial env + if(original_envs.find(std::string{ itr }) == original_envs.end()) return; + } + + if(_prepend || _append) { if(std::string_view{ itr }.find(join("", _env_val)) == std::string_view::npos) { auto _val = std::string{ itr }.substr(_key.length()); free(itr); - if(_env_var == "LD_PRELOAD") - itr = strdup( - join('=', _env_var, join(":", _val, _env_val)).c_str()); + if(_prepend) + itr = + strdup(join('=', _env_var, join(_join_delim, _val, _env_val)) + .c_str()); else - itr = strdup( - join('=', _env_var, join(":", _env_val, _val)).c_str()); + itr = + strdup(join('=', _env_var, join(_join_delim, _env_val, _val)) + .c_str()); } } else @@ -783,10 +792,10 @@ parse_args(int argc, char** argv, std::vector& _env) _update("OMNITRACE_TRACE_THREAD_SPIN_LOCKS", _v.count("spin-locks") > 0); if(_v.count("all") > 0 || _v.count("ompt") > 0) - update_env(_env, "OMP_TOOL_LIBRARIES", _dl_libpath, true); + update_env(_env, "OMP_TOOL_LIBRARIES", _dl_libpath, UPD_APPEND); if(_v.count("all") > 0 || _v.count("kokkosp") > 0) - update_env(_env, "KOKKOS_PROFILE_LIBRARY", _omni_libpath, true); + update_env(_env, "KOKKOS_PROFILE_LIBRARY", _omni_libpath, UPD_APPEND); }); parser.add_argument({ "-E", "--exclude" }, "Exclude data from these backends") diff --git a/source/bin/omnitrace-sample/omnitrace-sample.hpp b/source/bin/omnitrace-sample/omnitrace-sample.hpp index e09588b1a..5c446a40a 100644 --- a/source/bin/omnitrace-sample/omnitrace-sample.hpp +++ b/source/bin/omnitrace-sample/omnitrace-sample.hpp @@ -26,13 +26,22 @@ #include #include +enum update_mode : int +{ + UPD_REPLACE = 0x1, + UPD_PREPEND = 0x2, + UPD_APPEND = 0x3, + UPD_WEAK = 0x4, +}; + std::string -get_realpath(const std::string&); +get_realpath(const std::string& _fpath); void print_command(const std::vector& _argv); -void print_updated_environment(std::vector); +void +print_updated_environment(std::vector _env); std::vector get_initial_environment(); @@ -42,10 +51,11 @@ get_internal_libpath(const std::string& _lib); template void -update_env(std::vector&, std::string_view, Tp&&, bool _append = false); +update_env(std::vector& _environ, std::string_view _env_var, Tp&& _env_val, + update_mode&& _mode = UPD_REPLACE, std::string_view _join_delim = ":"); void -remove_env(std::vector&, std::string_view); +remove_env(std::vector& _environ, std::string_view _env_var); std::vector -parse_args(int argc, char** argv, std::vector&); +parse_args(int argc, char** argv, std::vector& envp); diff --git a/source/lib/core/argparse.cpp b/source/lib/core/argparse.cpp index c140aaa50..c91ad674b 100644 --- a/source/lib/core/argparse.cpp +++ b/source/lib/core/argparse.cpp @@ -248,6 +248,15 @@ add_ld_preload(parser_data& _data) return _data; } +parser_data& +add_ld_library_path(parser_data& _data) +{ + auto _libdir = filepath::dirname(_data.dl_libpath); + if(filepath::exists(_libdir)) + update_env(_data, "LD_LIBRARY_PATH", _libdir, UPD_APPEND); + return _data; +} + parser_data& add_core_arguments(parser_t& _parser, parser_data& _data) { diff --git a/source/lib/core/argparse.hpp b/source/lib/core/argparse.hpp index 7a7d55ce3..0e30463d3 100644 --- a/source/lib/core/argparse.hpp +++ b/source/lib/core/argparse.hpp @@ -81,6 +81,9 @@ init_parser(parser_data&); parser_data& add_ld_preload(parser_data&); +parser_data& +add_ld_library_path(parser_data&); + parser_data& add_core_arguments(parser_t&, parser_data&);