[AGENT] Some hacks to make pyperf work

agent/src/ebpf/Makefile

@@ -75,6 +75,7 @@ OBJS := user/elf.o \
 	user/symbol.o \
 	user/go_tracer.o \
 	user/ssl_tracer.o \
+	user/python_probe.o \
 	user/ring.o \
 	user/btf_vmlinux.o \
 	user/load.o \

agent/src/ebpf/kernel/Makefile

@@ -48,6 +48,8 @@ EBPF_CLAGS ?= -I. -Ivmlinux -Iinclude -D__BPF_TRACING__ -D GROUP_LEADER_OFFSET_O
         -DSTART_BOOTTIME_OFFSET_OVERRIDE=0      \
         -DSTART_BOOTTIME_VARNAME=real_start_time
 
+socket_trace.bpf.c: uprobe_base.bpf.c go_tls.bpf.c go_http2.bpf.c openssl.bpf.c python_uprobe.bpf.c
+
 %.elf: %.c
 	$(call msg,BPF,$@,$(CORE))	
 	$(Q)$(CLANG) $(EBPF_CLAGS) $(EXTRA_EBPF_CLAGS) -std=gnu99 -Wimplicit-function-declaration \

agent/src/ebpf/kernel/perf_profiler.bpf.c

@@ -72,9 +72,9 @@ MAP_PERARRAY(heap, __u32, unwind_state_t, 1)
 MAP_PERARRAY(python_symbol_index, __u32, __u32, 1)
 MAP_HASH(python_stack, __u64, stack_trace_t, STACK_MAP_ENTRIES)
 
