Feature #TASK_QT-18250 grpc-clien端到端数据trace_id与server数据采集到的trace_id相关联。

ebpftracer/ebpf/utrace/go/net/grpc.client.probe.bpf.c

@@ -44,7 +44,7 @@ struct {
 struct {
     __uint(type, BPF_MAP_TYPE_HASH);
     __type(key, u32);
-    __type(value, struct span_context);
+    __type(value, struct apm_trace_key_t);
     __uint(max_entries, MAX_CONCURRENT);
 } streamid_to_span_contexts SEC(".maps");
 
@@ -94,7 +94,7 @@ volatile const bool write_status_supported;
 // func (cc *ClientConn) Invoke(ctx context.Context, method string, args, reply interface{}, opts ...CallOption) error
 SEC("uprobe/ClientConn_Invoke")
 int uprobe_ClientConn_Invoke(struct pt_regs *ctx) {
-    // bpf_printk("enter the uprobe_ClientConn_Invoke \n");
+    bpf_printk("enter the uprobe_ClientConn_Invoke \n");
     // positions
     u64 clientconn_pos = 1;
     u64 method_ptr_pos = 4;
@@ -110,6 +110,7 @@ int uprobe_ClientConn_Invoke(struct pt_regs *ctx) {
 
     // Get key
     void *key = (void *)GOROUTINE(ctx);
+    bpf_printk("enter the uprobe_ClientConn_Invoke key is 0x%llx\n", (u64)key);
     void *grpcReq_ptr = bpf_map_lookup_elem(&grpc_client_events, &key);
     if (grpcReq_ptr != NULL) {
         bpf_printk("uprobe/ClientConn_Invoke already tracked with the current context");
@@ -169,7 +170,7 @@ int uprobe_ClientConn_Invoke(struct pt_regs *ctx) {
 // func (cc *ClientConn) Invoke(ctx context.Context, method string, args, reply interface{}, opts ...CallOption) error
 SEC("uprobe/ClientConn_Invoke")
 int uprobe_ClientConn_Invoke_Returns(struct pt_regs *ctx) {
-    // bpf_printk("enter the uprobe_ClientConn_Invoke_Returns \n");
+    bpf_printk("enter the uprobe_ClientConn_Invoke_Returns \n");
     void *key = (void *)GOROUTINE(ctx);
     struct grpc_client_request_t *grpc_span = bpf_map_lookup_elem(&grpc_client_events, &key);
     if (grpc_span == NULL) {
@@ -370,8 +371,12 @@ cw_append_item_to_slice(void *new_item, u32 item_size, void *slice_user_ptr) {
 
 SEC("uprobe/loopyWriter_headerHandler")
 int uprobe_LoopyWriter_HeaderHandler(struct pt_regs *ctx) {
+
+    void *key = (void *)GOROUTINE(ctx);
+    bpf_printk("enter the uprobe_LoopyWriter_HeaderHandler key is 0x%llx\n", (u64)key);
+
     void *headerFrame_ptr = get_argument(ctx, 2);
-    // bpf_printk("enter the get header handler storage\n");
+    bpf_printk("enter the get header handler storage\n");
 
     __u64 pid_tgid = bpf_get_current_pid_tgid();
 	__u32 tgid = pid_tgid >> 32;
@@ -381,6 +386,14 @@ int uprobe_LoopyWriter_HeaderHandler(struct pt_regs *ctx) {
     {
         return 0;
     }
+
+    u32 stream_id = 0;
+    bpf_probe_read(
+        &stream_id, sizeof(stream_id), (void *)(headerFrame_ptr + (headerFrame_streamid_pos)));
+    struct apm_trace_key_t *sc_ptr = bpf_map_lookup_elem(&streamid_to_span_contexts, &stream_id);
+    if (sc_ptr == NULL) {
+        return 0;
+    }
     
     // Get storage from per-cpu map
     u32 zero = 0;
@@ -407,7 +420,8 @@ int uprobe_LoopyWriter_HeaderHandler(struct pt_regs *ctx) {
     grpcClientReq->start_time = bpf_ktime_get_ns();
 
     // bpf_printk("enter the uprobe_LoopyWriter_HeaderHandler444444\n");
-	struct apm_span_context *cw_psc = cw_get_parent_tracking_span();
+	// struct apm_span_context *cw_psc = cw_get_parent_tracking_span();
+    struct apm_span_context *cw_psc = cw_get_parent_tracking_span_by_trace_key(*sc_ptr);
 	if(cw_psc){
 		bpf_probe_read(&grpcClientReq->apm_psc, sizeof(grpcClientReq->apm_psc), cw_psc);
 		copy_byte_arrays(grpcClientReq->apm_psc.trace_id, grpcClientReq->apm_sc.trace_id, APM_TRACE_ID_SIZE);
@@ -476,19 +490,17 @@ int uprobe_LoopyWriter_HeaderHandler(struct pt_regs *ctx) {
 SEC("uprobe/http2Client_NewStream")
 // func (t *http2Client) NewStream(ctx context.Context, callHdr *CallHdr) (*Stream, error)
 int uprobe_http2Client_NewStream(struct pt_regs *ctx) {
-    // bpf_printk("enter the uprobe_http2Client_NewStream \n");
-    // struct go_iface go_context = {0};
-    // get_Go_context(ctx, 2, 0, true);
-    // void *httpclient_ptr = get_argument(ctx, 1);
-    // u32 nextid = 0;
-    // bpf_probe_read(&nextid, sizeof(nextid), (void *)(httpclient_ptr + (httpclient_nextid_pos)));
-    // // Get the span context from go context. The mapping is created in the Invoke probe,
-    // // the context here is derived from the Invoke context.
+    bpf_printk("enter the uprobe_http2Client_NewStream \n");
+    struct go_iface go_context = {0};
+    get_Go_context(ctx, 2, 0, true);
+    void *httpclient_ptr = get_argument(ctx, 1);
+    u32 nextid = 0;
+    bpf_probe_read(&nextid, sizeof(nextid), (void *)(httpclient_ptr + (httpclient_nextid_pos)));
+    // Get the span context from go context. The mapping is created in the Invoke probe,
+    // the context here is derived from the Invoke context.
     // struct span_context *current_span_context = get_parent_span_context(&go_context);
-    // if (current_span_context != NULL) {
-    //     bpf_printk("uprobe_http2Client_NewStream streamid is %d\n", nextid);
-    //     bpf_map_update_elem(&streamid_to_span_contexts, &nextid, current_span_context, 0);
-    // }
+    struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
+    bpf_map_update_elem(&streamid_to_span_contexts, &nextid, &trace_key, 0);
 
     return 0;
 }