add request ID injection to context for ruler queries to track split queries

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

Author: erlan-z

pkg/api/middlewares.go

@@ -46,7 +46,7 @@ func (h HTTPHeaderMiddleware) InjectRequestIdHeader(r *http.Request) *http.Reque
 	if reqId == "" {
 		reqId = uuid.NewString()
 	}
-	ctx := requestutil.ContextWithRequestIdHeader(r.Context(), reqId)
+	ctx := requestutil.ContextWithRequestId(r.Context(), reqId)
 	return r.WithContext(ctx)
 }
 

pkg/api/middlewares_test.go

@@ -122,7 +122,7 @@ func TestExistingRequestIdIsPreserved(t *testing.T) {
 		RequestIdHeader: "X-Request-ID",
 	}
 
-	ctx := requestutil.ContextWithRequestIdHeader(context.Background(), existingID)
+	ctx := requestutil.ContextWithRequestId(context.Background(), existingID)
 
 	h := http.Header{}
 	h.Add("X-Request-ID", "should-be-ignored")

pkg/querier/worker/frontend_processor.go

@@ -143,7 +143,7 @@ func (fp *frontendProcessor) runRequest(ctx context.Context, request *httpgrpc.H
 	}
 	ctx = util_log.ContextWithHeaderMap(ctx, headerMap)
 	requestId := headers[textproto.CanonicalMIMEHeaderKey(requestutil.RequestIdPropagationString)]
-	ctx = requestutil.ContextWithRequestIdHeader(ctx, requestId)
+	ctx = requestutil.ContextWithRequestId(ctx, requestId)
 
 	logger := util_log.WithContext(ctx, fp.log)
 	if statsEnabled {

pkg/querier/worker/scheduler_processor.go

@@ -151,7 +151,7 @@ func (sp *schedulerProcessor) querierLoop(c schedulerpb.SchedulerForQuerier_Quer
 			}
 			ctx = util_log.ContextWithHeaderMap(ctx, headerMap)
 			requestId := headers[textproto.CanonicalMIMEHeaderKey(requestutil.RequestIdPropagationString)]
-			ctx = requestutil.ContextWithRequestIdHeader(ctx, requestId)
+			ctx = requestutil.ContextWithRequestId(ctx, requestId)
 
 			tracer := opentracing.GlobalTracer()
 			// Ignore errors here. If we cannot get parent span, we just don't create new one.

pkg/ruler/compat.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
@@ -27,6 +28,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring/client"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	promql_util "github.com/cortexproject/cortex/pkg/util/promql"
+	requestutil "github.com/cortexproject/cortex/pkg/util/request"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
@@ -183,6 +185,9 @@ func EngineQueryFunc(engine promql.QueryEngine, frontendClient *frontendClient,
 			}
 		}
 
+		// Add request ID to the context so that it can be used in logs and metrics for split queries.
+		ctx = requestutil.ContextWithRequestId(ctx, uuid.NewString())
+
 		if frontendClient != nil {
 			v, err := frontendClient.InstantQuery(ctx, qs, t)
 			if err != nil {

pkg/util/grpcutil/util.go

@@ -71,7 +71,7 @@ func extractRequestIdFromMetadata(ctx context.Context) context.Context {
 	if len(requestIds) == 0 {
 		return ctx
 	}
-	return requestutil.ContextWithRequestIdHeader(ctx, requestIds[0])
+	return requestutil.ContextWithRequestId(ctx, requestIds[0])
 }
 
 // HTTPHeaderPropagationClientInterceptor allows for propagation of HTTP Request headers across gRPC calls - works

pkg/util/request/request_id.go

@@ -13,7 +13,7 @@ const (
 	RequestIdPropagationString string = "x-request-id-forwarding-string"
 )
 
-func ContextWithRequestIdHeader(ctx context.Context, reqId string) context.Context {
+func ContextWithRequestId(ctx context.Context, reqId string) context.Context {
 	return context.WithValue(ctx, RequestIdKey, reqId)
 }