Merge pull request #162 from coroot/partial_zk_reads

handle partial Zookeeper reads
coroot · Dec 25, 2024 · 1dbf538 · 1dbf538
2 parents 1da58fa + 8ed2571
commit 1dbf538
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 15 deletions.
diff --git a/ebpftracer/ebpf.go b/ebpftracer/ebpf.go
diff --git a/ebpftracer/ebpf/l7/l7.c b/ebpftracer/ebpf/l7/l7.c
@@ -477,7 +477,11 @@ int trace_exit_read(void *ctx, __u64 id, __u32 pid, __u16 is_tls, long int ret)
             return 0; // keeping the query in the map
         }
     } else if (e->protocol == PROTOCOL_ZOOKEEPER) {
-        response = is_zk_response(payload, total_size, &e->status);
+        response = is_zk_response(payload, total_size, &e->status, req->partial);
+        if (response == 2) { // partial
+            req->partial = 1;
+            return 0; // keeping the query in the map
+        }
     } else if (e->protocol == PROTOCOL_DUBBO2) {
         response = is_dubbo2_response(payload, &e->status);
     }

diff --git a/ebpftracer/ebpf/l7/zookeeper.c b/ebpftracer/ebpf/l7/zookeeper.c
@@ -49,11 +49,19 @@ int is_zk_request(char *buf, __u64 buf_size) {
 }
 
 static __always_inline
-int is_zk_response(char *buf, __u64 buf_size, __s32 *status) {
+int is_zk_response(char *buf, __u64 buf_size, __s32 *status, __u8 partial) {
+    if (partial == 0 && buf_size == 4) { //partial read
+        return 2;
+    }
     struct zk_response_header resp = {};
-    bpf_read(buf, resp);
-    if (bpf_ntohl(resp.length)+4 != buf_size) {
-        return 0;
+    if (partial) {
+        bpf_read(buf, resp.xid);
+        bpf_read(buf+12, resp.err_code);
+    } else {
+        bpf_read(buf, resp);
+        if (bpf_ntohl(resp.length)+4 != buf_size) {
+            return 0;
+        }
     }
     __s32 xid = bpf_ntohl(resp.xid);
     if (xid < 0 && xid != -1 && xid != -2) {