add request ID injection to context to enable tracking reqeusts across downstream services

Signed-off-by: Erlan Zholdubai uulu <[email protected]>

Author: erlan-z

CHANGELOG.md

@@ -56,6 +56,7 @@
 * [ENHANCEMENT] Distributor: Add native histograms max sample size bytes limit validation. #6834
 * [ENHANCEMENT] Querier: Support caching parquet labels file in parquet queryable. #6835
 * [ENHANCEMENT] Querier: Support query limits in parquet queryable. #6870
+* [ENHANCEMENT] API: add request ID injection to context to enable tracking reqeusts across downstream services. #6895
 * [BUGFIX] Ingester: Avoid error or early throttling when READONLY ingesters are present in the ring #6517
 * [BUGFIX] Ingester: Fix labelset data race condition. #6573
 * [BUGFIX] Compactor: Cleaner should not put deletion marker for blocks with no-compact marker. #6576

docs/configuration/config-file-reference.md

@@ -102,6 +102,10 @@ api:
   # CLI flag: -api.http-request-headers-to-log
   [http_request_headers_to_log: <list of string> | default = []]
 
+  # HTTP header that can be used as request id
+  # CLI flag: -api.request-id-header
+  [request_id_header: <string> | default = ""]
+
   # Regex for CORS origin. It is fully anchored. Example:
   # 'https?://(domain1|domain2)\.com'
   # CLI flag: -server.cors-origin

go.mod

