From 120880635a75a36c03c2bfb091afa7e2a761afc0 Mon Sep 17 00:00:00 2001
From: Filip Petkovski <filip.petkovsky@gmail.com>
Date: Thu, 23 May 2024 13:47:08 +0200
Subject: [PATCH] Add request ID to gRPC and HTTP client spans

This commit adds the request ID as a span tag to outgoing (client)
http and gRPC requests. This would allow easier correlation
of traces and logs using the request ID.

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>
---
 cmd/thanos/query_frontend.go | 16 ++++++++--------
 pkg/api/api.go               |  8 ++++----
 pkg/query/remote_engine.go   | 25 +++++++++++++++----------
 pkg/tracing/http.go          |  5 +++++
 4 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/cmd/thanos/query_frontend.go b/cmd/thanos/query_frontend.go
index 5fa7cf3c5e..78d1ec3744 100644
--- a/cmd/thanos/query_frontend.go
+++ b/cmd/thanos/query_frontend.go
@@ -350,19 +350,19 @@ func runQueryFrontend(
 				if !cfg.webDisableCORS {
 					api.SetCORS(w)
 				}
-				tracing.HTTPMiddleware(
-					tracer,
-					name,
-					logger,
-					ins.NewHandler(
+				middleware.RequestID(
+					tracing.HTTPMiddleware(
+						tracer,
 						name,
-						gzhttp.GzipHandler(
-							middleware.RequestID(
+						logger,
+						ins.NewHandler(
+							name,
+							gzhttp.GzipHandler(
 								logMiddleware.HTTPMiddleware(name, f),
 							),
 						),
+						// Cortex frontend middlewares require orgID.
 					),
-					// Cortex frontend middlewares require orgID.
 				).ServeHTTP(w, r.WithContext(user.InjectOrgID(r.Context(), orgId)))
 			})
 			return hf
diff --git a/pkg/api/api.go b/pkg/api/api.go
index d9dc5f4700..6c2ef037f6 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -221,10 +221,10 @@ func GetInstr(
 			}
 		})
 
-		return tracing.HTTPMiddleware(tracer, name, logger,
-			ins.NewHandler(name,
-				gzhttp.GzipHandler(
-					middleware.RequestID(
+		return middleware.RequestID(
+			tracing.HTTPMiddleware(tracer, name, logger,
+				ins.NewHandler(name,
+					gzhttp.GzipHandler(
 						logMiddleware.HTTPMiddleware(name, hf),
 					),
 				),
diff --git a/pkg/query/remote_engine.go b/pkg/query/remote_engine.go
index 16e4b7cb31..f1cf286e0d 100644
--- a/pkg/query/remote_engine.go
+++ b/pkg/query/remote_engine.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/tracing"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/model/labels"
@@ -25,6 +26,7 @@ import (
 
 	"github.com/thanos-io/thanos/pkg/api/query/querypb"
 	"github.com/thanos-io/thanos/pkg/info/infopb"
+	"github.com/thanos-io/thanos/pkg/server/http/middleware"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb/prompb"
 )
@@ -231,17 +233,20 @@ func (r *remoteQuery) Exec(ctx context.Context) *promql.Result {
 	r.cancel = cancel
 	defer cancel()
 
-	queryRange := r.end.Sub(r.start)
-	span, qctx := opentracing.StartSpanFromContext(qctx, "remote_query_exec", opentracing.Tags{
-		"query":            r.plan.String(),
-		"remote_address":   r.remoteAddr,
-		"start":            r.start.UTC().String(),
-		"end":              r.end.UTC().String(),
-		"interval_seconds": r.interval.Seconds(),
-		"range_seconds":    queryRange.Seconds(),
-		"range_human":      queryRange,
+	var (
+		queryRange   = r.end.Sub(r.start)
+		requestID, _ = middleware.RequestIDFromContext(qctx)
+	)
+	qctx = tracing.ClientAddContextTags(qctx, opentracing.Tags{
+		"query.expr":             r.plan.String(),
+		"query.remote_address":   r.remoteAddr,
+		"query.start":            r.start.UTC().String(),
+		"query.end":              r.end.UTC().String(),
+		"query.interval_seconds": r.interval.Seconds(),
+		"query.range_seconds":    queryRange.Seconds(),
+		"query.range_human":      queryRange,
+		"request_id":             requestID,
 	})
-	defer span.Finish()
 
 	var maxResolution int64
 	if r.opts.AutoDownsample {
diff --git a/pkg/tracing/http.go b/pkg/tracing/http.go
index 21b6663f4f..0ea07c9e00 100644
--- a/pkg/tracing/http.go
+++ b/pkg/tracing/http.go
@@ -14,6 +14,8 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
+
+	"github.com/thanos-io/thanos/pkg/server/http/middleware"
 	"github.com/thanos-io/thanos/pkg/tracing/migration"
 )
 
@@ -32,6 +34,9 @@ func HTTPMiddleware(tracer opentracing.Tracer, name string, logger log.Logger, n
 		}
 
 		opts := []opentracing.StartSpanOption{ext.RPCServerOption(wireContext)}
+		if requestID, ok := middleware.RequestIDFromContext(r.Context()); ok {
+			opts = append(opts, opentracing.Tag{Key: "request_id", Value: requestID})
+		}
 		// Check for force tracing header and add it as a tag at the start of span.
 		// This is required for the OpenTelemetry sampler to force tracing.
 		if r.Header.Get(ForceTracingBaggageKey) != "" {