From 9971ef17cb2643a548ec2b302232a25e4e85026c Mon Sep 17 00:00:00 2001
From: Natalia Goc <goc_natalia@ma-al.com>
Date: Wed, 24 Jul 2024 15:19:51 +0200
Subject: [PATCH] feat: add trace context propagation to http headers

Trace context should now be propagated correctly in headers of requests
and responses touched by the fiber_tracing middleware. This should enable
true distributed tracing between multiple services.
---
 pkg/fiber_tracing/middleware.go | 48 ++++++++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/pkg/fiber_tracing/middleware.go b/pkg/fiber_tracing/middleware.go
index 6882bcf..20d096c 100644
--- a/pkg/fiber_tracing/middleware.go
+++ b/pkg/fiber_tracing/middleware.go
@@ -8,6 +8,7 @@ import (
 	"git.ma-al.com/maal-libraries/observer/pkg/attr"
 	"github.com/gofiber/fiber/v2"
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.25.0"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -45,6 +46,38 @@ func configDefaults(config ...middlewareConfig) middlewareConfig {
 	return cfg
 }
 
+// A helper implementing `propagation.TextMapCarrier`
+type headersCarrier struct {
+	*fiber.Ctx
+}
+
+func (h headersCarrier) Get(key string) string {
+	c := h.Ctx
+	headers := c.GetReqHeaders()
+	if val, ok := headers[key]; ok {
+		if len(val) > 0 {
+			return val[len(val)-1]
+		}
+	}
+	return ""
+}
+
+func (h headersCarrier) Set(key string, value string) {
+	c := h.Ctx
+	c.Set(key, value)
+	return
+}
+
+func (h headersCarrier) Keys() []string {
+	c := h.Ctx
+	headers := c.GetRespHeaders()
+	keys := make([]string, 0)
+	for k, _ := range headers {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
 func new(config ...middlewareConfig) fiber.Handler {
 	// Set default config
 	cfg := configDefaults(config...)
@@ -54,7 +87,6 @@ func new(config ...middlewareConfig) fiber.Handler {
 		spanStartAttributes := []attr.KeyValue{
 			semconv.HTTPMethod(c.Method()),
 			semconv.HTTPTarget(string(c.Request().RequestURI())),
-			semconv.HTTPRoute(c.Route().Path),
 			semconv.HTTPURL(c.OriginalURL()),
 			semconv.HTTPUserAgent(string(c.Request().Header.UserAgent())),
 			semconv.HTTPRequestContentLength(c.Request().Header.ContentLength()),
@@ -69,22 +101,32 @@ func new(config ...middlewareConfig) fiber.Handler {
 		}
 		opts = append(opts, cfg.TracerStartAttributes...)
 
+		// Init context using values from request headers if possible:
+		headersCarrier := headersCarrier{c}
+		ctx := c.UserContext()
+		tctx := propagation.TraceContext{}
+		ctx = tctx.Extract(ctx, headersCarrier)
+
 		otelCtx, span := Tracer.Start(
-			c.UserContext(),
+			ctx,
 			c.Method()+" "+c.OriginalURL(),
 			opts...,
 		)
 
 		c.SetUserContext(otelCtx)
+
 		defer span.End()
 
 		err := c.Next()
 
 		statusCode := c.Response().StatusCode()
 		attrs := semconv.HTTPResponseStatusCode(statusCode)
-
 		span.SetAttributes(attrs)
 
+		// Return with trace context added to headers
+		ctx = c.UserContext()
+		tctx.Inject(ctx, headersCarrier)
+
 		return err
 	}
 }