gitlab.com/gitlab-org/labkit@v1.21.0/tracing/correlation/baggage_handler.go (about)

     1  package tracingcorrelation
     2  
     3  import (
     4  	"net/http"
     5  
     6  	opentracing "github.com/opentracing/opentracing-go"
     7  	"gitlab.com/gitlab-org/labkit/correlation"
     8  )
     9  
    10  // BaggageHandler will set opentracing baggage items with the current correlation_id.
    11  func BaggageHandler(h http.Handler) http.Handler {
    12  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    13  		ctx := r.Context()
    14  
    15  		if span := opentracing.SpanFromContext(ctx); span != nil {
    16  			correlationID := correlation.ExtractFromContext(ctx)
    17  			if correlationID != "" {
    18  				span.SetBaggageItem(correlation.FieldName, correlationID)
    19  			} else {
    20  				// If the span contains the correlation_id, but the context doesn't
    21  				// inject it from the span
    22  				correlationID = span.BaggageItem(correlation.FieldName)
    23  				if correlationID != "" {
    24  					ctx = correlation.ContextWithCorrelation(ctx, correlationID)
    25  					r = r.WithContext(ctx)
    26  				}
    27  			}
    28  		}
    29  
    30  		h.ServeHTTP(w, r)
    31  	})
    32  }