add kprobe for destroy sock to improve map cleanup

DataDog · Oct 31, 2024 · 76e86c6 · 76e86c6
1 parent a414970
commit 76e86c6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 1 deletion.
diff --git a/pkg/network/ebpf/c/tracer.c b/pkg/network/ebpf/c/tracer.c
@@ -275,7 +275,7 @@ int BPF_BYPASSABLE_KPROBE(kprobe__tcp_close, struct sock *sk) {
     skp_conn_tuple_t skp_conn = {.sk = sk, .tup = t};
     skp_conn.tup.pid = 0;
 
-    bpf_map_delete_elem(&tcp_ongoing_connect_pid, &skp_conn);
+    // bpf_map_delete_elem(&tcp_ongoing_connect_pid, &skp_conn);
 
     if (!tcp_failed_connections_enabled()) {
         cleanup_conn(ctx, &t, sk);
@@ -324,6 +324,27 @@ int BPF_KRETPROBE(kretprobe__tcp_close_flush) {
     return 0;
 }
 
+SEC("kprobe/tcp_v4_destroy_sock")
+int BPF_BYPASSABLE_KPROBE(kprobe__tcp_v4_destroy_sock, struct sock *sk) {
+    conn_tuple_t t = {};
+    u64 pid_tgid = bpf_get_current_pid_tgid();
+
+    // Get network namespace id
+    log_debug("adamk kprobe/tcp_v4_destroy_sock: tgid: %llu, pid: %llu", pid_tgid >> 32, pid_tgid & 0xFFFFFFFF);
+    if (!read_conn_tuple(&t, sk, pid_tgid, CONN_TYPE_TCP)) {
+        return 0;
+    }
+    log_debug("adamk kprobe/tcp_v4_destroy_sock: netns: %u, sport: %u, dport: %u", t.netns, t.sport, t.dport);
+
+    bpf_map_delete_elem(&conn_close_flushed, &t);
+
+    skp_conn_tuple_t skp_conn = {.sk = sk, .tup = t};
+    skp_conn.tup.pid = 0;
+
+    bpf_map_delete_elem(&tcp_ongoing_connect_pid, &skp_conn);
+    return 0;
+}
+
 #if !defined(COMPILE_RUNTIME) || defined(FEATURE_UDPV6_ENABLED)
 
 static __always_inline void fl6_saddr(struct flowi6 *fl6, u64 *addr_h, u64 *addr_l) {

diff --git a/pkg/network/ebpf/probes/probes.go b/pkg/network/ebpf/probes/probes.go
@@ -82,6 +82,8 @@ const (
 
 	// TCPClose traces the tcp_close() system call
 	TCPClose ProbeFuncName = "kprobe__tcp_close"
+	// TCPDestroySock does stuff
+	TCPDestroySock ProbeFuncName = "kprobe__tcp_v4_destroy_sock"
 	// TCPDone traces the tcp_done() system call
 	TCPDone ProbeFuncName = "kprobe__tcp_done"
 	// TCPDoneFlushReturn traces the return of the tcp_done() system call

diff --git a/pkg/network/tracer/connection/kprobe/config.go b/pkg/network/tracer/connection/kprobe/config.go
@@ -76,6 +76,7 @@ func enabledProbes(c *config.Config, runtimeTracer, coreTracer bool) (map[probes
 		enableProbe(enabled, probes.TCPReadSock)
 		enableProbe(enabled, probes.TCPReadSockReturn)
 		enableProbe(enabled, probes.TCPClose)
+		enableProbe(enabled, probes.TCPDestroySock)
 		enableProbe(enabled, probes.TCPCloseFlushReturn)
 		enableProbe(enabled, probes.TCPConnect)
 		enableProbe(enabled, probes.TCPDone)

diff --git a/pkg/network/tracer/connection/kprobe/manager.go b/pkg/network/tracer/connection/kprobe/manager.go
@@ -31,6 +31,7 @@ var mainProbes = []probes.ProbeFuncName{
 	probes.TCPReadSock,
 	probes.TCPReadSockReturn,
 	probes.TCPClose,
+	probes.TCPDestroySock,
 	probes.TCPDone,
 	probes.TCPDoneFlushReturn,
 	probes.TCPCloseCleanProtocolsReturn,