Fix context propagation across goroutines (#118)

* refactor global eBPF maps

* rollback emojivoto yaml

* convert gin instrumentation to new maps

* convert gorillamux and grpc client to new maps

* fix compilation errors

* update generated files

* add changelog and generated files

* fix verifier errors

* revert logsize

* working on fixing context propagation

* update offsets

* fix context propagation

* commit generated files

* change net/http client instrumentation to new context propagation

* add generated files

* implement pr comments

* fix CHANGELOG

CHANGELOG.md

@@ -8,6 +8,10 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix context propagation across different goroutines. ([#118](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/118))
+
 ## [v0.2.2-alpha] - 2023-07-12
 
 ### Added

include/arguments.h

@@ -64,12 +64,20 @@ void *get_argument(struct pt_regs *ctx, int index)
  return get_argument_by_stack(ctx, index);
 }
 
-inline void *get_goroutine_address(struct pt_regs *ctx, int go_ctx_index)
+// Every span created by the auto instrumentation should contain end timestamp.
+// This end timestamp is recorded at the end of probed function by editing the struct that was created at the beginning.
+// Usually instrumentors create an eBPF map to store the span struct and retrieve it at the end of the function.
+// Consistent key is used as a key for that map.
+// For Go < 1.17: consistent key is the address of context.Context.
+// For Go >= 1.17: consistent key is the goroutine address.
+static __always_inline void *get_consistent_key(struct pt_regs *ctx, void *contextContext)
 {
  if (is_registers_abi)
  {
  return (void *)GOROUTINE(ctx);
  }
 
- return get_argument_by_stack(ctx, go_ctx_index);
-}
+ void *ctx_ptr = 0;
+ bpf_probe_read(&ctx_ptr, sizeof(ctx_ptr), contextContext);
+ return ctx_ptr;
+}

include/go_context.h

@@ -16,22 +16,71 @@
 
 #define MAX_DISTANCE 10
 
-static __always_inline void *find_context_in_map(void *ctx, void *context_map)
+struct
 {
+ __uint(type, BPF_MAP_TYPE_HASH);
+ __type(key, void *);
+ __type(value, struct span_context);
+ __uint(max_entries, MAX_CONCURRENT_SPANS);
+ __uint(pinning, LIBBPF_PIN_BY_NAME);
+} tracked_spans SEC(".maps");
+
+struct
+{
+ __uint(type, BPF_MAP_TYPE_HASH);
+ __type(key, struct span_context);
+ __type(value, void *);
+ __uint(max_entries, MAX_CONCURRENT_SPANS);
+ __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;
  for (int i = 0; i < MAX_DISTANCE; i++)
  {
- void *found_in_map = bpf_map_lookup_elem(context_map, &data);
+ void *found_in_map = bpf_map_lookup_elem(map, &data);
  if (found_in_map != NULL)
  {
  return data;
  }
 
- // We assume context.Context implementation containens Parent context.Context member
+ // We assume context.Context implementation contains Parent context.Context member
  // Since the parent is also an interface, we need to read the data part of it
  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)
+ {
+ return NULL;
+ }
+
+ struct span_context *parent_sc = bpf_map_lookup_elem(&tracked_spans, &parent_ctx);
+ if (parent_sc == NULL)
+ {
+ return NULL;
+ }
+
+ 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 stop_tracking_span(struct span_context *sc) {
+ void *ctx = bpf_map_lookup_elem(&tracked_spans_by_sc, sc);
+ if (ctx == NULL)
+ {
+ return;
+ }
+
+ bpf_map_delete_elem(&tracked_spans, &ctx);
+ bpf_map_delete_elem(&tracked_spans_by_sc, sc);
+}

include/span_context.h

@@ -23,15 +23,6 @@ struct span_context
  unsigned char SpanID[SPAN_ID_SIZE];
 };
 
-struct
-{
- __uint(type, BPF_MAP_TYPE_HASH);
- __type(key, void *);
- __type(value, struct span_context);
- __uint(max_entries, MAX_CONCURRENT_SPANS);
- __uint(pinning, LIBBPF_PIN_BY_NAME);
-} spans_in_progress SEC(".maps");
-
 static __always_inline struct span_context generate_span_context()
 {
  struct span_context context = {};

offsets-tracker/main.go

@@ -109,6 +109,18 @@ func main() {
  StructName: "golang.org/x/net/http2.FrameHeader",
  Field: "StreamID",
  },
+ {
+ StructName: "google.golang.org/grpc/internal/transport.http2Client",
+ Field: "nextID",
+ },
+ {
+ StructName: "google.golang.org/grpc/internal/transport.headerFrame",
+ Field: "streamID",
+ },
+ {
+ StructName: "google.golang.org/grpc/internal/transport.headerFrame",
+ Field: "hf",
+ },
  })
 
  if err != nil {

pkg/inject/offset_results.json

@@ -100,6 +100,70 @@
  ]
  }
  },
