fix-tracing: use custom nats header carrier for otel (#2149)

closes:
https://linear.app/overmind/issue/ENG-348/handlequery-does-not-correctly-link-into-main-trace
GitOrigin-RevId: eda2adcbba634b87b177167bdf1aeb1524ed52dc

Author: carabasdaniel

sdp-go/tracing.go

@@ -6,8 +6,8 @@ import (
 	"connectrpc.com/connect"
 	"github.com/getsentry/sentry-go"
 	"github.com/nats-io/nats.go"
+	"github.com/overmindtech/cli/tracing"
 	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -21,7 +21,7 @@ func NewOtelExtractingHandler(spanName string, h CtxMsgHandler, t trace.Tracer,
 	return func(msg *nats.Msg) {
 		ctx := context.Background()
 
-		ctx = otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(msg.Header))
+		ctx = otel.GetTextMapPropagator().Extract(ctx, tracing.NewNatsHeaderCarrier(msg.Header))
 
 		// don't start a span when we have no spanName
 		if spanName != "" {
@@ -44,7 +44,7 @@ func NewAsyncOtelExtractingHandler(spanName string, h CtxMsgHandler, t trace.Tra
 			defer sentry.Recover()
 
 			ctx := context.Background()
-			ctx = otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(msg.Header))
+			ctx = otel.GetTextMapPropagator().Extract(ctx, tracing.NewNatsHeaderCarrier(msg.Header))
 
 			// don't start a span when we have no spanName
 			if spanName != "" {
@@ -62,7 +62,8 @@ func InjectOtelTraceContext(ctx context.Context, msg *nats.Msg) {
 	if msg.Header == nil {
 		msg.Header = make(nats.Header)
 	}
-	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(msg.Header))
+
+	otel.GetTextMapPropagator().Inject(ctx, tracing.NewNatsHeaderCarrier(msg.Header))
 }
 
 type sentryInterceptor struct{}

tracing/header_carrier.go

@@ -0,0 +1,31 @@
+package tracing
+
+import "github.com/nats-io/nats.go"
+
+// HeaderCarrier is a custom wrapper on top of nats.Headers for otel's TextMapCarrier.
+type HeaderCarrier struct {
+	headers nats.Header
+}
+
+// NewNatsHeaderCarrier creates a new HeaderCarrier.
+func NewNatsHeaderCarrier(h nats.Header) *HeaderCarrier {
+	return &HeaderCarrier{
+		headers: h,
+	}
+}
+
+func (c *HeaderCarrier) Get(key string) string {
+	return c.headers.Get(key)
+}
+
+func (c *HeaderCarrier) Set(key, value string) {
+	c.headers.Set(key, value)
+}
+
+func (c *HeaderCarrier) Keys() []string {
+	keys := make([]string, 0, len(c.headers))
+	for key := range c.headers {
+		keys = append(keys, key)
+	}
+	return keys
+}