Use struct with pid and Go routine addr for Go BPF maps

bpf/go_common.h

@@ -37,14 +37,19 @@ char __license[] SEC("license") = "Dual MIT/GPL";
 // Then it is retrieved in the return uprobes and used to know the HTTP call duration as well as its
 // attributes (method, path, and status code).
 
+typedef struct goroutine_key {
+    u32 pid;  // PID of the process
+    u64 addr; // Address of the goroutine
+} goroutine_key_t;
+
 typedef struct goroutine_metadata_t {
-    u64 parent;
+    goroutine_key_t parent;
     u64 timestamp;
 } goroutine_metadata;
 
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __type(key, void *);               // key: pointer to the goroutine
+    __type(key, goroutine_key_t);      // key: pointer to the goroutine
     __type(value, goroutine_metadata); // value: timestamp of the goroutine creation
     __uint(max_entries, MAX_CONCURRENT_SHARED_REQUESTS);
     __uint(pinning, LIBBPF_PIN_BY_NAME);
@@ -73,18 +78,31 @@ struct {
     __uint(pinning, LIBBPF_PIN_BY_NAME);
 } go_trace_map SEC(".maps");
 
-static __always_inline u64 find_parent_goroutine(void *goroutine_addr) {
-    void *r_addr = goroutine_addr;
+static __always_inline void goroutine_key_from_id(goroutine_key_t *current, void *goroutine) {
+    u64 pid_tid = bpf_get_current_pid_tgid();
+    u32 pid = pid_from_pid_tgid(pid_tid);
+
+    current->addr = (u64)goroutine;
+    current->pid = pid;
+}
+
+static __always_inline u64 find_parent_goroutine(goroutine_key_t *current) {
+    if (!current) {
+        return 0;
+    }
+
+    u64 r_addr = current->addr;
+
     int attempts = 0;
     do {
-        void *p_inv = bpf_map_lookup_elem(&go_trace_map, &r_addr);
+        tp_info_t *p_inv = bpf_map_lookup_elem(&go_trace_map, &r_addr);
         if (!p_inv) { // not this goroutine running the server request processing
             // Let's find the parent scope
             goroutine_metadata *g_metadata =
-                (goroutine_metadata *)bpf_map_lookup_elem(&ongoing_goroutines, &r_addr);
+                (goroutine_metadata *)bpf_map_lookup_elem(&ongoing_goroutines, current);
             if (g_metadata) {
                 // Lookup now to see if the parent was a request
-                r_addr = (void *)g_metadata->parent;
+                r_addr = (u64)g_metadata->parent.addr;
             } else {
                 break;
             }
@@ -202,8 +220,10 @@ static __always_inline u8 client_trace_parent(void *goroutine_addr,
 
     if (!found_trace_id) {
         tp_info_t *tp = 0;
+        goroutine_key_t g_key = {};
+        goroutine_key_from_id(&g_key, goroutine_addr);
 
-        u64 parent_id = find_parent_goroutine(goroutine_addr);
+        u64 parent_id = find_parent_goroutine(&g_key);
 
         if (parent_id) { // we found a parent request
             tp = (tp_info_t *)bpf_map_lookup_elem(&go_trace_map, &parent_id);

bpf/go_grpc.h

@@ -199,7 +199,10 @@ int uprobe_server_handler_transport_handle_streams(struct pt_regs *ctx) {
                tr,
                goroutine_addr);
 
-    void *parent_go = (void *)find_parent_goroutine(goroutine_addr);
+    goroutine_key_t g_key = {};
+    goroutine_key_from_id(&g_key, goroutine_addr);
+
+    void *parent_go = (void *)find_parent_goroutine(&g_key);
     if (parent_go) {
         bpf_dbg_printk("found parent goroutine for transport handler [%llx]", parent_go);
         connection_info_t *conn = bpf_map_lookup_elem(&ongoing_server_connections, &parent_go);
@@ -253,10 +256,12 @@ int uprobe_server_handleStream_return(struct pt_regs *ctx) {
     trace->content_length = 0;
     trace->method[0] = 0;
 
-    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &goroutine_addr);
+    goroutine_key_t g_key = {.addr = (u64)goroutine_addr, .pid = trace->pid.host_pid};
+
+    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &g_key);
     if (g_metadata) {
         trace->go_start_monotime_ns = g_metadata->timestamp;
-        bpf_map_delete_elem(&ongoing_goroutines, &goroutine_addr);
+        bpf_map_delete_elem(&ongoing_goroutines, &g_key);
     } else {
         trace->go_start_monotime_ns = invocation->start_monotime_ns;
     }

bpf/go_nethttp.h

@@ -201,16 +201,21 @@ int uprobe_readRequestReturns(struct pt_regs *ctx) {
     void *goroutine_addr = GOROUTINE_PTR(ctx);
     bpf_dbg_printk("goroutine_addr %lx", goroutine_addr);
 
+    u64 pid_tid = bpf_get_current_pid_tgid();
+    u32 pid = pid_from_pid_tgid(pid_tid);
+
+    goroutine_key_t g_key = {.addr = (u64)goroutine_addr, .pid = pid};
+
     // This code is here for keepalive support on HTTP requests. Since the connection is not
     // established everytime, we set the initial goroutine start on the new read initiation.
-    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &goroutine_addr);
+    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &g_key);
     if (!g_metadata) {
         goroutine_metadata metadata = {
             .timestamp = bpf_ktime_get_ns(),
-            .parent = (u64)goroutine_addr,
+            .parent = g_key,
         };
 
-        if (bpf_map_update_elem(&ongoing_goroutines, &goroutine_addr, &metadata, BPF_ANY)) {
+        if (bpf_map_update_elem(&ongoing_goroutines, &g_key, &metadata, BPF_ANY)) {
             bpf_dbg_printk("can't update active goroutine");
         }
     }
@@ -229,7 +234,10 @@ int uprobe_ServeHTTPReturns(struct pt_regs *ctx) {
         bpf_map_lookup_elem(&ongoing_http_server_requests, &goroutine_addr);
 
     if (invocation == NULL) {
-        void *parent_go = (void *)find_parent_goroutine(goroutine_addr);
+        goroutine_key_t g_key = {};
+        goroutine_key_from_id(&g_key, goroutine_addr);
+
+        void *parent_go = (void *)find_parent_goroutine(&g_key);
         if (parent_go) {
             bpf_dbg_printk("found parent goroutine for header [%llx]", parent_go);
             invocation = bpf_map_lookup_elem(&ongoing_http_server_requests, &parent_go);
@@ -252,10 +260,12 @@ int uprobe_ServeHTTPReturns(struct pt_regs *ctx) {
     trace->start_monotime_ns = invocation->start_monotime_ns;
     trace->end_monotime_ns = bpf_ktime_get_ns();
 
-    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &goroutine_addr);
+    goroutine_key_t g_key = {.addr = (u64)goroutine_addr, .pid = trace->pid.host_pid};
+
+    goroutine_metadata *g_metadata = bpf_map_lookup_elem(&ongoing_goroutines, &g_key);
     if (g_metadata) {
         trace->go_start_monotime_ns = g_metadata->timestamp;
-        bpf_map_delete_elem(&ongoing_goroutines, &goroutine_addr);
+        bpf_map_delete_elem(&ongoing_goroutines, &g_key);
     } else {
         trace->go_start_monotime_ns = invocation->start_monotime_ns;
     }
@@ -510,7 +520,10 @@ int uprobe_http2ResponseWriterStateWriteHeader(struct pt_regs *ctx) {
         bpf_map_lookup_elem(&ongoing_http_server_requests, &goroutine_addr);
 
     if (invocation == NULL) {
-        void *parent_go = (void *)find_parent_goroutine(goroutine_addr);
+        goroutine_key_t g_key = {};
+        goroutine_key_from_id(&g_key, goroutine_addr);
+
+        void *parent_go = (void *)find_parent_goroutine(&g_key);
         if (parent_go) {
             bpf_dbg_printk("found parent goroutine for header [%llx]", parent_go);
             invocation = bpf_map_lookup_elem(&ongoing_http_server_requests, &parent_go);

bpf/go_runtime.h

@@ -62,12 +62,17 @@ int uprobe_proc_newproc1_ret(struct pt_regs *ctx) {
     void *goroutine_addr = (void *)GO_PARAM1(ctx);
     bpf_dbg_printk("goroutine_addr %lx", goroutine_addr);
 
+    u64 pid_tid = bpf_get_current_pid_tgid();
+    u32 pid = pid_from_pid_tgid(pid_tid);
+
+    goroutine_key_t g_key = {.addr = (u64)goroutine_addr, .pid = pid};
+
     goroutine_metadata metadata = {
         .timestamp = bpf_ktime_get_ns(),
-        .parent = (u64)parent_goroutine,
+        .parent = {.addr = (u64)parent_goroutine, .pid = pid},
     };
 
-    if (bpf_map_update_elem(&ongoing_goroutines, &goroutine_addr, &metadata, BPF_ANY)) {
+    if (bpf_map_update_elem(&ongoing_goroutines, &g_key, &metadata, BPF_ANY)) {
         bpf_dbg_printk("can't update active goroutine");
     }
 
@@ -84,7 +89,12 @@ int uprobe_proc_goexit1(struct pt_regs *ctx) {
     void *goroutine_addr = GOROUTINE_PTR(ctx);
     bpf_dbg_printk("goroutine_addr %lx", goroutine_addr);
 
-    bpf_map_delete_elem(&ongoing_goroutines, &goroutine_addr);
+    u64 pid_tid = bpf_get_current_pid_tgid();
+    u32 pid = pid_from_pid_tgid(pid_tid);
+
+    goroutine_key_t g_key = {.addr = (u64)goroutine_addr, .pid = pid};
+
+    bpf_map_delete_elem(&ongoing_goroutines, &g_key);
     // We also clean-up the go routine based trace map, it's an LRU
     // but at this point we are sure we don't need the data.
     bpf_map_delete_elem(&go_trace_map, &goroutine_addr);