From ab0b382dfe8cd22e70af549d4f63e99dcd4b18b7 Mon Sep 17 00:00:00 2001
From: changlehung <changlehung@gmail.com>
Date: Sat, 24 May 2025 23:06:24 +0800
Subject: [PATCH 1/3] Add update_tls to fix tls info

---
 src/viztracer/modules/snaptrace.c | 48 +++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
index 772b52c2..c07cadef 100644
--- a/src/viztracer/modules/snaptrace.c
+++ b/src/viztracer/modules/snaptrace.c
@@ -175,6 +175,24 @@ clear_stack(struct FunctionNode** stack_top) {
 // =============================================================================
 // Thread info related functions
 // =============================================================================
+static unsigned long
+get_thread_id()
+{
+#if _WIN32  
+    return GetCurrentThreadId();
+#elif defined(__APPLE__)
+    __uint64_t tid = 0;
+    if (pthread_threadid_np(NULL, &tid)) {
+        return (unsigned long)pthread_self();
+    } else {
+        return tid;
+    }
+#elif defined(__FreeBSD__)
+    return pthread_getthreadid_np();
+#else
+    return syscall(SYS_gettid);
+#endif
+}
 
 static struct ThreadInfo*
 snaptrace_createthreadinfo(TracerObject* self) {
@@ -1964,6 +1982,35 @@ tracer_get_sync_marker(TracerObject* self, PyObject* Py_UNUSED(unused))
     return PyFloat_FromDouble(ts_sync_marker);
 }
 
+static PyObject* 
+tracer_update_tls(TracerObject* self, PyObject* Py_UNUSED(unused))
+{
+    struct ThreadInfo* info = get_thread_info(self);
+    unsigned long old_tid = info->tid;
+    info->tid = get_thread_id();
+
+    struct EventNode* curr = self->buffer + self->buffer_head_idx;
+    while (curr != self->buffer + self->buffer_tail_idx) {
+        struct EventNode* node = curr;
+        if (node->tid == old_tid) {
+            node->tid = info->tid;
+        }
+        curr = curr + 1;
+        if (curr == self->buffer + self->buffer_size) {
+            curr = self->buffer;
+        }
+    }
+
+    struct MetadataNode* node = self->metadata_head;
+    while (node != NULL) {
+        if (node->tid == old_tid) {
+            node->tid = info->tid;
+        }
+        node = node->next;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef Tracer_methods[] = {
     {"threadtracefunc", (PyCFunction)tracer_threadtracefunc, METH_VARARGS, "trace function"},
     {"start", (PyCFunction)tracer_start, METH_NOARGS, "start profiling"},
@@ -1986,6 +2033,7 @@ static PyMethodDef Tracer_methods[] = {
     {"setignorestackcounter", (PyCFunction)tracer_setignorestackcounter, METH_O, "reset ignore stack depth"},
     {"set_sync_marker", (PyCFunction)tracer_set_sync_marker, METH_NOARGS, "set current timestamp to synchronization marker"},
     {"get_sync_marker", (PyCFunction)tracer_get_sync_marker, METH_NOARGS, "get synchronization marker or None if not set"},
+    {"update_tls", (PyCFunction)tracer_update_tls, METH_NOARGS, "fix thread local storage"},
     {NULL, NULL, 0, NULL}
 };
 

From eef2dd9fad811cd4ed48955ee979225b4a3d5e1d Mon Sep 17 00:00:00 2001
From: changlehung <changlehung@gmail.com>
Date: Sat, 24 May 2025 23:06:53 +0800
Subject: [PATCH 2/3] Add update_tls interface

---
 src/viztracer/snaptrace.pyi | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/viztracer/snaptrace.pyi b/src/viztracer/snaptrace.pyi
index e5f2074d..d759c3c3 100644
--- a/src/viztracer/snaptrace.pyi
+++ b/src/viztracer/snaptrace.pyi
@@ -75,3 +75,6 @@ class Tracer:
     def get_sync_marker(self) -> Optional[float]:
         """get synchronization marker or None if not set"""
         ...
+
+    def update_tls(self) -> None:
+        ...

From 47d86a5e29a4439eaef1479bb0504037b4619d00 Mon Sep 17 00:00:00 2001
From: changlehung <changlehung@gmail.com>
Date: Sat, 24 May 2025 23:08:02 +0800
Subject: [PATCH 3/3] Fix tls in multiprocessing programs

---
 src/viztracer/patch.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/viztracer/patch.py b/src/viztracer/patch.py
index d8897ef3..b4302bf3 100644
--- a/src/viztracer/patch.py
+++ b/src/viztracer/patch.py
@@ -119,6 +119,7 @@ def func_after_fork(tracer: VizTracer):
 
         tracer.clear()
         tracer.reset_stack()
+        tracer.update_tls()
 
         if tracer._afterfork_cb:
             tracer._afterfork_cb(tracer, *tracer._afterfork_args, **tracer._afterfork_kwargs)
@@ -292,6 +293,7 @@ def callback():
                     import multiprocessing.util
                     multiprocessing.util.Finalize(tracer, tracer.exit_routine, exitpriority=-1)
                     tracer.label_file_to_write()
+                    tracer.update_tls()
             os.register_at_fork(after_in_child=callback)  # type: ignore
 
         if tracer.log_audit is not None: