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

@@ -59,6 +59,7 @@
 * [ENHANCEMENT] Querier: Support query limits in parquet queryable. #6870
 * [ENHANCEMENT] Ring: Add zone label to ring_members metric. #6900
 * [ENHANCEMENT] Ingester: Add new metric `cortex_ingester_push_errors_total` to track reasons for ingester request failures. #6901
+* [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. It will always be included in logs
+	// 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,40 +1,51 @@
 package api
 
 import (
-	"context"
 	"net/http"
 
-	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/google/uuid"
+
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 )
 
 // 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 {
-	headerMap := make(map[string]string)
+// injectRequestContext injects request related metadata into the request context
+func (h HTTPHeaderMiddleware) injectRequestContext(r *http.Request) *http.Request {
+	requestContextMap := make(map[string]string)
 
-	// Check to make sure that Headers have not already been injected
-	checkMapInContext := util_log.HeaderMapFromContext(r.Context())
+	// Check to make sure that request context have not already been injected
+	checkMapInContext := requestmeta.MapFromContext(r.Context())
 	if checkMapInContext != nil {
-		return r.Context()
+		return r
 	}
 
 	for _, target := range h.TargetHeaders {
 		contents := r.Header.Get(target)
 		if contents != "" {
-			headerMap[target] = contents
+			requestContextMap[target] = contents
 		}
 	}
-	return util_log.ContextWithHeaderMap(r.Context(), headerMap)
+	requestContextMap[requestmeta.LoggingHeadersKey] = requestmeta.LoggingHeaderKeysToString(h.TargetHeaders)
+
+	reqId := r.Header.Get(h.RequestIdHeader)
+	if reqId == "" {
+		reqId = uuid.NewString()
+	}
+	requestContextMap[requestmeta.RequestIdKey] = reqId
+
+	ctx := requestmeta.ContextWithRequestMetadataMap(r.Context(), requestContextMap)
+	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.injectRequestContext(r)
+		next.ServeHTTP(w, r)
 	})
 }

pkg/api/middlewares_test.go

@@ -7,12 +7,11 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 )
 
-var HTTPTestMiddleware = HTTPHeaderMiddleware{TargetHeaders: []string{"TestHeader1", "TestHeader2", "Test3"}}
-
 func TestHeaderInjection(t *testing.T) {
+	middleware := HTTPHeaderMiddleware{TargetHeaders: []string{"TestHeader1", "TestHeader2", "Test3"}}
 	ctx := context.Background()
 	h := http.Header{}
 	contentsMap := make(map[string]string)
@@ -32,12 +31,12 @@ func TestHeaderInjection(t *testing.T) {
 	}
 
 	req = req.WithContext(ctx)
-	ctx = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
+	req = middleware.injectRequestContext(req)
 
-	headerMap := util_log.HeaderMapFromContext(ctx)
+	headerMap := requestmeta.MapFromContext(req.Context())
 	require.NotNil(t, headerMap)
 
-	for _, header := range HTTPTestMiddleware.TargetHeaders {
+	for _, header := range middleware.TargetHeaders {
 		require.Equal(t, contentsMap[header], headerMap[header])
 	}
 	for header, contents := range contentsMap {
@@ -46,6 +45,7 @@ func TestHeaderInjection(t *testing.T) {
 }
 
 func TestExistingHeaderInContextIsNotOverridden(t *testing.T) {
+	middleware := HTTPHeaderMiddleware{TargetHeaders: []string{"TestHeader1", "TestHeader2", "Test3"}}
 	ctx := context.Background()
 
 	h := http.Header{}
@@ -58,7 +58,7 @@ func TestExistingHeaderInContextIsNotOverridden(t *testing.T) {
 	h.Add("TestHeader2", "Fail2")
 	h.Add("Test3", "Fail3")
 
-	ctx = util_log.ContextWithHeaderMap(ctx, contentsMap)
+	ctx = requestmeta.ContextWithRequestMetadataMap(ctx, contentsMap)
 	req := &http.Request{
 		Method:     "GET",
 		RequestURI: "/HTTPHeaderTest",
@@ -67,8 +67,77 @@ func TestExistingHeaderInContextIsNotOverridden(t *testing.T) {
 	}
 
 	req = req.WithContext(ctx)
-	ctx = HTTPTestMiddleware.InjectTargetHeadersIntoHTTPRequest(req)
+	req = middleware.injectRequestContext(req)
+
+	require.Equal(t, contentsMap, requestmeta.MapFromContext(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.injectRequestContext(req)
+
+	requestID := requestmeta.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)
 
-	require.Equal(t, contentsMap, util_log.HeaderMapFromContext(ctx))
+	req := &http.Request{
+		Method:     "GET",
+		RequestURI: "/test",
+		Body:       http.NoBody,
+		Header:     h,
+	}
+	req = req.WithContext(context.Background())
+	req = middleware.injectRequestContext(req)
+
+	requestID := requestmeta.RequestIdFromContext(req.Context())
+	require.Equal(t, providedID, requestID, "Request ID from header should be used")
+}
+
+func TestTargetHeaderAndRequestIdHeaderOverlap(t *testing.T) {
+	const headerKey = "X-Request-ID"
+	const providedID = "overlap-id-456"
+
+	middleware := HTTPHeaderMiddleware{
+		TargetHeaders:   []string{headerKey, "Other-Header"},
+		RequestIdHeader: headerKey,
+	}
+
+	h := http.Header{}
+	h.Add(headerKey, providedID)
+	h.Add("Other-Header", "some-value")
+
+	req := &http.Request{
+		Method:     "GET",
+		RequestURI: "/test",
+		Body:       http.NoBody,
+		Header:     h,
+	}
+	req = req.WithContext(context.Background())
+	req = middleware.injectRequestContext(req)
 
+	ctxMap := requestmeta.MapFromContext(req.Context())
+	requestID := requestmeta.RequestIdFromContext(req.Context())
+	require.Equal(t, providedID, ctxMap[headerKey], "Header value should be correctly stored")
+	require.Equal(t, providedID, requestID, "Request ID should come from the overlapping header")
 }

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/distributor/distributor.go

@@ -40,6 +40,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	util_math "github.com/cortexproject/cortex/pkg/util/math"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
@@ -892,9 +893,9 @@ func (d *Distributor) doBatch(ctx context.Context, req *cortexpb.WriteRequest, s
 	if sp := opentracing.SpanFromContext(ctx); sp != nil {
 		localCtx = opentracing.ContextWithSpan(localCtx, sp)
 	}
-	// Get any HTTP headers that are supposed to be added to logs and add to localCtx for later use
-	if headerMap := util_log.HeaderMapFromContext(ctx); headerMap != nil {
-		localCtx = util_log.ContextWithHeaderMap(localCtx, headerMap)
+	// Get any HTTP request metadata that are supposed to be added to logs and add to localCtx for later use
+	if requestContextMap := requestmeta.MapFromContext(ctx); requestContextMap != nil {
+		localCtx = requestmeta.ContextWithRequestMetadataMap(localCtx, requestContextMap)
 	}
 	// Get clientIP(s) from Context and add it to localCtx
 	source := util.GetSourceIPsFromOutgoingCtx(ctx)