[buf] fix multi-thread stack depth problem

linuxkerneltravel · Aug 12, 2023 · 644e732 · 644e732
1 parent 184d632
commit 644e732
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 58 deletions.
diff --git a/eBPF_Supermarket/User_Function_Tracer/README.md b/eBPF_Supermarket/User_Function_Tracer/README.md
@@ -38,4 +38,3 @@ Examples:
 ### TODO
 - simplify c++ symbols
 - more tests
-- multithread
diff --git a/eBPF_Supermarket/User_Function_Tracer/src/utrace.bpf.c b/eBPF_Supermarket/User_Function_Tracer/src/utrace.bpf.c
@@ -26,68 +26,96 @@
 
 char LICENSE[] SEC("license") = "Dual BSD/GPL";
 
+static int global_sz;
+static int current_pid;
 struct {
   __uint(type, BPF_MAP_TYPE_HASH);
-  __uint(max_entries, 8192);
-  __type(key, s32);
-  __type(value, u64);
+  __type(key, __u32);
+  __type(value, __u64);
+  __uint(max_entries, MAX_STACK_DEPTH);
 } function_start SEC(".maps");
 
+struct {
+  __uint(type, BPF_MAP_TYPE_HASH);
+  __type(key, __u32);
+  __type(value, __u32);
+  __uint(max_entries, MAX_THREAD_NUM);
+} stack_depth SEC(".maps");
+
 struct {
   __uint(type, BPF_MAP_TYPE_RINGBUF);
-  __uint(max_entries, 8192 * 1024);
+  __uint(max_entries, 2048 * 4096);
 } records SEC(".maps");
 
 SEC("uprobe/trace")
 int BPF_KPROBE(uprobe) {
   struct profile_record *r;
-  pid_t pid, tid, cpu_id;
-  u64 id, ts;
+  pid_t tid, cpu_id;
+  __u64 ts;
+  __u32 *depth_ptr;
+  __u32 depth;
 
-  id = bpf_get_current_pid_tgid();
-  tid = (u32)id;
-  pid = id >> 32;
+  current_pid = tid = (__u32)bpf_get_current_pid_tgid();
   cpu_id = bpf_get_smp_processor_id();
 
   r = bpf_ringbuf_reserve(&records, sizeof(*r), 0);
   if (!r) return 0;
+
   r->tid = tid;
   r->cpu_id = cpu_id;
+  r->timestamp = 0;
   r->duration_ns = 0;
-  r->ustack_sz =
-      bpf_get_stack(ctx, r->ustack, sizeof(r->ustack), BPF_F_USER_STACK) / sizeof(r->ustack[0]);
+  bpf_get_stack(ctx, &r->ustack, sizeof(r->ustack), BPF_F_USER_STACK);
+  depth_ptr = bpf_map_lookup_elem(&stack_depth, &tid);
+  if (depth_ptr)
+    r->ustack_sz = depth = *depth_ptr + 1;
+  else
+    r->ustack_sz = depth = 0;
+  r->global_sz = global_sz;
   r->exit = 0;
-  s32 ustack_sz = r->ustack_sz;
-  ts = bpf_ktime_get_ns();
-  bpf_map_update_elem(&function_start, &ustack_sz, &ts, BPF_ANY);
+
   bpf_ringbuf_submit(r, 0);
+
+  ts = bpf_ktime_get_ns();
+  bpf_map_update_elem(&function_start, &global_sz, &ts, BPF_NOEXIST);
+  ++global_sz;
+  bpf_map_update_elem(&stack_depth, &tid, &depth, BPF_ANY);
   return 0;
 }
 
 SEC("uretprobe/trace")
 int BPF_KRETPROBE(uretprobe) {
   struct profile_record *r;
-  pid_t pid, tid, cpu_id;
-  u64 id, end_ts = bpf_ktime_get_ns();
-  u64 *start_ts;
-  s32 ustack_sz;
+  pid_t tid, cpu_id;
+  __u64 *start_ts_ptr;
+  __u64 end_ts = bpf_ktime_get_ns();
+  __u32 *depth_ptr;
 
-  id = bpf_get_current_pid_tgid();
+  current_pid = tid = (__u32)bpf_get_current_pid_tgid();
   cpu_id = bpf_get_smp_processor_id();
-  pid = id >> 32;
-  tid = (u32)id;
+
+  depth_ptr = bpf_map_lookup_elem(&stack_depth, &tid);
+  if (!depth_ptr) return 0;
 
   r = bpf_ringbuf_reserve(&records, sizeof(*r), 0);
   if (!r) return 0;
+
+  --global_sz;
+
   r->tid = tid;
   r->cpu_id = cpu_id;
-  r->ustack_sz = bpf_get_stack(ctx, r->ustack, sizeof(r->ustack), BPF_F_USER_STACK) /
-                 sizeof(sizeof(r->ustack[0]));
-  ustack_sz = r->ustack_sz;
-  start_ts = bpf_map_lookup_elem(&function_start, &ustack_sz);
-  if (start_ts) r->duration_ns = end_ts - *start_ts;
-  bpf_map_delete_elem(&function_start, &ustack_sz);
+  r->timestamp = 0;
+  start_ts_ptr = bpf_map_lookup_elem(&function_start, &global_sz);
+  if (start_ts_ptr) r->duration_ns = end_ts - *start_ts_ptr;
+  bpf_get_stack(ctx, &r->ustack, sizeof(r->ustack), BPF_F_USER_STACK);
+  r->ustack_sz = *depth_ptr;
+  r->global_sz = global_sz;
   r->exit = 1;
+
   bpf_ringbuf_submit(r, 0);
+
+  bpf_map_delete_elem(&function_start, &global_sz);
+  --*depth_ptr;
+  bpf_map_update_elem(&stack_depth, &tid, depth_ptr, BPF_ANY);
   return 0;
 }
