feat: add TPBASE and TSD extraction for multi-thread Python unwind

Author: kylewanginchina

agent/crates/trace-utils/cbindgen.toml

@@ -71,7 +71,8 @@ include = [
   "PhpUnwindTable",
   "V8UnwindInfo",
   "V8Offsets",
-  "V8UnwindTable"
+  "V8UnwindTable",
+  "TSDInfo"
 ]
 exclude = [
   "bpf_update_elem",
@@ -91,6 +92,7 @@ ProcessShardList = "process_shard_list_t"
 UnwindEntry = "unwind_entry_t"
 UnwindEntryShard = "unwind_entry_shard_t"
 UnwindTable = "unwind_table_t"
+TSDInfo = "tsd_info_t"
 PyCframe = "py_cframe_t"
 PyCodeObject = "py_code_object_t"
 PyFrameObject = "py_frame_object_t"

agent/crates/trace-utils/src/error.rs

@@ -37,6 +37,8 @@ pub enum Error {
     InvalidPointer(u64),
     #[error("Invalid or corrupted data")]
     InvalidData,
+    #[error("{0}")]
+    Msg(String),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;

agent/crates/trace-utils/src/lib.rs

@@ -18,6 +18,7 @@ pub mod btf;
 pub mod error;
 pub mod maps;
 pub mod remote_memory;
+pub mod tsd;
 pub mod unwind;
 pub mod utils;
 pub use utils::protect_cpu_affinity;
@@ -29,6 +30,7 @@ use std::io::Write;
 use log::info;
 
 // Crate internal modules
+pub use tsd::TSDInfo;
 use unwind::UnwindTable;
 
 // ============================================================================

agent/crates/trace-utils/src/maps.rs

@@ -21,7 +21,7 @@ use std::path::PathBuf;
 
 use log::trace;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct MemoryArea {
     pub m_start: u64,
     pub mx_start: u64, // start address of executable section

agent/crates/trace-utils/src/trace_utils.h

@@ -173,8 +173,41 @@ typedef struct {
 
 #if defined(DF_ENTERPRISE)
 typedef struct {
-    uint64_t thread_state_address;
+    /**
+     * Offset from thread pointer base (TPBASE) to TSD storage
+     */
+    int16_t offset;
+    /**
+     * TSD key multiplier (glibc=16, musl=8)
+     */
+    uint8_t multiplier;
+    /**
+     * Whether indirect addressing is needed (musl=1, glibc=0)
+     */
+    uint8_t indirect;
+} tsd_info_t;
+
+typedef struct {
+    /**
+     * Address of autoTLSkey variable in Python runtime
+     */
+    uint64_t auto_tls_key_addr;
+    /**
+     * Python version encoded as 0xMMmm (e.g., 0x030A for 3.10)
+     */
+    uint16_t version;
+    /**
+     * Thread Specific Data info for multi-threading support
+     */
+    tsd_info_t tsd_info;
+    /**
+     * ID for looking up python_offsets in the offsets map
+     */
     uint8_t offsets_id;
+    /**
+     * Padding for alignment
+     */
+    uint8_t _padding[5];
 } python_unwind_info_t;
 #endif
 
@@ -709,6 +742,12 @@ void rust_log_wrapper(LogLevel level,
                       const char *file_path,
                       int line_number);
 
+/**
+ * Read TPBASE offset from kernel functions
+ * Returns the offset of fsbase/tpidr in task_struct, or -1 on failure
+ */
+int64_t read_tpbase_offset(void);
+
 int rustc_demangle(const char *mangled, char *out, size_t out_size);
 
 unwind_table_t *unwind_table_create(int32_t process_shard_list_map_fd,

agent/crates/trace-utils/src/tsd.rs

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) 2024 Yunshan Networks
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/// Thread Specific Data info for accessing per-thread PyThreadState.
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Default)]
+pub struct TSDInfo {
+    /// Offset from thread pointer base (TPBASE) to TSD storage
+    pub offset: i16,
+    /// TSD key multiplier (glibc=16, musl=8)
+    pub multiplier: u8,
+    /// Whether indirect addressing is needed (musl=1, glibc=0)
+    pub indirect: u8,
+}

agent/crates/trace-utils/src/unwind.rs

@@ -16,6 +16,7 @@
 
 pub mod dwarf;
 pub mod elf_utils;
+pub mod tpbase;
 
 use std::alloc::{alloc, dealloc, handle_alloc_error, Layout};
 use std::collections::{hash_map::Entry, HashMap, HashSet};