diff --git a/libkineto/include/ThreadUtil.h b/libkineto/include/ThreadUtil.h index 4178ae4a1..ca52c6343 100644 --- a/libkineto/include/ThreadUtil.h +++ b/libkineto/include/ThreadUtil.h @@ -27,4 +27,8 @@ std::string processName(int32_t pid); // and its parents. std::vector> pidCommandPairsOfAncestors(); +// Resets all cached Thread local state, this must be done on +// forks to prevent stale values from being retained. +void resetTLS(); + } // namespace libkineto diff --git a/libkineto/include/libkineto.h b/libkineto/include/libkineto.h index 6fc571b34..a122a77a5 100644 --- a/libkineto/include/libkineto.h +++ b/libkineto/include/libkineto.h @@ -117,6 +117,10 @@ class LibkinetoApi { suppressLibkinetoLogMessages(); } + void resetKinetoTLS() { + resetTLS(); + } + // Provides access to profier configuration manaegement ConfigLoader& configLoader() { return configLoader_; diff --git a/libkineto/src/ActivityProfilerProxy.cpp b/libkineto/src/ActivityProfilerProxy.cpp index 8400fd052..fe17a31b5 100644 --- a/libkineto/src/ActivityProfilerProxy.cpp +++ b/libkineto/src/ActivityProfilerProxy.cpp @@ -12,6 +12,7 @@ #include "ActivityProfilerController.h" #include "Config.h" #include "Logger.h" +#include "ThreadUtil.h" namespace KINETO_NAMESPACE { @@ -31,6 +32,7 @@ void ActivityProfilerProxy::init() { } void ActivityProfilerProxy::scheduleTrace(const std::string& configStr) { + resetTLS(); Config config; config.parse(configStr); controller_->scheduleTrace(config); diff --git a/libkineto/src/ConfigLoader.cpp b/libkineto/src/ConfigLoader.cpp index 42b0b8163..4fe40c607 100644 --- a/libkineto/src/ConfigLoader.cpp +++ b/libkineto/src/ConfigLoader.cpp @@ -172,7 +172,9 @@ void ConfigLoader::stopThread() { std::lock_guard lock(updateThreadMutex_); updateThreadCondVar_.notify_one(); } - updateThread_->join(); + if (updateThread_->joinable()) { + updateThread_->join(); + } updateThread_ = nullptr; } } diff --git a/libkineto/src/ThreadUtil.cpp b/libkineto/src/ThreadUtil.cpp index f9ec041e2..6f9429d16 100644 --- a/libkineto/src/ThreadUtil.cpp +++ b/libkineto/src/ThreadUtil.cpp @@ -92,6 +92,14 @@ int32_t threadId() { return _tid; } +// Resets all cached Thread local state, this must be done on +// forks to prevent stale values from being retained. +void resetTLS() { + _pid = 0; + _tid = 0; + _sysTid = 0; +} + namespace { static constexpr size_t kMaxThreadNameLength = 16; diff --git a/libkineto/src/init.cpp b/libkineto/src/init.cpp index 6246db8e3..51dd332f4 100644 --- a/libkineto/src/init.cpp +++ b/libkineto/src/init.cpp @@ -16,6 +16,7 @@ #include "ConfigLoader.h" #include "DaemonConfigLoader.h" #include "DeviceUtil.h" +#include "ThreadUtil.h" #ifdef HAS_CUPTI #include "CuptiActivityApi.h" #include "CuptiCallbackApi.h"