diff --git a/eBPF_Supermarket/User_Function_Tracer/src/utrace.c b/eBPF_Supermarket/User_Function_Tracer/src/utrace.c
@@ -170,65 +170,61 @@ bool symbolize(size_t addr) {
 }
 
 static int handle_event(void *ctx, void *data, size_t data_sz) {
+  static struct profile_record pending_r;
+  static int pending = 0;
+  static int status = -1; /**< 0: exec 1: exit */
+
   const struct profile_record *r = data;
 
   if (r->exit) {
     // LOG("EXIT");
     // for (int i = 0; i < r->ustack_sz; i++) LOG(" %llx", r->ustack[i]);
     // LOG(" %s\n", stack_func[r->ustack_sz]);
     // return 0;
-    if (status == 0 && r->ustack_sz == pre_ustack_sz) {
-      log_cpuid(r->cpu_id);
+    if (status == 0) {
+      log_cpuid(pending_r.cpu_id);
       log_split();
-      log_tid(r->tid);
+      log_tid(pending_r.tid);
       log_split();
       log_time(r->duration_ns);
       log_split();
-      log_char(' ', 2 * r->ustack_sz - 2);
-      LOG("%s();\n", stack_func[r->ustack_sz]);
-      status = 1;
-      pre_ustack_sz = r->ustack_sz;
-    } else if (status == 1 && r->ustack_sz == pre_ustack_sz - 1) {
+      log_char(' ', 2 * pending_r.ustack_sz);
+      LOG("%s();\n", stack_func[pending_r.global_sz]);
+    } else {  // status == 1
       log_cpuid(r->cpu_id);
       log_split();
       log_tid(r->tid);
       log_split();
       log_time(r->duration_ns);
       log_split();
-      log_char(' ', 2 * r->ustack_sz - 2);
-      LOG("} ");
-      log_color(TERM_GRAY);
-      LOG("/* %s */", stack_func[r->ustack_sz]);
-      log_color(TERM_NC);
-      LOG("\n");
-      status = 1;
-      pre_ustack_sz = r->ustack_sz;
+      log_char(' ', 2 * r->ustack_sz);
+      LOG("} " TERM_GRAY "/* %s */" TERM_NC "\n", stack_func[r->global_sz]);
     }
+    pending = 0;
+    status = 1;
   } else {
     // LOG("EXEC");
     // for (int i = 0; i < r->ustack_sz; i++) LOG(" %llx", r->ustack[i]);
     // if (symbolize(r->ustack[0])) {
-    //   memcpy(stack_func[r->ustack_sz], buf, sizeof(buf));
+    //    memcpy(stack_func[r->ustack_sz], buf, sizeof(buf));
     // }
     // LOG(" %s\n", stack_func[r->ustack_sz]);
     // return 0;
-    if (status == -1 && r->ustack_sz != 1) return 0;
-    if (status == 0 && r->ustack_sz != pre_ustack_sz + 1) return 0;
-    if (status == 1 && r->ustack_sz != pre_ustack_sz) return 0;
     if (status == 0) {
-      log_cpuid(r->cpu_id);
+      log_cpuid(pending_r.cpu_id);
       log_split();
-      log_tid(r->tid);
+      log_tid(pending_r.tid);
       log_split();
       log_char(' ', 11);
       log_split();
-      log_char(' ', 2 * r->ustack_sz - 4);
-      LOG("%s() {\n", stack_func[r->ustack_sz - 1]);
+      log_char(' ', 2 * pending_r.ustack_sz);
+      LOG("%s() {\n", stack_func[pending_r.global_sz]);
     }
-    if (symbolize(r->ustack[0])) {
-      memcpy(stack_func[r->ustack_sz], buf, sizeof(buf));
+    if (symbolize(r->ustack)) {
+      memcpy(stack_func[r->global_sz], buf, sizeof(buf));
+      pending_r = *r;
+      pending = 1;
       status = 0;
-      pre_ustack_sz = r->ustack_sz;
     }
   }
   return 0;

diff --git a/eBPF_Supermarket/User_Function_Tracer/src/utrace.h b/eBPF_Supermarket/User_Function_Tracer/src/utrace.h
@@ -19,20 +19,23 @@
 #ifndef UTRACE_UTRACE_H
 #define UTRACE_UTRACE_H
 
+#define MAX_THREAD_NUM 32
 #define MAX_STACK_DEPTH 128
 #define MAX_MANGLED_LEN 128
 #define MAX_SYMBOL_LEN 1024
 #define MAX_PATH_LEN 256
 
-typedef unsigned long long stack_trace_t[MAX_STACK_DEPTH];
+typedef unsigned long long stack_trace_t;
 
 /**
  * @brief 内核态传给用户态的数据
  */
 struct profile_record {
-  unsigned int tid;    /**< 线程编号 */
-  unsigned int cpu_id; /**< CPU编号 */
+  unsigned int tid;      /**< 线程编号 */
+  unsigned int next_tid; /**< 切换后的线程编号 */
+  unsigned int cpu_id;   /**< CPU编号 */
 
+  unsigned long long timestamp;   /**< 时间戳 */
   unsigned long long duration_ns; /**< 函数时延 */
 
   unsigned int kstack_sz; /**< 内核栈大小 */
@@ -41,6 +44,8 @@ struct profile_record {
   unsigned int ustack_sz; /**< 用户栈大小 */
   stack_trace_t ustack;   /**< 用户栈 */
 
+  unsigned int global_sz; /**< 当前函数深度（考虑了多线程） */
+
   int exit; /**< 是否为函数退出时 */
 };