Refactor span creation

internal/include/go_context.h

@@ -29,7 +29,7 @@ struct
     __type(value, struct span_context);
     __uint(max_entries, MAX_CONCURRENT_SPANS);
     __uint(pinning, LIBBPF_PIN_BY_NAME);
-} tracked_spans SEC(".maps");
+} go_context_to_sc SEC(".maps");
 
 struct
 {
@@ -40,8 +40,8 @@ struct
     __uint(pinning, LIBBPF_PIN_BY_NAME);
 } tracked_spans_by_sc SEC(".maps");
 
-static __always_inline void *get_parent_go_context(void *ctx, void *map) {
-    void *data = ctx;
+static __always_inline void *get_parent_go_context(struct go_iface *go_context, void *map) {
+    void *data = go_context->data;
     for (int i = 0; i < MAX_DISTANCE; i++)
     {
         if (data == NULL)
@@ -60,18 +60,17 @@ static __always_inline void *get_parent_go_context(void *ctx, void *map) {
         bpf_probe_read(&data, sizeof(data), data + 8);
     }
 
-    bpf_printk("context %lx not found in context map", ctx);
     return NULL;
 }
 
-static __always_inline struct span_context *get_parent_span_context(void *ctx) {
-    void *parent_ctx = get_parent_go_context(ctx, &tracked_spans);
-    if (parent_ctx == NULL)
+static __always_inline struct span_context *get_parent_span_context(struct go_iface *go_context) {
+    void *parent_go_ctx = get_parent_go_context(go_context, &go_context_to_sc);
+    if (parent_go_ctx == NULL)
     {
         return NULL;
     }
 
-    struct span_context *parent_sc = bpf_map_lookup_elem(&tracked_spans, &parent_ctx);
+    struct span_context *parent_sc = bpf_map_lookup_elem(&go_context_to_sc, &parent_go_ctx);
     if (parent_sc == NULL)
     {
         return NULL;
@@ -82,7 +81,7 @@ static __always_inline struct span_context *get_parent_span_context(void *ctx) {
 
 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);
+    err = bpf_map_update_elem(&go_context_to_sc, &contextContext, sc, BPF_ANY);
     if (err != 0)
     {
         bpf_printk("Failed to update tracked_spans map: %ld", err);
@@ -114,7 +113,7 @@ static __always_inline void stop_tracking_span(struct span_context *sc, struct s
     if (parent_ctx == NULL)
     {
         // No parent span, delete the context
-        bpf_map_delete_elem(&tracked_spans, ctx);
+        bpf_map_delete_elem(&go_context_to_sc, ctx);
     } else 
     {
         void *ctx_val = 0;
@@ -125,37 +124,38 @@ static __always_inline void stop_tracking_span(struct span_context *sc, struct s
         if (ctx_val != parent_ctx_val)
         {
             // Parent with different context, delete the context
-            bpf_map_delete_elem(&tracked_spans, ctx);
+            bpf_map_delete_elem(&go_context_to_sc, 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_update_elem(&go_context_to_sc, 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);
+//  context_pos:
+//      The argument position of the context.Context type 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, struct go_iface *contextContext) {
+    // Read the argument which is either the context.Context type pointer or pointer to a struct containing the context.Context
+    void *ctx_type_or_struct = get_argument(ctx, context_pos);
     if (passed_as_arg) {
-        return arg;
+        contextContext->type = ctx_type_or_struct;
+        contextContext->data = get_argument(ctx, context_pos + 1);
+    } else {
+        void *context_struct_ptr = (void*)(ctx_type_or_struct + context_offset);
+        bpf_probe_read(&contextContext->type, sizeof(contextContext->type),context_struct_ptr );
+        bpf_probe_read(&contextContext->data, sizeof(contextContext->data), get_go_interface_instance(context_struct_ptr));
     }
-    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

@@ -38,7 +38,7 @@ struct go_slice
 
 struct go_iface
 {
-    void *tab;
+    void *type;
     void *data;
 };
 

internal/include/trace/sampling.h

@@ -0,0 +1,32 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _SAMPLING_H_
+#define _SAMPLING_H_
+
+#include "common.h"
+#include "span_context.h"
+
+typedef struct sampling_parameters {
+    struct span_context *psc;
+    u8 *trace_id;
+    // TODO: add more fields
+} sampling_parameters_t;
+
+static __always_inline bool should_sample(sampling_parameters_t *params) {
+    // TODO
+    return true;
+}
+
+#endif

internal/include/span_context.h → internal/include/trace/span_context.h

@@ -42,24 +42,12 @@ struct span_context
 static __always_inline void get_span_context_from_parent(struct span_context *parent, struct span_context *child) {
     copy_byte_arrays(parent->TraceID, child->TraceID, TRACE_ID_SIZE);
     generate_random_bytes(child->SpanID, SPAN_ID_SIZE);
-    child->TraceFlags = parent->TraceFlags;
 }
 
 // Fill the passed span context as root span context
 static __always_inline void get_root_span_context(struct span_context *sc) {
     generate_random_bytes(sc->TraceID, TRACE_ID_SIZE);
     generate_random_bytes(sc->SpanID, SPAN_ID_SIZE);
-    // currently we always sample
-    sc->TraceFlags = FLAG_SAMPLED;
-}
-
-// TODO: remove this function once all the probes move to the above functions
-static __always_inline struct span_context generate_span_context()
-{
-    struct span_context context = {};
-    generate_random_bytes(context.TraceID, TRACE_ID_SIZE);
-    generate_random_bytes(context.SpanID, SPAN_ID_SIZE);
-    return context;
 }
 
 static __always_inline void span_context_to_w3c_string(struct span_context *ctx, char *buff)

internal/include/span_output.h → internal/include/trace/span_output.h

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "span_context.h"
+#include "trace/span_context.h"
 #include "common.h"
-#include "span_context.h"
+#include "trace/span_context.h"
 
 #ifndef _SPAN_OUTPUT_H_
 #define _SPAN_OUTPUT_H_

internal/include/trace/start_span.h

@@ -0,0 +1,75 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _START_SPAN_H_
+#define _START_SPAN_H_
+
+#include "common.h"
+#include "span_context.h"
+#include "sampling.h"
+
+// function for getting the parent span context, the result is stored in the passed span context.
+// the function should return 0 if the parent span context is found, negative value otherwise.
+// Each probe can potentially have a different way of getting the parent span context,
+// this is useful for incoming requests (http, kafka, etc.) where the parent span context needs to be extracted from the
+// incoming request.
+// The handle param can be used to pass any data needed to get the parent span context.
+typedef long (*get_parent_sc_fn)(void *handle, struct span_context *psc);
+
+typedef struct start_span_params {
+    struct pt_regs *ctx;
+    struct go_iface *go_context;
+    struct span_context *psc;
+    struct span_context *sc;
+    // function for getting the parent span context, the result is stored in the passed span context.
+    get_parent_sc_fn get_parent_span_context_fn;
+    // argument to be passed to the get_parent_span_context_fn
+    void *get_parent_span_context_arg;
+} start_span_params_t;
+
+// Start a new span, setting the parent span context if found.
+// Generate a new span context for the new span. Perform sampling decision and set the TraceFlags accordingly.
+static __always_inline void start_span(start_span_params_t *params) {
+    long found_parent = -1;
+    if (params->get_parent_span_context_fn != NULL) {
+        found_parent = params->get_parent_span_context_fn(params->get_parent_span_context_arg, params->psc);
+    } else {
+        struct span_context *local_psc = get_parent_span_context(params->go_context);
+        if (local_psc != NULL) {
+            found_parent = 0;
+            *(params->psc) = *local_psc;
+        }
+    }
+
+    u8 parent_trace_flags = 0;
+    if (found_parent == 0) {
+        get_span_context_from_parent(params->psc, params->sc);
+        parent_trace_flags = params->psc->TraceFlags;
+    } else {
+        get_root_span_context(params->sc);
+    }
+
+    sampling_parameters_t sampling_params = {
+        .trace_id = params->sc->TraceID,
+        .psc = (found_parent == 0) ? params->psc : NULL,
+    };
+    bool sample = should_sample(&sampling_params);
+    if (sample) {
+        params->sc->TraceFlags = (parent_trace_flags) | (FLAG_SAMPLED);
+    } else {
+        params->sc->TraceFlags = (parent_trace_flags) & (~FLAG_SAMPLED);
+    }
+}
+
+#endif

internal/include/uprobe.h

@@ -16,10 +16,10 @@
 #define _UPROBE_H_
 
 #include "common.h"
-#include "span_context.h"
+#include "trace/span_context.h"
 #include "go_context.h"
 #include "go_types.h"
-#include "span_output.h"
+#include "trace/span_output.h"
 
 #define BASE_SPAN_PROPERTIES \
     u64 start_time;          \
@@ -37,8 +37,9 @@
 #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);                                                               \
+    struct go_iface go_context = {0};                                                                               \
+    get_Go_context(ctx, context_pos, context_offset, passed_as_arg, &go_context);                                   \
+    void *key = get_consistent_key(ctx, go_context.data);                                                           \
     event_type *event = bpf_map_lookup_elem(&uprobe_context_map, &key);                                             \
     if (event == NULL) {                                                                                            \
         bpf_printk("event is NULL in ret probe");                                                                   \

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

@@ -13,10 +13,11 @@
 // limitations under the License.
 
 #include "arguments.h"
-#include "span_context.h"
+#include "trace/span_context.h"
 #include "go_context.h"
 #include "go_types.h"
 #include "uprobe.h"
+#include "trace/start_span.h"
 
 char __license[] SEC("license") = "Dual MIT/GPL";
 
@@ -58,22 +59,23 @@ int uprobe_queryDC(struct pt_regs *ctx) {
         bpf_probe_read(sql_request.query, query_size, query_str_ptr);
     }
 
-    // Get parent if exists
-    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
-        bpf_probe_read(&sql_request.psc, sizeof(sql_request.psc), span_ctx);
-        get_span_context_from_parent(&sql_request.psc, &sql_request.sc);
-    } else {
-       get_root_span_context(&sql_request.sc);
-    }
+    struct go_iface go_context = {0};
+    get_Go_context(ctx, 2, 0, true, &go_context);
+    start_span_params_t start_span_params = {
+        .ctx = ctx,
+        .go_context = &go_context,
+        .psc = &sql_request.psc,
+        .sc = &sql_request.sc,
+        .get_parent_span_context_fn = NULL,
+        .get_parent_span_context_arg = NULL,
+    };
+    start_span(&start_span_params);
 
     // Get key
-    void *key = get_consistent_key(ctx, context_ptr_val);
+    void *key = get_consistent_key(ctx, go_context.data);
 
     bpf_map_update_elem(&sql_events, &key, &sql_request, 0);
-    start_tracking_span(context_ptr_val, &sql_request.sc);
+    start_tracking_span(go_context.data, &sql_request.sc);
     return 0;
 }
 
@@ -101,22 +103,23 @@ int uprobe_execDC(struct pt_regs *ctx) {
         bpf_probe_read(sql_request.query, query_size, query_str_ptr);
     }
 
-    // Get parent if exists
-    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
-        bpf_probe_read(&sql_request.psc, sizeof(sql_request.psc), span_ctx);
-        get_span_context_from_parent(&sql_request.psc, &sql_request.sc);
-    } else {
-        get_root_span_context(&sql_request.sc);
-    }
+    struct go_iface go_context = {0};
+    get_Go_context(ctx, 2, 0, true, &go_context);
+    start_span_params_t start_span_params = {
+        .ctx = ctx,
+        .go_context = &go_context,
+        .psc = &sql_request.psc,
+        .sc = &sql_request.sc,
+        .get_parent_span_context_fn = NULL,
+        .get_parent_span_context_arg = NULL,
+    };
+    start_span(&start_span_params);
 
     // Get key
-    void *key = get_consistent_key(ctx, context_ptr_val);
+    void *key = get_consistent_key(ctx, go_context.data);
 
     bpf_map_update_elem(&sql_events, &key, &sql_request, 0);
-    start_tracking_span(context_ptr_val, &sql_request.sc);
+    start_tracking_span(go_context.data, &sql_request.sc);
     return 0;
 }