-static inline __attribute__((always_inline)) bool comm_eq(char *a, char *b) {
+static inline __attribute__((always_inline)) bool comm_eq_n(char *a, char *b, int n) {
 #pragma unroll
-	for (int i = 0; i < TASK_COMM_LEN; i++) {
+	for (int i = 0; i < TASK_COMM_LEN && i < n; i++) {
 		if (a[i] == '\0' || b[i] == '\0') {
 			return a[i] == b[i];
 		}
@@ -419,31 +419,47 @@ int bpf_perf_event(struct bpf_perf_event_data *ctx)
 	/*
 	 * CPU idle stacks will not be collected.
 	 */
-	if (key->tgid == key->pid && key->pid == 0)
+	if (key->tgid == key->pid && key->pid == 0) {
 		return 0;
+	}
 
 	key->cpu = bpf_get_smp_processor_id();
 	bpf_get_current_comm(&key->comm, sizeof(key->comm));
 	key->timestamp = bpf_ktime_get_ns();
 
-	if (comm_eq(key->comm, "python3")) {
+	if (comm_eq_n(key->comm, "python3", 7) || comm_eq_n(key->comm, "pt_main", 7)) {
+		__builtin_memcpy(key->comm, "python3", 7);
+		key->comm[7] = '\0';
 		bpf_tail_call(ctx, &NAME(progs_jmp_perf_map), PROG_PYTHON_FRAME_PTR_IDX);
 	}
 
 	return get_stack_and_output_perf(ctx, state);
 }
 
+MAP_HASH(python_thread_state_map, __u32, __u64, 65536)
+
 PROGPE(python_frame_ptr)(struct bpf_perf_event_data *ctx) {
 	__u32 zero = 0;
 	unwind_state_t *state = heap__lookup(&zero);
 	if (state == NULL) {
 		return 0;
 	}
 
-	__u64 thread_state_addr = 140737353904984;
+	// __u64 thread_state_addr = 140737353904984;
+	__u64 thread_state_addr = 9851320;
 	if (bpf_probe_read_user(&state->thread_state, sizeof(void *), (void *)thread_state_addr) != 0) {
 		goto finish;
 	}
+	if (state->thread_state != NULL) {
+		python_thread_state_map__update(&state->key.tgid, (__u64 *)&state->thread_state);
+	} else {
+		__u64 *entry = python_thread_state_map__lookup(&state->key.tgid);
+		if (entry) {
+			state->thread_state = (void *)*entry;
+		} else {
+			goto finish;
+		}
+	}
 
 	if (bpf_probe_read_user(&state->key.itid, sizeof(__u32), state->thread_state + py_offsets.py_thread_state.thread_id) != 0) {
 		goto finish;
@@ -456,7 +472,7 @@ PROGPE(python_frame_ptr)(struct bpf_perf_event_data *ctx) {
 	bpf_tail_call(ctx, &NAME(progs_jmp_perf_map), PROG_PYTHON_WALK_STACK_IDX);
 
 finish:
-	return 0;
+	return get_stack_and_output_perf(ctx, state);
 }
 
 PROGPE(python_walk_stack)(struct bpf_perf_event_data *ctx) {

agent/src/ebpf/kernel/python_uprobe.bpf.c

@@ -0,0 +1,17 @@
+#if 0
+SEC("uprobe/python_save_thread_state_address")
+int uprobe_python_save_thread_state_address(struct pt_regs *ctx) {
+	__u64 zero = 0;
+	python_thread_state_map__update(&zero, &zero);
+	bpf_debug("PyEval_SaveThread");
+	return 0;
+}
+
+SEC("uprobe/pyeval_evalframedefault")
+int uprobe_pyeval_evalframedefault(struct pt_regs *ctx) {
+	__u64 zero = 0;
+	python_thread_state_map__update(&zero, &zero);
+	bpf_debug("_PyEval_EvalFrameDefault");
+	return 0;
+}
+#endif

agent/src/ebpf/kernel/socket_trace.bpf.c

@@ -2595,3 +2595,4 @@ PROGTP(io_event) (void *ctx) {
 #include "go_tls.bpf.c"
 #include "go_http2.bpf.c"
 #include "openssl.bpf.c"
+#include "python_uprobe.bpf.c"

agent/src/ebpf/user/log.c

@@ -95,7 +95,7 @@ void _ebpf_error(int how_to_die, char *function_name, char *file_path,
 {
 	char msg[MSG_SZ] = {};
 	uint16_t len = 0;
-	uint16_t max = MSG_SZ;
+	int16_t max = MSG_SZ;
 	va_list va;
 
 	if (function_name) {
@@ -130,7 +130,7 @@ void _ebpf_info(char *fmt, ...)
 {
 	char msg[MSG_SZ] = {};
 	uint16_t len = 0;
-	uint16_t max = MSG_SZ;
+	int16_t max = MSG_SZ;
 	va_list va;
 
 	len += snprintf(msg + len, max - len, "[eBPF] INFO ");

agent/src/ebpf/user/profile/stringifier.c

@@ -553,6 +553,58 @@ static inline char *alloc_stack_trace_str(int len)
 	return trace_str;
 }
 
+bool is_python_to_c(char *fname) {
+	return strcmp(fname, "PyCFunction_Call") == 0;
+}
+
+bool is_c_to_python(char *fname) {
+	return strcmp(fname, "PyVectorcall_Call") == 0;
+}
+
+void merge_interpreter_and_user_trace(char *trace_str, int len, char *i_trace, char *u_trace)
+{
+	char *i_sp = NULL, *u_sp = NULL, *i_frame = NULL, *u_frame = NULL;
+	int offset = 0;
+	bool in_c = false, seen_eval_frame = false;
+
+	// ebpf_warning("\n%s\n%s", i_trace, u_trace);
+	// <module>
+	i_frame = strtok_r(i_trace, ";", &i_sp);
+	offset += snprintf(trace_str + offset, len - offset, "%s;", i_frame);
+
+	i_frame = strtok_r(NULL, ";", &i_sp);
+	u_frame = strtok_r(u_trace, ";", &u_sp);
+	while (u_frame) {
+		if (in_c) {
+			if (is_c_to_python(u_frame)) {
+				in_c = false;
+				seen_eval_frame = true;
+			} else {
+				offset += snprintf(trace_str + offset, len - offset, "%s;", u_frame);
+			}
+		} else {
+			if (strcmp(u_frame, "_PyEval_EvalFrameDefault") == 0) {
+				if (seen_eval_frame && i_frame) {
+					offset += snprintf(trace_str + offset, len - offset, "%s;", i_frame);
+					i_frame = strtok_r(NULL, ";", &i_sp);
+				}
+				seen_eval_frame = true;
+			} else if (is_python_to_c(u_frame)) {
+				in_c = true;
+				seen_eval_frame = false;
+			}
+		}
+		u_frame = strtok_r(NULL, ";", &u_sp);
+	}
+	if (in_c) {
+		while (u_frame) {
+			offset += snprintf(trace_str + offset, len - offset, "%s;", u_frame);
+			u_frame = strtok_r(NULL, ";", &u_sp);
+		}
+	}
+	trace_str[offset - 1] = '\0';
+}
+
 char *resolve_and_gen_stack_trace_str(struct bpf_tracer *t,
 				      struct stack_trace_key_t *v,
 				      const char *stack_map_name,
@@ -667,6 +719,14 @@ char *resolve_and_gen_stack_trace_str(struct bpf_tracer *t,
 		return trace_str;
 	}
 
+	if (v->intpstack != 0) {
+		if (i_trace_str) {
+			len += strlen(i_trace_str);
+		} else {
+			len += strlen(i_err_tag);
+		}
+	}
+
 	if (v->kernstack >= 0) {
 		if (k_trace_str) {
 			len += strlen(k_trace_str);
@@ -683,14 +743,6 @@ char *resolve_and_gen_stack_trace_str(struct bpf_tracer *t,
 		}
 	}
 
-	if (v->intpstack != 0) {
-		if (i_trace_str) {
-			len += strlen(i_trace_str);
-		} else {
-			len += strlen(i_err_tag);
-		}
-	}
-
 	trace_str = alloc_stack_trace_str(len);
 	if (trace_str == NULL) {
 		ebpf_warning("No available memory space.\n");
@@ -700,16 +752,22 @@ char *resolve_and_gen_stack_trace_str(struct bpf_tracer *t,
 	int offset = 0;
 	bool last_stack = false;
 
-	if (v->userstack >= 0) {
-		offset += snprintf(trace_str + offset, len - offset, "%s%s", last_stack ? ";" : "", u_trace_str ? u_trace_str : u_err_tag);
-		last_stack = true;
-	}
-	if (v->intpstack != 0) {
-		offset += snprintf(trace_str + offset, len - offset, "%s%s", last_stack ? ";" : "", i_trace_str ? i_trace_str : i_err_tag);
+	if (i_trace_str && u_trace_str) {
+		merge_interpreter_and_user_trace(trace_str + offset, len - offset, i_trace_str, u_trace_str);
 		last_stack = true;
+	} else {
+		if (v->intpstack != 0) {
+			offset += snprintf(trace_str + offset, len - offset, "%s%s", last_stack ? ";" : "", i_trace_str ? i_trace_str : i_err_tag);
+			last_stack = true;
+		}
+		if (v->userstack >= 0) {
+			offset += snprintf(trace_str + offset, len - offset, "%s%s", last_stack ? ";" : "", u_trace_str ? u_trace_str : u_err_tag);
+			last_stack = true;
+		}
 	}
 	if (v->kernstack >= 0) {
 		offset += snprintf(trace_str + offset, len - offset, "%s%s", last_stack ? ";" : "", k_trace_str ? k_trace_str : k_err_tag);
+		last_stack = true;
 	}
 
 	return trace_str;