feat: added response x-trace-id header

instrumentation/github.com/gorilla/mux/otelmux/config.go

@@ -10,14 +10,17 @@ import (
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
+const defaultTraceResponseHeaderKey = "X-Trace-Id"
+
 // config is used to configure the mux middleware.
 type config struct {
-	TracerProvider    oteltrace.TracerProvider
-	Propagators       propagation.TextMapPropagator
-	spanNameFormatter func(string, *http.Request) string
-	PublicEndpoint    bool
-	PublicEndpointFn  func(*http.Request) bool
-	Filters           []Filter
+	TracerProvider         oteltrace.TracerProvider
+	Propagators            propagation.TextMapPropagator
+	spanNameFormatter      func(string, *http.Request) string
+	PublicEndpoint         bool
+	PublicEndpointFn       func(*http.Request) bool
+	Filters                []Filter
+	TraceResponseHeaderKey string
 }
 
 // Option specifies instrumentation configuration options.
@@ -55,6 +58,19 @@ func WithPublicEndpointFn(fn func(*http.Request) bool) Option {
 	})
 }
 
+// WithTraceIDResponseHeader enables adding trace id into response header.
+// It accepts a function that generates the header key name. If this parameter
+// function set to `nil` the default header key which is `X-Trace-Id` will be used.
+func WithTraceIDResponseHeader(headerKeyFunc func() string) Option {
+	return optionFunc(func(cfg *config) {
+		if headerKeyFunc == nil {
+			cfg.TraceResponseHeaderKey = defaultTraceResponseHeaderKey // use default trace header
+		} else {
+			cfg.TraceResponseHeaderKey = headerKeyFunc()
+		}
+	})
+}
+
 // WithPropagators specifies propagators to use for extracting
 // information from the HTTP requests. If none are specified, global
 // ones will be used.

instrumentation/github.com/gorilla/mux/otelmux/mux.go

@@ -48,27 +48,29 @@ func Middleware(service string, opts ...Option) mux.MiddlewareFunc {
 
 	return func(handler http.Handler) http.Handler {
 		return traceware{
-			service:           service,
-			tracer:            tracer,
-			propagators:       cfg.Propagators,
-			handler:           handler,
-			spanNameFormatter: cfg.spanNameFormatter,
-			publicEndpoint:    cfg.PublicEndpoint,
-			publicEndpointFn:  cfg.PublicEndpointFn,
-			filters:           cfg.Filters,
+			service:                service,
+			tracer:                 tracer,
+			propagators:            cfg.Propagators,
+			handler:                handler,
+			spanNameFormatter:      cfg.spanNameFormatter,
+			publicEndpoint:         cfg.PublicEndpoint,
+			publicEndpointFn:       cfg.PublicEndpointFn,
+			filters:                cfg.Filters,
+			traceResponseHeaderKey: cfg.TraceResponseHeaderKey,
 		}
 	}
 }
 
 type traceware struct {
-	service           string
-	tracer            trace.Tracer
-	propagators       propagation.TextMapPropagator
-	handler           http.Handler
-	spanNameFormatter func(string, *http.Request) string
-	publicEndpoint    bool
-	publicEndpointFn  func(*http.Request) bool
-	filters           []Filter
+	service                string
+	tracer                 trace.Tracer
+	propagators            propagation.TextMapPropagator
+	handler                http.Handler
+	spanNameFormatter      func(string, *http.Request) string
+	publicEndpoint         bool
+	publicEndpointFn       func(*http.Request) bool
+	filters                []Filter
+	traceResponseHeaderKey string
 }
 
 type recordingResponseWriter struct {
@@ -164,6 +166,11 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	spanName := tw.spanNameFormatter(routeStr, r)
 	ctx, span := tw.tracer.Start(ctx, spanName, opts...)
 	defer span.End()
+
+	if len(tw.traceResponseHeaderKey) > 0 && span.SpanContext().HasTraceID() {
+		w.Header().Add(tw.traceResponseHeaderKey, span.SpanContext().TraceID().String())
+	}
+
 	r2 := r.WithContext(ctx)
 	rrw := getRRW(w)
 	defer putRRW(rrw)