@@ -79,6 +79,7 @@ require (
 	github.com/bboreham/go-loser v0.0.0-20230920113527-fcc2c21820a3
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/google/go-cmp v0.7.0
+	github.com/google/uuid v1.6.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822
 	github.com/oklog/ulid/v2 v2.1.1
@@ -170,7 +171,6 @@ require (
 	github.com/google/btree v1.1.3 // indirect
 	github.com/google/pprof v0.0.0-20250607225305-033d6d78b36a // indirect
 	github.com/google/s2a-go v0.1.9 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect
 	github.com/googleapis/gax-go/v2 v2.14.1 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2 // indirect

pkg/api/api.go

@@ -71,6 +71,10 @@ type Config struct {
 	// Allows and is used to configure the addition of HTTP Header fields to logs
 	HTTPRequestHeadersToLog flagext.StringSlice `yaml:"http_request_headers_to_log"`
 
+	// HTTP header that can be used as request id.
+	// If it's not provided, or this header is empty, then random requestId will be generated
+	RequestIdHeader string `yaml:"request_id_header"`
+
 	// This sets the Origin header value
 	corsRegexString string `yaml:"cors_origin"`
 
@@ -87,6 +91,7 @@ var (
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.ResponseCompression, "api.response-compression-enabled", false, "Use GZIP compression for API responses. Some endpoints serve large YAML or JSON blobs which can benefit from compression.")
 	f.Var(&cfg.HTTPRequestHeadersToLog, "api.http-request-headers-to-log", "Which HTTP Request headers to add to logs")
+	f.StringVar(&cfg.RequestIdHeader, "api.request-id-header", "", "HTTP header that can be used as request id")
 	f.BoolVar(&cfg.buildInfoEnabled, "api.build-info-enabled", false, "If enabled, build Info API will be served by query frontend or querier.")
 	f.StringVar(&cfg.QuerierDefaultCodec, "api.querier-default-codec", "json", "Choose default codec for querier response serialization. Supports 'json' and 'protobuf'.")
 	cfg.RegisterFlagsWithPrefix("", f)
@@ -169,8 +174,9 @@ func New(cfg Config, serverCfg server.Config, s *server.Server, logger log.Logge
 	if cfg.HTTPAuthMiddleware == nil {
 		api.AuthMiddleware = middleware.AuthenticateUser
 	}
-	if len(cfg.HTTPRequestHeadersToLog) > 0 {
-		api.HTTPHeaderMiddleware = &HTTPHeaderMiddleware{TargetHeaders: cfg.HTTPRequestHeadersToLog}
+	api.HTTPHeaderMiddleware = &HTTPHeaderMiddleware{
+		TargetHeaders:   cfg.HTTPRequestHeadersToLog,
+		RequestIdHeader: cfg.RequestIdHeader,
 	}
 
 	return api, nil

pkg/api/api_test.go

@@ -89,6 +89,7 @@ func TestNewApiWithHeaderLogging(t *testing.T) {
 
 }
 
+// HTTPHeaderMiddleware should be added even if no headers are specified to log because it also handles request ID injection.
 func TestNewApiWithoutHeaderLogging(t *testing.T) {
 	cfg := Config{
 		HTTPRequestHeadersToLog: []string{},
@@ -102,7 +103,8 @@ func TestNewApiWithoutHeaderLogging(t *testing.T) {
 
 	api, err := New(cfg, serverCfg, server, &FakeLogger{})
 	require.NoError(t, err)
-	require.Nil(t, api.HTTPHeaderMiddleware)
+	require.NotNil(t, api.HTTPHeaderMiddleware)
+	require.Empty(t, api.HTTPHeaderMiddleware.TargetHeaders)
 
 }
 

pkg/api/middlewares.go

@@ -1,25 +1,28 @@
 package api
 
 import (
-	"context"
 	"net/http"
 
+	"github.com/google/uuid"
+
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	requestutil "github.com/cortexproject/cortex/pkg/util/request"
 )
 
 // HTTPHeaderMiddleware adds specified HTTPHeaders to the request context
 type HTTPHeaderMiddleware struct {
-	TargetHeaders []string
+	TargetHeaders   []string
+	RequestIdHeader string
 }
 
 // InjectTargetHeadersIntoHTTPRequest injects specified HTTPHeaders into the request context
-func (h HTTPHeaderMiddleware) InjectTargetHeadersIntoHTTPRequest(r *http.Request) context.Context {
+func (h HTTPHeaderMiddleware) InjectTargetHeadersIntoHTTPRequest(r *http.Request) *http.Request {
 	headerMap := make(map[string]string)
 
 	// Check to make sure that Headers have not already been injected
 	checkMapInContext := util_log.HeaderMapFromContext(r.Context())
 	if checkMapInContext != nil {
-		return r.Context()
+		return r
 	}
 
 	for _, target := range h.TargetHeaders {
@@ -28,13 +31,30 @@ func (h HTTPHeaderMiddleware) InjectTargetHeadersIntoHTTPRequest(r *http.Request
 			headerMap[target] = contents
 		}
 	}
-	return util_log.ContextWithHeaderMap(r.Context(), headerMap)
+	ctx := util_log.ContextWithHeaderMap(r.Context(), headerMap)
+	return r.WithContext(ctx)
+}
+
+// InjectRequestIdHeader injects a request ID into the request context if one is not already present
+func (h HTTPHeaderMiddleware) InjectRequestIdHeader(r *http.Request) *http.Request {
+
+	if requestutil.RequestIdFromContext(r.Context()) != "" {
+		return r
+	}
+
+	reqId := r.Header.Get(h.RequestIdHeader)
+	if reqId == "" {
+		reqId = uuid.NewString()
+	}
+	ctx := requestutil.ContextWithRequestIdHeader(r.Context(), reqId)
+	return r.WithContext(ctx)
 }
 
 // Wrap implements Middleware
 func (h HTTPHeaderMiddleware) Wrap(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := h.InjectTargetHeadersIntoHTTPRequest(r)
-		next.ServeHTTP(w, r.WithContext(ctx))
+		r = h.InjectTargetHeadersIntoHTTPRequest(r)
+		r = h.InjectRequestIdHeader(r)
+		next.ServeHTTP(w, r)
 	})
 }

pkg/api/middlewares_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	requestutil "github.com/cortexproject/cortex/pkg/util/request"
 )
 
 var HTTPTestMiddleware = HTTPHeaderMiddleware{TargetHeaders: []string{"TestHeader1", "TestHeader2", "Test3"}}
@@ -32,9 +33,9 @@ func TestHeaderInjection(t *testing.T) {
 	}
 
 	req = req.WithContext(ctx)
-	ctx = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
+	req = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
 
-	headerMap := util_log.HeaderMapFromContext(ctx)
+	headerMap := util_log.HeaderMapFromContext(req.Context())
 	require.NotNil(t, headerMap)
 
 	for _, header := range HTTPTestMiddleware.TargetHeaders {
@@ -67,8 +68,74 @@ func TestExistingHeaderInContextIsNotOverridden(t *testing.T) {
 	}
 
 	req = req.WithContext(ctx)
-	ctx = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
+	req = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
 
-	require.Equal(t, contentsMap, util_log.HeaderMapFromContext(ctx))
+	require.Equal(t, contentsMap, util_log.HeaderMapFromContext(req.Context()))
 
 }
+
+func TestRequestIdInjection(t *testing.T) {
+	middleware := HTTPHeaderMiddleware{
+		RequestIdHeader: "X-Request-ID",
+	}
+
+	req := &http.Request{
+		Method:     "GET",
+		RequestURI: "/test",
+		Body:       http.NoBody,
+		Header:     http.Header{},
+	}
+	req = req.WithContext(context.Background())
+	req = middleware.InjectRequestIdHeader(req)
+
+	requestID := requestutil.RequestIdFromContext(req.Context())
+	require.NotEmpty(t, requestID, "Request ID should be generated if not provided")
+}
+
+func TestRequestIdFromHeaderIsUsed(t *testing.T) {
+	const providedID = "my-test-id-123"
+
+	middleware := HTTPHeaderMiddleware{
+		RequestIdHeader: "X-Request-ID",
+	}
+
+	h := http.Header{}
+	h.Add("X-Request-ID", providedID)
+
+	req := &http.Request{
+		Method:     "GET",
+		RequestURI: "/test",
+		Body:       http.NoBody,
+		Header:     h,
+	}
+	req = req.WithContext(context.Background())
+	req = middleware.InjectRequestIdHeader(req)
+
+	requestID := requestutil.RequestIdFromContext(req.Context())
+	require.Equal(t, providedID, requestID, "Request ID from header should be used")
+}
+
+func TestExistingRequestIdIsPreserved(t *testing.T) {
+	const existingID = "already-present-id"
+
+	middleware := HTTPHeaderMiddleware{
+		RequestIdHeader: "X-Request-ID",
+	}
+
+	ctx := requestutil.ContextWithRequestIdHeader(context.Background(), existingID)
+
+	h := http.Header{}
+	h.Add("X-Request-ID", "should-be-ignored")
+
+	req := &http.Request{
+		Method:     "GET",
+		RequestURI: "/test",
+		Body:       http.NoBody,
+		Header:     h,
+	}
+	req = req.WithContext(ctx)
+	req = middleware.InjectRequestIdHeader(req)
+
+	requestID := requestutil.RequestIdFromContext(req.Context())
+	require.Equal(t, existingID, requestID, "Existing request ID in context should be preserved")
+}

pkg/cortex/cortex.go

@@ -392,10 +392,8 @@ func (t *Cortex) setupThanosTracing() {
 // setupGRPCHeaderForwarding appends a gRPC middleware used to enable the propagation of
 // HTTP Headers through child gRPC calls
 func (t *Cortex) setupGRPCHeaderForwarding() {
-	if len(t.Cfg.API.HTTPRequestHeadersToLog) > 0 {
-		t.Cfg.Server.GRPCMiddleware = append(t.Cfg.Server.GRPCMiddleware, grpcutil.HTTPHeaderPropagationServerInterceptor)
-		t.Cfg.Server.GRPCStreamMiddleware = append(t.Cfg.Server.GRPCStreamMiddleware, grpcutil.HTTPHeaderPropagationStreamServerInterceptor)
-	}
+	t.Cfg.Server.GRPCMiddleware = append(t.Cfg.Server.GRPCMiddleware, grpcutil.HTTPHeaderPropagationServerInterceptor)
+	t.Cfg.Server.GRPCStreamMiddleware = append(t.Cfg.Server.GRPCStreamMiddleware, grpcutil.HTTPHeaderPropagationStreamServerInterceptor)
 }
 
 func (t *Cortex) setupRequestSigning() {

pkg/cortex/modules.go

@@ -402,9 +402,7 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 		// request context.
 		internalQuerierRouter = t.API.AuthMiddleware.Wrap(internalQuerierRouter)
 
-		if len(t.Cfg.API.HTTPRequestHeadersToLog) > 0 {
-			internalQuerierRouter = t.API.HTTPHeaderMiddleware.Wrap(internalQuerierRouter)
-		}
+		internalQuerierRouter = t.API.HTTPHeaderMiddleware.Wrap(internalQuerierRouter)
 	}
 
 	// If neither frontend address or scheduler address is configured, no worker is needed.

pkg/querier/tripperware/roundtrip.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	requestutil "github.com/cortexproject/cortex/pkg/util/request"
 )
 
 const (
@@ -269,6 +270,10 @@ func (q roundTripper) Do(ctx context.Context, r Request) (Response, error) {
 		return nil, httpgrpc.Errorf(http.StatusBadRequest, "%s", err.Error())
 	}
 
+	if requestId := requestutil.RequestIdFromContext(ctx); requestId != "" {
+		requestutil.InjectRequestIdIntoHTTPRequest(requestId, request)
+	}
+
 	response, err := q.next.RoundTrip(request)
 	if err != nil {
 		return nil, err