Tracking the NodeJS event loop

bpf/http_sock.c

@@ -121,7 +121,6 @@ int BPF_KPROBE(kprobe_tcp_rcv_established, struct sock *sk, struct sk_buff *skb)
         // If the source port for a client call is lower, we'll get this wrong.
         // TODO: Need to fix this. 
         pid_info.orig_dport = pid_info.p_conn.conn.s_port,
-        task_tid(&pid_info.c_tid);
         bpf_map_update_elem(&pid_tid_to_conn, &id, &pid_info, BPF_ANY); // to support SSL on missing handshake, respect the original info if there
     }
 
@@ -167,7 +166,6 @@ int BPF_KRETPROBE(kretprobe_sys_accept4, uint fd)
         sort_connection_info(&info.p_conn.conn);
         info.p_conn.pid = pid_from_pid_tgid(id);
         info.orig_dport = orig_dport;
-        task_tid(&info.c_tid);
 
         bpf_map_update_elem(&pid_tid_to_conn, &id, &info, BPF_ANY); // to support SSL on missing handshake
     }
@@ -234,7 +232,6 @@ int BPF_KRETPROBE(kretprobe_sys_connect, int fd)
         sort_connection_info(&info.p_conn.conn);
         info.p_conn.pid = pid_from_pid_tgid(id);
         info.orig_dport = orig_dport;
-        task_tid(&info.c_tid);
 
         bpf_map_update_elem(&pid_tid_to_conn, &id, &info, BPF_ANY); // to support SSL 
     }
@@ -336,7 +333,6 @@ int BPF_KPROBE(kprobe_tcp_sendmsg, struct sock *sk, struct msghdr *msg, size_t s
             .orig_dport = orig_dport,
         };
         bpf_memcpy(&ssl_conn.p_conn, &s_args.p_conn, sizeof(pid_connection_info_t));
-        task_tid(&ssl_conn.c_tid);
         bpf_map_update_elem(&ssl_to_conn, &ssl, &ssl_conn, BPF_ANY);
     }
 
@@ -628,8 +624,8 @@ int BPF_KPROBE(kprobe_sys_exit, int status) {
         return 0;
     }
 
-    pid_key_t task = {0};
-    task_tid(&task);
+    trace_key_t task = {0};
+    task_tid(&task.p_key);
 
     bpf_dbg_printk("sys_exit %d, pid=%d, valid_pid(id)=%d", id, pid_from_pid_tgid(id), valid_pid(id));
 
@@ -641,7 +637,9 @@ int BPF_KPROBE(kprobe_sys_exit, int status) {
         bpf_map_delete_elem(&active_ssl_connections, &s_args->p_conn);
     }
 
-    bpf_map_delete_elem(&clone_map, &task);
+    bpf_map_delete_elem(&clone_map, &task.p_key);
+    // This won't delete trace ids for traces with extra_id, like NodeJS. But, 
+    // we expect that it doesn't matter, since NodeJS main thread won't exit. 
     bpf_map_delete_elem(&server_traces, &task);
 
     return 0;

bpf/http_ssl.h

@@ -82,7 +82,12 @@ static __always_inline void cleanup_ssl_trace_info(http_info_t *info, void *ssl)
         if (ssl_info) {
             bpf_dbg_printk("Looking to delete server trace for ssl = %llx, info->type = %d", ssl, info->type);
             //dbg_print_http_connection_info(&ssl_info->conn.conn); // commented out since GitHub CI doesn't like this call
-            delete_server_trace_tid(&ssl_info->c_tid);
+            trace_key_t t_key = {0};
+            t_key.extra_id = info->extra_id;
+            t_key.p_key.ns = info->pid.ns;
+            t_key.p_key.pid = info->task_tid;
+
+            delete_server_trace(&t_key);
         }
     }
 
