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

bpf/go_common.h

@@ -32,60 +32,80 @@ 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 {
+    u64 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);
 } ongoing_goroutines SEC(".maps");
 
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __type(key, void *); // key: pointer to the request goroutine
+    __type(key, goroutine_key_t); // key: pointer to the request goroutine
     __type(value, connection_info_t);
     __uint(max_entries, MAX_CONCURRENT_SHARED_REQUESTS);
     __uint(pinning, LIBBPF_PIN_BY_NAME);
 } ongoing_server_connections SEC(".maps");
 
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __type(key, void *); // key: pointer to the request goroutine
+    __type(key, goroutine_key_t); // key: pointer to the request goroutine
     __type(value, connection_info_t);
     __uint(max_entries, MAX_CONCURRENT_REQUESTS);
 } ongoing_client_connections SEC(".maps");
 
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __type(key, void *);      // key: pointer to the goroutine
-    __type(value, tp_info_t); // value: traceparent info
+    __type(key, goroutine_key_t); // key: pointer to the goroutine
+    __type(value, tp_info_t);     // value: traceparent info
     __uint(max_entries, MAX_CONCURRENT_SHARED_REQUESTS);
     __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;
+    goroutine_key_t *parent = current;
+
     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, parent);
         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, parent);
             if (g_metadata) {
                 // Lookup now to see if the parent was a request
-                r_addr = (void *)g_metadata->parent;
+                r_addr = g_metadata->parent.addr;
+                parent = &g_metadata->parent;
             } else {
                 break;
             }
         } else {
             bpf_dbg_printk("Found parent %lx", r_addr);
-            return (u64)r_addr;
+            return r_addr;
         }
 
         attempts++;
@@ -131,6 +151,8 @@ server_trace_parent(void *goroutine_addr, tp_info_t *tp, void *req_header) {
     tp->flags = 1;
     // Get traceparent from the Request.Header
     void *traceparent_ptr = extract_traceparent_from_req_headers(req_header);
+    goroutine_key_t g_key = {};
+    goroutine_key_from_id(&g_key, goroutine_addr);
     if (traceparent_ptr != NULL) {
         unsigned char buf[TP_MAX_VAL_LENGTH];
         long res = bpf_probe_read(buf, sizeof(buf), traceparent_ptr);
@@ -143,7 +165,7 @@ server_trace_parent(void *goroutine_addr, tp_info_t *tp, void *req_header) {
             decode_go_traceparent(buf, tp->trace_id, tp->parent_id, &tp->flags);
         }
     } else {
-        connection_info_t *info = bpf_map_lookup_elem(&ongoing_server_connections, &goroutine_addr);
+        connection_info_t *info = bpf_map_lookup_elem(&ongoing_server_connections, &g_key);
         u8 found_info = 0;
 
         if (info) {
@@ -169,7 +191,7 @@ server_trace_parent(void *goroutine_addr, tp_info_t *tp, void *req_header) {
     }
 
     urand_bytes(tp->span_id, SPAN_ID_SIZE_BYTES);
-    bpf_map_update_elem(&go_trace_map, &goroutine_addr, tp, BPF_ANY);
+    bpf_map_update_elem(&go_trace_map, &g_key, tp, BPF_ANY);
 
     unsigned char tp_buf[TP_MAX_VAL_LENGTH];
     make_tp_string(tp_buf, tp);
@@ -201,11 +223,15 @@ 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);
+        goroutine_key_t p_key = {};
+        goroutine_key_from_id(&p_key, (void *)parent_id);
 
         if (parent_id) { // we found a parent request
-            tp = (tp_info_t *)bpf_map_lookup_elem(&go_trace_map, &parent_id);
+            tp = (tp_info_t *)bpf_map_lookup_elem(&go_trace_map, &p_key);
         }
 
         if (tp) {