github.com/thanos-io/thanos@v0.32.5/pkg/server/http/middleware/request_id.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package middleware
     5  
     6  import (
     7  	"context"
     8  	"math/rand"
     9  	"net/http"
    10  	"time"
    11  
    12  	"github.com/oklog/ulid"
    13  )
    14  
    15  type ctxKey int
    16  
    17  const reqIDKey = ctxKey(0)
    18  
    19  // newContextWithRequestID creates a context with a request id.
    20  func newContextWithRequestID(ctx context.Context, rid string) context.Context {
    21  	return context.WithValue(ctx, reqIDKey, rid)
    22  }
    23  
    24  // RequestIDFromContext returns the request id from context.
    25  func RequestIDFromContext(ctx context.Context) (string, bool) {
    26  	rid, ok := ctx.Value(reqIDKey).(string)
    27  	return rid, ok
    28  }
    29  
    30  // RequestID sets a unique request id for each request.
    31  func RequestID(h http.Handler) http.HandlerFunc {
    32  	return func(w http.ResponseWriter, r *http.Request) {
    33  		reqID := r.Header.Get("X-Request-ID")
    34  		if reqID == "" {
    35  			entropy := ulid.Monotonic(rand.New(rand.NewSource(time.Now().UnixNano())), 0)
    36  			reqID := ulid.MustNew(ulid.Timestamp(time.Now()), entropy)
    37  			r.Header.Set("X-Request-ID", reqID.String())
    38  		}
    39  		ctx := newContextWithRequestID(r.Context(), reqID)
    40  		h.ServeHTTP(w, r.WithContext(ctx))
    41  	}
    42  }