Refactor http2_grpc to leverage tail calls

grafana · Dec 13, 2024 · 380fbc1 · 380fbc1
1 parent 85b6cdd
commit 380fbc1
Show file tree

Hide file tree

Showing 31 changed files with 1,254 additions and 761 deletions.
diff --git a/bpf/http2_grpc.h b/bpf/http2_grpc.h
@@ -8,10 +8,6 @@
 
 #define HTTP2_GRPC_PREFACE "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
 
-#define FRAME_HEADER_LEN 9
-
-#define FLAG_DATA_END_STREAM 0x1
-
 typedef enum {
     FrameData = 0x0,
     FrameHeaders = 0x1,
@@ -33,10 +29,14 @@ typedef struct frame_header {
     u32 stream_id : 31;
 } __attribute__((packed)) frame_header_t;
 
+enum { k_flag_data_end_stream = 0x1, k_frame_header_len = 9 };
+
+_Static_assert(sizeof(frame_header_t) == k_frame_header_len, "frame_header_t size mismatch");
+
 static __always_inline u8 read_http2_grpc_frame_header(frame_header_t *frame,
                                                        const unsigned char *p,
                                                        u32 len) {
-    if (len < FRAME_HEADER_LEN) {
+    if (len < k_frame_header_len) {
         return 0;
     }
 
@@ -61,7 +61,7 @@ static __always_inline u8 is_settings_frame(unsigned char *p, u32 len) {
     return frame.type == FrameSettings && !frame.stream_id;
 }
 
-static __always_inline u8 is_headers_frame(frame_header_t *frame) {
+static __always_inline u8 is_headers_frame(const frame_header_t *frame) {
     return frame->type == FrameHeaders && frame->stream_id;
 }
 
@@ -87,20 +87,20 @@ static __always_inline u8 is_http2_or_grpc(unsigned char *p, u32 len) {
     return has_preface(p, len) || is_settings_frame(p, len);
 }
 
-static __always_inline u8 http_grpc_stream_ended(frame_header_t *frame) {
+static __always_inline u8 http_grpc_stream_ended(const frame_header_t *frame) {
     return is_headers_frame(frame) &&
-           ((frame->flags & FLAG_DATA_END_STREAM) == FLAG_DATA_END_STREAM);
+           ((frame->flags & k_flag_data_end_stream) == k_flag_data_end_stream);
 }
 
-static __always_inline u8 is_invalid_frame(frame_header_t *frame) {
+static __always_inline u8 is_invalid_frame(const frame_header_t *frame) {
     return frame->length == 0 && frame->type == FrameData;
 }
 
-static __always_inline u8 is_data_frame(frame_header_t *frame) {
+static __always_inline u8 is_data_frame(const frame_header_t *frame) {
     return frame->length && frame->type == FrameData;
 }
 
-static __always_inline u8 is_flags_only_frame(frame_header_t *frame) {
+static __always_inline u8 is_flags_only_frame(const frame_header_t *frame) {
     return frame->length <= 2;
 }
 

diff --git a/bpf/k_tracer_defs.h b/bpf/k_tracer_defs.h
@@ -4,6 +4,7 @@
 #include "vmlinux.h"
 #include "bpf_helpers.h"
 #include "http_types.h"
+#include "k_tracer_tailcall.h"
 #include "ringbuf.h"
 #include "pid.h"
 #include "trace_common.h"
@@ -25,13 +26,6 @@ typedef struct recv_args {
     u8 iovec_ctx[sizeof(iovec_iter_ctx)];
 } recv_args_t;
 
-struct bpf_map_def SEC("maps") jump_table = {
-    .type = BPF_MAP_TYPE_PROG_ARRAY,
-    .key_size = sizeof(__u32),
-    .value_size = sizeof(__u32),
-    .max_entries = 8,
-};
-
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __type(key, egress_key_t);
@@ -40,24 +34,20 @@ struct {
     __uint(pinning, BEYLA_PIN_INTERNAL);
 } msg_buffers SEC(".maps");
 
-#define TAIL_PROTOCOL_HTTP 0
-#define TAIL_PROTOCOL_HTTP2 1
-#define TAIL_PROTOCOL_TCP 2
-
 static __always_inline void handle_buf_with_args(void *ctx, call_protocol_args_t *args) {
     bpf_dbg_printk(
         "buf=[%s], pid=%d, len=%d", args->small_buf, args->pid_conn.pid, args->bytes_len);
 
     if (is_http(args->small_buf, MIN_HTTP_SIZE, &args->packet_type)) {
-        bpf_tail_call(ctx, &jump_table, TAIL_PROTOCOL_HTTP);
+        bpf_tail_call(ctx, &jump_table, k_tail_protocol_http);
     } else if (is_http2_or_grpc(args->small_buf, MIN_HTTP2_SIZE)) {
         bpf_dbg_printk("Found HTTP2 or gRPC connection");
         u8 is_ssl = args->ssl;
         bpf_map_update_elem(&ongoing_http2_connections, &args->pid_conn, &is_ssl, BPF_ANY);
     } else {
         u8 *h2g = bpf_map_lookup_elem(&ongoing_http2_connections, &args->pid_conn);
         if (h2g && *h2g == args->ssl) {
-            bpf_tail_call(ctx, &jump_table, TAIL_PROTOCOL_HTTP2);
+            bpf_tail_call(ctx, &jump_table, k_tail_protocol_http2);
         } else { // large request tracking
             http_info_t *info = bpf_map_lookup_elem(&ongoing_http, &args->pid_conn);
 
@@ -70,7 +60,7 @@ static __always_inline void handle_buf_with_args(void *ctx, call_protocol_args_t
                     // Essentially, when a packet is extended by our sock_msg program and
                     // passed down another service, the receiving side may reassemble the
                     // packets into one buffer or not. If they are reassembled, then the
-                    // call to bpf_tail_call(ctx, &jump_table, TAIL_PROTOCOL_HTTP); will
+                    // call to bpf_tail_call(ctx, &jump_table, k_tail_protocol_http); will
                     // scan for the incoming 'Traceparent' header. If they are not reassembled
                     // we'll see something like this:
                     // [before the injected header],[70 bytes for 'Traceparent...'],[the rest].
@@ -109,7 +99,7 @@ static __always_inline void handle_buf_with_args(void *ctx, call_protocol_args_t
                 // we are processing SSL request. HTTP2 is already checked in handle_buf_with_connection.
                 http_info_t *http_info = bpf_map_lookup_elem(&ongoing_http, &args->pid_conn);
                 if (!http_info) {
-                    bpf_tail_call(ctx, &jump_table, TAIL_PROTOCOL_TCP);
+                    bpf_tail_call(ctx, &jump_table, k_tail_protocol_tcp);
                 }
             }
         }

diff --git a/bpf/k_tracer_tailcall.h b/bpf/k_tracer_tailcall.h
@@ -0,0 +1,23 @@
+#ifndef K_TRACER_TAILCALL_H
+#define K_TRACER_TAILCALL_H
+
+#include "vmlinux.h"
+#include "bpf_helpers.h"
+
+struct bpf_map_def SEC("maps") jump_table = {
+    .type = BPF_MAP_TYPE_PROG_ARRAY,
+    .key_size = sizeof(__u32),
+    .value_size = sizeof(__u32),
+    .max_entries = 8,
+};
+
+enum {
+    k_tail_protocol_http = 0,
+    k_tail_protocol_http2 = 1,
+    k_tail_protocol_tcp = 2,
+    k_tail_protocol_http2_grpc_frames = 3,
+    k_tail_protocol_http2_grpc_handle_start_frame = 4,
+    k_tail_protocol_http2_grpc_handle_end_frame = 5,
+};
+
+#endif
diff --git a/bpf/protocol_http.h b/bpf/protocol_http.h
@@ -307,7 +307,7 @@ static __always_inline void handle_http_response(unsigned char *small_buf,
     bpf_map_delete_elem(&active_ssl_connections, pid_conn);
 }
 
-// TAIL_PROTOCOL_HTTP
+// k_tail_protocol_http
 SEC("kprobe/http")
 int protocol_http(void *ctx) {
     call_protocol_args_t *args = protocol_args();