tracing/correlation/baggage_handler.go

package tracingcorrelation

import (
	"net/http"

	opentracing "github.com/opentracing/opentracing-go"
	"gitlab.com/gitlab-org/labkit/correlation"
)

// BaggageHandler will set opentracing baggage items with the current correlation_id.
func BaggageHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		if span := opentracing.SpanFromContext(ctx); span != nil {
			correlationID := correlation.ExtractFromContext(ctx)
			if correlationID != "" {
				span.SetBaggageItem(correlation.FieldName, correlationID)
			} else {
				// If the span contains the correlation_id, but the context doesn't
				// inject it from the span
				correlationID = span.BaggageItem(correlation.FieldName)
				if correlationID != "" {
					ctx = correlation.ContextWithCorrelation(ctx, correlationID)
					r = r.WithContext(ctx)
				}
			}
		}

		h.ServeHTTP(w, r)
	})
}