From a2f124e60d80a2783d3309fed66ed2ad979303be Mon Sep 17 00:00:00 2001
From: Shivam Raikundalia <sraikund@meta.com>
Date: Tue, 1 Oct 2024 12:06:19 -0700
Subject: [PATCH] Do Not Cache PID/TID in Logging (#994)

Summary:
Pull Request resolved: https://github.com/pytorch/kineto/pull/994

S451588 was caused by the LOG macro caching values that would then be copied to other processes via forking. In general, we should probably use fork handlers to clear out said variables, but from a hygiene point of view we should also not be changing control flow based on logging. For this reason, the pid/tid retrieval in logging should get the cached variable if it exists, but never do the caching itself.

Differential Revision: D63668265
---
 libkineto/include/ThreadUtil.h |  6 +++---
 libkineto/src/Logger.cpp       |  2 +-
 libkineto/src/ThreadUtil.cpp   | 39 ++++++++++++++++++++++------------
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/libkineto/include/ThreadUtil.h b/libkineto/include/ThreadUtil.h
index 12b37da55..4260eac85 100644
--- a/libkineto/include/ThreadUtil.h
+++ b/libkineto/include/ThreadUtil.h
@@ -15,12 +15,12 @@
 
 namespace libkineto {
 
-int32_t systemThreadId();
-int32_t threadId();
+int32_t systemThreadId(bool cache=true);
+int32_t threadId(bool cache=true);
 bool setThreadName(const std::string& name);
 std::string getThreadName();
 
-int32_t processId();
+int32_t processId(bool cache=true);
 std::string processName(int32_t pid);
 
 // Return a list of pids and process names for the current process
diff --git a/libkineto/src/Logger.cpp b/libkineto/src/Logger.cpp
index d4886506f..7a4b771d9 100644
--- a/libkineto/src/Logger.cpp
+++ b/libkineto/src/Logger.cpp
@@ -39,7 +39,7 @@ Logger::Logger(int severity, int line, const char* filePath, int errnum)
       std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
   const char* file = strrchr(filePath, '/');
   buf_ << fmt::format("{:%Y-%m-%d %H:%M:%S}", fmt::localtime(tt)) << " "
-       << processId() << ":" << systemThreadId() << " "
+       << processId(false) << ":" << systemThreadId(false) << " "
        << (file ? file + 1 : filePath) << ":" << line << "] ";
 }
 
diff --git a/libkineto/src/ThreadUtil.cpp b/libkineto/src/ThreadUtil.cpp
index 3fdc22b74..49e67411b 100644
--- a/libkineto/src/ThreadUtil.cpp
+++ b/libkineto/src/ThreadUtil.cpp
@@ -39,45 +39,58 @@ thread_local int32_t _tid = 0;
 thread_local int32_t _sysTid = 0;
 }
 
-int32_t processId() {
+int32_t processId(bool cache) {
+  int32_t pid = 0;
   if (!_pid) {
 #ifndef _WIN32
-    _pid = (int32_t)getpid();
+    pid = (int32_t)getpid();
 #else
-    _pid = (int32_t)GetCurrentProcessId();
+    pid = (int32_t)GetCurrentProcessId();
 #endif
+    if (cache) {
+      _pid = pid;
+    }
+    return pid;
   }
   return _pid;
 }
 
-int32_t systemThreadId() {
+int32_t systemThreadId(bool cache) {
+  int32_t sysTid = 0;
   if (!_sysTid) {
 #ifdef __APPLE__
-    _sysTid = (int32_t)syscall(SYS_thread_selfid);
+    sysTid = (int32_t)syscall(SYS_thread_selfid);
 #elif defined _WIN32
-    _sysTid = (int32_t)GetCurrentThreadId();
+    sysTid = (int32_t)GetCurrentThreadId();
 #elif defined __FreeBSD__
-    syscall(SYS_thr_self, &_sysTid);
+    syscall(SYS_thr_self, &sysTid);
 #else
-    _sysTid = (int32_t)syscall(SYS_gettid);
+    sysTid = (int32_t)syscall(SYS_gettid);
 #endif
+    if (cache) {
+      _sysTid = sysTid;
+    }
+    return sysTid;
   }
   return _sysTid;
 }
 
-int32_t threadId() {
+int32_t threadId(bool cache) {
+  int64_t tid = 0;
   if (!_tid) {
 #ifdef __APPLE__
-    uint64_t tid;
     pthread_threadid_np(nullptr, &tid);
-    _tid = tid;
 #elif defined _WIN32
-  _tid = (int32_t)GetCurrentThreadId();
+  tid = (int32_t)GetCurrentThreadId();
 #else
   pthread_t pth = pthread_self();
   int32_t* ptr = reinterpret_cast<int32_t*>(&pth);
-  _tid = *ptr;
+  tid = *ptr;
 #endif
+    if (cache) {
+      _tid = (int32_t)tid;
+    }
+    return (int32_t)tid;
   }
   return _tid;
 }