@@ -165,7 +170,6 @@ static __always_inline void handle_ssl_buf(void *ctx, u64 id, ssl_args_t *args,
             bpf_memcpy(&p_c.p_conn.conn.s_addr, &ssl, sizeof(void *));
             p_c.p_conn.conn.d_port = p_c.p_conn.conn.s_port = p_c.orig_dport = 0;
             p_c.p_conn.pid = pid_from_pid_tgid(id);
-            task_tid(&p_c.c_tid);
 
             bpf_map_update_elem(&ssl_to_conn, &ssl, &p_c, BPF_ANY);
             conn = bpf_map_lookup_elem(&ssl_to_conn, &ssl);

bpf/http_types.h

@@ -62,7 +62,6 @@ typedef struct http_pid_connection_info {
 typedef struct ssl_pid_connection_info {
     pid_connection_info_t p_conn;
     u16 orig_dport;
-    pid_key_t c_tid;
 } ssl_pid_connection_info_t;
 
 typedef struct tp_info {
@@ -96,6 +95,8 @@ typedef struct http_info {
     // with other instrumented processes
     pid_info pid;
     tp_info_t tp;
+    u64 extra_id;
+    u32 task_tid;
 } http_info_t;
 
 // Here we track unknown TCP requests that are not HTTP, HTTP2 or gRPC

bpf/nodejs.c

@@ -0,0 +1,91 @@
+#include "vmlinux.h"
+#include "bpf_helpers.h"
+#include "bpf_dbg.h"
+#include "pid.h"
+#include "ringbuf.h"
+#include "nodejs.h"
+
+char __license[] SEC("license") = "Dual MIT/GPL";
+
+volatile const s32 async_wrap_async_id_off = 0;
+volatile const s32 async_wrap_trigger_async_id_off = 0;
+
+struct {
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __type(key, u64);   // the pid_tid 
+    __type(value, u64); // the last AsyncWrap *
+    __uint(max_entries, 1000); // 1000 nodejs services, small number, nodejs is single threaded
+    __uint(pinning, LIBBPF_PIN_BY_NAME);
+} async_reset_args SEC(".maps");
+
+SEC("uprobe/node:AsyncReset")
+int async_reset(struct pt_regs *ctx) {
+    u64 id = bpf_get_current_pid_tgid();
+
+    if (!valid_pid(id)) {
+        return 0;
+    }
+
+    u64 wrap = (u64)PT_REGS_PARM1(ctx);
+
+    bpf_dbg_printk("=== uprobe AsyncReset id=%d wrap=%llx ===", id, wrap);
+    bpf_map_update_elem(&async_reset_args, &id, &wrap, BPF_ANY);
+
+    return 0;
+}
+
+SEC("uretprobe/node:AsyncReset")
+int async_reset_ret(struct pt_regs *ctx) {
+    u64 id = bpf_get_current_pid_tgid();
+
+    if (!valid_pid(id)) {
+        return 0;
+    }
+
+    bpf_dbg_printk("=== uprobe AsyncReset returns id=%d ===", id);
+    bpf_map_delete_elem(&async_reset_args, &id);
+
+    return 0;
+}
+
+SEC("uprobe/node:EmitAsyncInit")
+int emit_async_init(struct pt_regs *ctx) {
+    u64 id = bpf_get_current_pid_tgid();
+
+    if (!valid_pid(id)) {
+        return 0;
+    }
+
+    bpf_dbg_printk("=== uprobe EmitAsyncInit id=%d ===", id);
+
+    u64 *wrap_val = bpf_map_lookup_elem(&async_reset_args, &id);
+    bpf_dbg_printk("wrap_val = %llx", wrap_val);
+    if (wrap_val) {
+        u64 wrap = *wrap_val;
+        bpf_dbg_printk("wrap = %llx", wrap);
+
+        if (wrap) {
+            u64 async_id = 0;
+            u64 trigger_async_id = 0;
+
+            bpf_probe_read_user(&async_id, sizeof(u64), ((void *)wrap) + async_wrap_async_id_off);
+            bpf_probe_read_user(&trigger_async_id, sizeof(u64), ((void *)wrap) + async_wrap_trigger_async_id_off);
+
+            if (async_id) {
+                bpf_map_update_elem(&active_nodejs_ids, &id, &async_id, BPF_ANY);
+                if (trigger_async_id) {
+                    bpf_map_update_elem(&nodejs_parent_map, &async_id, &trigger_async_id, BPF_ANY);
+                    bpf_dbg_printk("async_id = %llx, trigger_async_id = %llx", async_id, trigger_async_id);
+                } else {
+                    bpf_dbg_printk("No trigger async id");
+                }
+            } else {
+                bpf_dbg_printk("No async id");
+            }
+        }
+    } else {
+        bpf_dbg_printk("No wrap value found");
+    }
+
+    return 0;
+}

bpf/nodejs.h

@@ -0,0 +1,25 @@
+#ifndef NODE_JS_H
+#define NODE_JS_H
+
+#include "vmlinux.h"
+#include "bpf_helpers.h"
+#include "bpf_builtins.h"
+#include "map_sizing.h"
+
+struct {
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __type(key, u64);   // the pid_tid 
+    __type(value, u64); // the last active async_id
+    __uint(max_entries, 1000); // 1000 nodejs services, small number, nodejs is single threaded
+    __uint(pinning, LIBBPF_PIN_BY_NAME);
+} active_nodejs_ids SEC(".maps");
+
+struct {
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __type(key, u64);   // child async_id
+    __type(value, u64); // parent async_id
+    __uint(max_entries, MAX_CONCURRENT_REQUESTS); 
+    __uint(pinning, LIBBPF_PIN_BY_NAME);
+} nodejs_parent_map SEC(".maps");
+
+#endif

bpf/pid_types.h

@@ -55,6 +55,19 @@ static __always_inline void task_pid(pid_info *pid) {
     pid->ns = (u32)ns.inum;
 }
 
+static __always_inline u32 get_task_tid() {
+    struct upid upid;
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+
+    // https://github.com/torvalds/linux/blob/556e2d17cae620d549c5474b1ece053430cd50bc/kernel/pid.c#L324 (type is )
+    // set user-side PID
+    unsigned int level = BPF_CORE_READ(task, nsproxy, pid_ns_for_children, level);
+    struct pid *ns_pid = (struct pid *)BPF_CORE_READ(task, thread_pid);
+    bpf_probe_read_kernel(&upid, sizeof(upid), &ns_pid->numbers[level]);
+
+    return (u32)upid.nr;
+}
+
 static __always_inline void task_tid(pid_key_t *tid) {
     struct upid upid;
     struct task_struct *task = (struct task_struct *)bpf_get_current_task();

bpf/protocol_http.h

@@ -7,6 +7,7 @@
 #include "http_types.h"
 #include "ringbuf.h"
 #include "pid.h"
+#include "runtime.h"
 #include "protocol_common.h"
 
 // http_info_t became too big to be declared as a variable in the stack.
@@ -166,6 +167,8 @@ static __always_inline void process_http_request(http_info_t *info, int len, htt
     info->start_monotime_ns = bpf_ktime_get_ns();
     info->status = 0;
     info->len = len;
+    info->extra_id = extra_runtime_id(); // required for deleting the trace information
+    info->task_tid = get_task_tid();     // required for deleting the trace information
 }
 
 static __always_inline void process_http_response(http_info_t *info, unsigned char *buf, int len) {
@@ -195,7 +198,11 @@ static __always_inline void handle_http_response(unsigned char *small_buf, pid_c
     }
 
     if (info->type == EVENT_HTTP_REQUEST) {
-        delete_server_trace();
+        trace_key_t t_key = {0};
+        t_key.extra_id = info->extra_id;
+        t_key.p_key.ns = info->pid.ns;
+        t_key.p_key.pid = info->task_tid;
+        delete_server_trace(&t_key);
     } else {
         //bpf_dbg_printk("Deleting client trace map for connection");
         //dbg_print_http_connection_info(&pid_conn->conn);

bpf/runtime.h

@@ -0,0 +1,30 @@
+#ifndef RUNTIME_SUPPORT_H
+#define RUNTIME_SUPPORT_H
+
+#include "vmlinux.h"
+#include "bpf_helpers.h"
+#include "bpf_builtins.h"
+#include "pid_types.h"
+#include "nodejs.h"
+
+static __always_inline u64 extra_runtime_id() {
+    u64 id = bpf_get_current_pid_tgid();
+
+    u64 *active_node_id = (u64 *)bpf_map_lookup_elem(&active_nodejs_ids, &id);
+    if (active_node_id) {
+        return *active_node_id;
+    }
+
+    return 0;
+}
+
+static __always_inline u64 parent_runtime_id(pid_key_t *p_key, u64 runtime_id) {
+    u64 *parent_id = (u64 *)bpf_map_lookup_elem(&nodejs_parent_map, &runtime_id);
+    if (parent_id) {
+        return *parent_id;
+    }
+
+    return 0;
+}
+
+#endif