+ "google.golang.org/grpc/internal/transport.headerFrame": {
+ "hf": {
+ "versions": {
+ "oldest": "1.3.0",
+ "newest": "1.58.0-dev"
+ },
+ "offsets": [
+ {
+ "offset": 8,
+ "since": "v1.3.0"
+ }
+ ]
+ },
+ "streamID": {
+ "versions": {
+ "oldest": "1.3.0",
+ "newest": "1.58.0-dev"
+ },
+ "offsets": [
+ {
+ "offset": 0,
+ "since": "v1.3.0"
+ }
+ ]
+ }
+ },
+ "google.golang.org/grpc/internal/transport.http2Client": {
+ "nextID": {
+ "versions": {
+ "oldest": "1.3.0",
+ "newest": "1.58.0-dev"
+ },
+ "offsets": [
+ {
+ "offset": 412,
+ "since": "v1.3.0"
+ },
+ {
+ "offset": 356,
+ "since": "v1.15.0"
+ },
+ {
+ "offset": 348,
+ "since": "v1.24.0"
+ },
+ {
+ "offset": 340,
+ "since": "v1.35.0"
+ },
+ {
+ "offset": 348,
+ "since": "v1.48.0"
+ },
+ {
+ "offset": 340,
+ "since": "v1.51.0"
+ },
+ {
+ "offset": 412,
+ "since": "v1.52.0"
+ }
+ ]
+ }
+ },
  "net/http.Request": {
  "Header": {
  "versions": {

pkg/instrumentors/bpf/github.com/gin-gonic/gin/bpf/probe.bpf.c

@@ -30,13 +30,12 @@ struct http_request_t {
  struct span_context sc;
 };
 
-// map key: pointer to the goroutine that handles the request
 struct {
  __uint(type, BPF_MAP_TYPE_HASH);
  __type(key, void *);
  __type(value, struct http_request_t);
  __uint(max_entries, MAX_CONCURRENT);
-} context_to_http_events SEC(".maps");
+} http_events SEC(".maps");
 
 struct {
  __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
@@ -79,25 +78,32 @@ int uprobe_GinEngine_ServeHTTP(struct pt_regs *ctx) {
  path_size = path_size < path_len ? path_size : path_len;
  bpf_probe_read(&httpReq.path, path_size, path_ptr);
 
- // Get goroutine pointer
- void *goroutine = get_goroutine_address(ctx, ctx_ptr_pos);
+ // Get key
+ void *req_ctx_ptr = 0;
+ bpf_probe_read(&req_ctx_ptr, sizeof(req_ctx_ptr), (void *)(req_ptr + ctx_ptr_pos));
+ void *key = get_consistent_key(ctx, (void *)(req_ptr + ctx_ptr_pos));
 
  // Write event
  httpReq.sc = generate_span_context();
- bpf_map_update_elem(&context_to_http_events, &goroutine, &httpReq, 0);
- bpf_map_update_elem(&spans_in_progress, &goroutine, &httpReq.sc, 0);
+ bpf_map_update_elem(&http_events, &key, &httpReq, 0);
+ start_tracking_span(req_ctx_ptr, &httpReq.sc);
  return 0;
 }
 
 SEC("uprobe/GinEngine_ServeHTTP")
 int uprobe_GinEngine_ServeHTTP_Returns(struct pt_regs *ctx) {
- void *goroutine = get_goroutine_address(ctx, ctx_ptr_pos);
- void *httpReq_ptr = bpf_map_lookup_elem(&context_to_http_events, &goroutine);
+ u64 request_pos = 4;
+ void *req_ptr = get_argument(ctx, request_pos);
+
+ // Get key
+ void *key = get_consistent_key(ctx, (void *)(req_ptr + ctx_ptr_pos));
+
+ void *httpReq_ptr = bpf_map_lookup_elem(&http_events, &key);
  struct http_request_t httpReq = {};
  bpf_probe_read(&httpReq, sizeof(httpReq), httpReq_ptr);
  httpReq.end_time = bpf_ktime_get_ns();
  bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &httpReq, sizeof(httpReq));
- bpf_map_delete_elem(&context_to_http_events, &goroutine);
- bpf_map_delete_elem(&spans_in_progress, &goroutine);
+ bpf_map_delete_elem(&http_events, &key);
+ stop_tracking_span(&httpReq.sc);
  return 0;
 }

pkg/instrumentors/bpf/github.com/gin-gonic/gin/probe.go

@@ -99,6 +99,11 @@ func (h *Instrumentor) Load(ctx *context.InstrumentorContext) error {
  StructName: "net/url.URL",
  Field: "Path",
  },
+ {
+ VarName: "ctx_ptr_pos",
+ StructName: "net/http.Request",
+ Field: "ctx",
+ },
  }, false)
 
  if err != nil {