Improve eBPF context propagation stability

CHANGELOG.md

@@ -17,6 +17,7 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 
 - Fix bug in the `net/http` server instrumentation which always created a new span context. ([#266]https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/266)
 - Fix runtime panic if OTEL_GO_AUTO_TARGET_EXE is not set. ([#339](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/339))
+- Improve eBPF context propagation stability ([#368]https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/368)
 
 ### Deprecated
 

instrumentation.go

@@ -57,6 +57,13 @@ var errUndefinedTarget = fmt.Errorf("undefined target Go binary, consider settin
 // NewInstrumentation returns a new [Instrumentation] configured with the
 // provided opts.
 func NewInstrumentation(opts ...InstrumentationOption) (*Instrumentation, error) {
+	if log.Logger.IsZero() {
+		err := log.Init()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	c := newInstConfig(opts)
 	if err := c.validate(); err != nil {
 		return nil, err
@@ -84,12 +91,11 @@ func NewInstrumentation(opts ...InstrumentationOption) (*Instrumentation, error)
 		return nil, err
 	}
 
-	if log.Logger.IsZero() {
-		err := log.Init()
-		if err != nil {
-			return nil, err
-		}
+	allocDetails, err := process.Allocate(pid)
+	if err != nil {
+		return nil, err
 	}
+	td.AllocationDetails = allocDetails
 
 	log.Logger.V(0).Info(
 		"target process analysis completed",

internal/include/alloc.h

@@ -102,6 +102,12 @@ static __always_inline void *write_target_data(void *data, s32 size)
 
     void *target = (void *)start;
     size = bound_number(size, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
+    u64 page_offset = (u64)target & 0xFFF;
+    u64 dist_to_next_page = 4096 - page_offset;
+    if (dist_to_next_page < size)
+    {
+        target += dist_to_next_page;
+    }
     long success = bpf_probe_write_user(target, data, size);
     if (success == 0)
     {

internal/include/arguments.h

@@ -57,7 +57,7 @@ void *get_argument_by_stack(struct pt_regs *ctx, int index)
     return ptr;
 }
 
-void *get_argument(struct pt_regs *ctx, int index)
+void __always_inline *get_argument(struct pt_regs *ctx, int index)
 {
     if (is_registers_abi)
     {
@@ -80,9 +80,7 @@ static __always_inline void *get_consistent_key(struct pt_regs *ctx, void *conte
         return (void *)GOROUTINE(ctx);
     }
 
-    void *ctx_ptr = 0;
-    bpf_probe_read(&ctx_ptr, sizeof(ctx_ptr), contextContext);
-    return ctx_ptr;
+    return contextContext;
 }
 
 #endif

internal/include/go_context.h

@@ -16,8 +16,10 @@
 #define _GO_CONTEXT_H_
 
 #include "bpf_helpers.h"
+#include "go_types.h"
 
 #define MAX_DISTANCE 10
+#define MAX_CONCURRENT_SPANS 1000
 
 struct
 {
@@ -41,6 +43,11 @@ static __always_inline void *get_parent_go_context(void *ctx, void *map) {
     void *data = ctx;
     for (int i = 0; i < MAX_DISTANCE; i++)
     {
+        if (data == NULL)
+        {
+            break;
+        }
+
         void *found_in_map = bpf_map_lookup_elem(map, &data);
         if (found_in_map != NULL)
         {
@@ -72,20 +79,82 @@ static __always_inline struct span_context *get_parent_span_context(void *ctx) {
     return parent_sc;
 }
 
-static __always_inline void start_tracking_span(void *ctx, struct span_context *sc) {
-    bpf_map_update_elem(&tracked_spans, &ctx, sc, BPF_ANY);
-    bpf_map_update_elem(&tracked_spans_by_sc, sc, &ctx, BPF_ANY);
+static __always_inline void start_tracking_span(void *contextContext, struct span_context *sc) {
+    long err = 0;
+    err = bpf_map_update_elem(&tracked_spans, &contextContext, sc, BPF_ANY);
+    if (err != 0)
+    {
+        bpf_printk("Failed to update tracked_spans map: %ld", err);
+        return;
+    }
+
+    err = bpf_map_update_elem(&tracked_spans_by_sc, sc, &contextContext, BPF_ANY);
+    if (err != 0)
+    {
+        bpf_printk("Failed to update tracked_spans_by_sc map: %ld", err);
+        return;
+    }
 }
 
-static __always_inline void stop_tracking_span(struct span_context *sc) {
+static __always_inline void stop_tracking_span(struct span_context *sc, struct span_context *psc) {
+    if (sc == NULL) {
+        bpf_printk("stop_tracking_span: sc is null");
+        return;
+    }
+
     void *ctx = bpf_map_lookup_elem(&tracked_spans_by_sc, sc);
     if (ctx == NULL)
     {
+        bpf_printk("stop_tracking_span: cant find span context");
         return;
     }
 
-    bpf_map_delete_elem(&tracked_spans, &ctx);
+    void *parent_ctx = ((psc == NULL) ? NULL : bpf_map_lookup_elem(&tracked_spans_by_sc, psc));
+    if (parent_ctx == NULL)
+    {
+        // No parent span, delete the context
+        bpf_map_delete_elem(&tracked_spans, ctx);
+    } else 
+    {
+        void *ctx_val = 0;
+        bpf_probe_read_user(&ctx_val, sizeof(ctx_val), ctx);
+        void *parent_ctx_val = 0;
+        bpf_probe_read_user(&parent_ctx_val, sizeof(parent_ctx_val), parent_ctx);
+
+        if (ctx_val != parent_ctx_val)
+        {
+            // Parent with different context, delete the context
+            bpf_map_delete_elem(&tracked_spans, ctx);
+        } else {
+            // Parent with the same context, update the entry to point to the parent span
+            bpf_map_update_elem(&tracked_spans, ctx, psc, BPF_ANY);
+        }
+    }
+
     bpf_map_delete_elem(&tracked_spans_by_sc, sc);
 }
 
+// Extract the go context.Context data pointer from the function arguments
+// context_pos:
+    // The argument position of the context.Context data pointer
+    // In case the context.Context is passed as an argument,
+    // this is the argument index of the pointer (starting from 1).
+    // In case the context.Context is a member of a struct,
+    // this is the argument index of the struct pointer (starting from 1). 
+// context_offset:
+    // In case the context.Context is a member of a struct,
+    // this is the offset of the context.Context member in the struct.
+// passed_as_arg:
+    // Indicates whether context.Context is passed as an argument or is a member of a struct
+static __always_inline void *get_Go_context(void *ctx, int context_pos, const volatile u64 context_offset, bool passed_as_arg) {
+    void *arg = get_argument(ctx, context_pos);
+    if (passed_as_arg) {
+        return arg;
+    }
+    void *ctx_addr = get_go_interface_instance((void*)(arg + context_offset));
+    void *ctx_val = 0;
+    bpf_probe_read_user(&ctx_val, sizeof(ctx_val), ctx_addr);
+    return ctx_val;
+}
+
 #endif

internal/include/go_types.h

@@ -54,13 +54,25 @@ struct map_bucket {
     void *overflow;
 };
 
+// In Go, interfaces are represented as a pair of pointers: a pointer to the
+// interface data, and a pointer to the interface table.
+// See: runtime.iface in https://golang.org/src/runtime/runtime2.go
+static __always_inline void* get_go_interface_instance(void *iface)
+{
+    return (void*)(iface + 8);
+}
+
 static __always_inline struct go_string write_user_go_string(char *str, u32 len)
 {
     // Copy chars to userspace
+    struct go_string new_string = {.str = NULL, .len = 0};
     char *addr = write_target_data((void *)str, len);
+    if (addr == NULL) {
+        bpf_printk("write_user_go_string: failed to copy string to userspace");
+        return new_string;
+    }
 
     // Build string struct in kernel space
-    struct go_string new_string = {};
     new_string.str = addr;
     new_string.len = len;
 
@@ -82,7 +94,7 @@ static __always_inline void append_item_to_slice(struct go_slice *slice, void *n
     }
     else
     {
-        // No room on current array - copy to new one of size item_size * (len + 1)
+        //No room on current array - copy to new one of size item_size * (len + 1)
         if (slice->len > MAX_DATA_SIZE || slice->len < 1)
         {
             return;
@@ -111,19 +123,63 @@ static __always_inline void append_item_to_slice(struct go_slice *slice, void *n
         }
 
         void *new_array = write_target_data(map_buff, new_array_size);
+        if (new_array == NULL)
+        {
+            bpf_printk("append_item_to_slice: failed to copy new array to userspace");
+            return;
+        }
 
         // Update array
         slice->array = new_array;
         long success = bpf_probe_write_user(slice_user_ptr->array, &slice->array, sizeof(slice->array));
+        if (success != 0)
+        {
+            bpf_printk("append_item_to_slice: failed to update array pointer in userspace");
+            return;
+        }
 
         // Update cap
         slice->cap++;
         success = bpf_probe_write_user(slice_user_ptr->cap, &slice->cap, sizeof(slice->cap));
+        if (success != 0)
+        {
+            bpf_printk("append_item_to_slice: failed to update cap in userspace");
+            return;
+        }
     }
 
     // Update len
     slice->len++;
     long success = bpf_probe_write_user(slice_user_ptr->len, &slice->len, sizeof(slice->len));
+    if (success != 0)
+    {
+        bpf_printk("append_item_to_slice: failed to update len in userspace");
+        return;
+    }
 }
 
+static __always_inline bool get_go_string_from_user_ptr(void *user_str_ptr, char *dst, u64 max_len)
+{
+    if (user_str_ptr == NULL)
+    {
+        return false;
+    }
+
+    struct go_string user_str = {0};
+    long success = 0;
+    success = bpf_probe_read(&user_str, sizeof(struct go_string), user_str_ptr);
+    if (success != 0 || user_str.len < 1)
+    {
+        return false;
+    }
+
+    u64 size_to_read = user_str.len > max_len ? max_len : user_str.len;
+    success = bpf_probe_read(dst, size_to_read, user_str.str);
+    if (success != 0)
+    {
+        return false;
+    }
+
+    return true;
+}
 #endif

internal/include/span_context.h

@@ -18,7 +18,6 @@
 #include "utils.h"
 
 #define SPAN_CONTEXT_STRING_SIZE 55
-#define MAX_CONCURRENT_SPANS 100
 
 struct span_context
 {

internal/include/uprobe.h

@@ -18,6 +18,7 @@
 #include "common.h"
 #include "span_context.h"
 #include "go_context.h"
+#include "go_types.h"
 
 #define BASE_SPAN_PROPERTIES \
     u64 start_time;          \
@@ -32,19 +33,22 @@
 // 4. Submit the constructed event to the agent code using perf buffer events_map
 // 5. Delete the span from the uprobe_context_map
 // 6. Delete the span from the global active spans map
-#define UPROBE_RETURN(name, event_type, ctx_struct_pos, ctx_struct_offset, uprobe_context_map, events_map) \
-SEC("uprobe/##name##")                                                                                     \
-int uprobe_##name##_Returns(struct pt_regs *ctx) {                                                         \
-    void *req_ptr = get_argument(ctx, ctx_struct_pos);                                                     \
-    void *key = get_consistent_key(ctx, (void *)(req_ptr + ctx_struct_offset));                            \
-    void *req_ptr_map = bpf_map_lookup_elem(&uprobe_context_map, &key);                                    \
-    event_type tmpReq = {};                                                                                \
-    bpf_probe_read(&tmpReq, sizeof(tmpReq), req_ptr_map);                                                  \
-    tmpReq.end_time = bpf_ktime_get_ns();                                                                  \
-    bpf_perf_event_output(ctx, &events_map, BPF_F_CURRENT_CPU, &tmpReq, sizeof(tmpReq));                   \
-    bpf_map_delete_elem(&uprobe_context_map, &key);                                                        \
-    stop_tracking_span(&tmpReq.sc);                                                                        \
-    return 0;                                                                                              \
-}
+#define UPROBE_RETURN(name, event_type, uprobe_context_map, events_map, context_pos, context_offset, passed_as_arg) \
+SEC("uprobe/##name##")                                                                                              \
+int uprobe_##name##_Returns(struct pt_regs *ctx) {                                                                  \
+    void *ctx_address = get_Go_context(ctx, context_pos, context_offset, passed_as_arg);                            \
+    void *key = get_consistent_key(ctx, ctx_address);                                                               \
+    void *req_ptr_map = bpf_map_lookup_elem(&uprobe_context_map, &key);                                             \
+    if (req_ptr_map == NULL) {                                                                                      \
+        return 0;                                                                                                   \
+    }                                                                                                               \
+    event_type tmpReq = {0};                                                                                        \
+    bpf_probe_read(&tmpReq, sizeof(tmpReq), req_ptr_map);                                                           \
+    tmpReq.end_time = bpf_ktime_get_ns();                                                                           \
+    bpf_perf_event_output(ctx, &events_map, BPF_F_CURRENT_CPU, &tmpReq, sizeof(tmpReq));                            \
+    bpf_map_delete_elem(&uprobe_context_map, &key);                                                                 \
+    stop_tracking_span(&tmpReq.sc, &tmpReq.psc);                                                                    \
+    return 0;                                                                                                       \
+} 
 
 #endif

internal/include/utils.h

@@ -83,3 +83,16 @@ static __always_inline void bpf_memset(unsigned char *dst, u32 size, unsigned ch
         dst[i] = value;
     }
 }
+
+static __always_inline bool bpf_is_zero(unsigned char *buff, u32 size)
+{
+    for (int i = 0; i < size; i++)
+    {
+        if (buff[i] != 0)
+        {
+            return false;
+        }
+    }
+
+    return true;
+}

internal/pkg/instrumentors/bpf/database/sql/bpf/probe.bpf.c

@@ -63,9 +63,7 @@ int uprobe_queryDC(struct pt_regs *ctx) {
     }
 
     // Get parent if exists
-    void *context_ptr = get_argument(ctx, context_ptr_pos);
-    void *context_ptr_val = 0;
-    bpf_probe_read(&context_ptr_val, sizeof(context_ptr_val), context_ptr);
+    void *context_ptr_val = get_Go_context(ctx, 3, 0, true);
     struct span_context *span_ctx = get_parent_span_context(context_ptr_val);
     if (span_ctx != NULL) {
         // Set the parent context
@@ -77,7 +75,7 @@ int uprobe_queryDC(struct pt_regs *ctx) {
     }
 
     // Get key
-    void *key = get_consistent_key(ctx, context_ptr);
+    void *key = get_consistent_key(ctx, context_ptr_val);
 
     bpf_map_update_elem(&sql_events, &key, &sql_request, 0);
     start_tracking_span(context_ptr_val, &sql_request.sc);
@@ -86,4 +84,4 @@ int uprobe_queryDC(struct pt_regs *ctx) {
 
 // This instrumentation attaches uprobe to the following function:
 // func (db *DB) queryDC(ctx, txctx context.Context, dc *driverConn, releaseConn func(error), query string, args []any)
-UPROBE_RETURN(queryDC, struct sql_request_t, 3, 0, sql_events, events)
+UPROBE_RETURN(queryDC, struct sql_request_t, sql_events, events, 3, 0, true)