From d57da896be15c76560a7987b5ff0273413a277b7 Mon Sep 17 00:00:00 2001
From: Shivam Raikundalia <sraikund@meta.com>
Date: Wed, 9 Oct 2024 13:39:52 -0700
Subject: [PATCH] Reset TLS on Profiling Entrance

Summary:
D63924780 broke some tests because of pthread_atfork having strange properties with subsequent calls. To fix this, lets deviate from this method and just reset the TLS whenever we enter the a profiling context. This will ensure that we will start "fresh" for all the PID/TID related content upon every profile. The drawbacks are:

1. 1 Extra Cache Miss per Profile - This is negligible because the cache miss is during the prepare stage for auto-trace and schedule for on-demand. To add to this, the cache miss penalty is very small

2. No Reset if Fork during Profile - If someone were to fork in the middle of a profile the TLS won't get reset. However, there are many other issues that could happen due to a fork midway through a profile such as undefined behavior with cupti, distorted profiling window etc. We shouldn't worry about this case as of today.

Differential Revision: D64120658
---
 libkineto/include/ThreadUtil.h          | 4 ++++
 libkineto/include/libkineto.h           | 4 ++++
 libkineto/src/ActivityProfilerProxy.cpp | 2 ++
 libkineto/src/ConfigLoader.cpp          | 4 +++-
 libkineto/src/ThreadUtil.cpp            | 8 ++++++++
 libkineto/src/init.cpp                  | 1 +
 6 files changed, 22 insertions(+), 1 deletion(-)

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<std::pair<int32_t, std::string>> 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<std::mutex> 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"