From 756bce099aff48f239a9ad280fe42bf00bcc1c9b Mon Sep 17 00:00:00 2001
From: Karsten Jeschkies <karsten.jeschkies@grafana.com>
Date: Mon, 20 Nov 2023 17:15:06 +0100
Subject: [PATCH] Propagate trace ID with HTTP gRPC request. (#11251)

The changes in https://github.com/grafana/loki/pull/10688 did not
propage the trace ID from the context. `Frontend.RoundTripGRPC` would
inject the trace ID into the request. That's not done in `Frontend.Do`.
This changes extends the `codec.EncodeRequest` to inject the trace ID
there. This is more inline with other metadata.

(cherry picked from commit 8d34f857bcb41120fca71948f9b7dc3504047228)
---
 pkg/loki/modules.go                  |  1 +
 pkg/querier/queryrange/codec.go      |  9 +++++++++
 pkg/querier/queryrange/instrument.go | 16 ++++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/pkg/loki/modules.go b/pkg/loki/modules.go
index 4c14a4872655f..e0e8ab4d1f88d 100644
--- a/pkg/loki/modules.go
+++ b/pkg/loki/modules.go
@@ -519,6 +519,7 @@ func (t *Loki) initQuerier() (services.Service, error) {
 	internalMiddlewares := []queryrangebase.Middleware{
 		serverutil.RecoveryMiddleware,
 		queryrange.Instrument{Metrics: t.Metrics},
+		queryrange.Tracer{},
 	}
 	if t.supportIndexDeleteRequest() && t.Cfg.CompactorConfig.RetentionEnabled {
 		internalMiddlewares = append(
diff --git a/pkg/querier/queryrange/codec.go b/pkg/querier/queryrange/codec.go
index d8adbfca2778a..04aefe5e9a839 100644
--- a/pkg/querier/queryrange/codec.go
+++ b/pkg/querier/queryrange/codec.go
@@ -579,6 +579,15 @@ func (c Codec) EncodeRequest(ctx context.Context, r queryrangebase.Request) (*ht
 	}
 	header.Set(user.OrgIDHeaderName, orgID)
 
+	// Propagate trace context in request.
+	tracer, span := opentracing.GlobalTracer(), opentracing.SpanFromContext(ctx)
+	if tracer != nil && span != nil {
+		carrier := opentracing.HTTPHeadersCarrier(header)
+		if err := tracer.Inject(span.Context(), opentracing.HTTPHeaders, carrier); err != nil {
+			return nil, err
+		}
+	}
+
 	switch request := r.(type) {
 	case *LokiRequest:
 		params := url.Values{
diff --git a/pkg/querier/queryrange/instrument.go b/pkg/querier/queryrange/instrument.go
index 8c32fad4ca304..497cfb2dd8a1a 100644
--- a/pkg/querier/queryrange/instrument.go
+++ b/pkg/querier/queryrange/instrument.go
@@ -8,6 +8,7 @@ import (
 	"github.com/grafana/dskit/httpgrpc"
 	"github.com/grafana/dskit/instrument"
 	"github.com/grafana/dskit/middleware"
+	"github.com/opentracing/opentracing-go"
 
 	"github.com/grafana/dskit/server"
 
@@ -52,3 +53,18 @@ func (i Instrument) observe(ctx context.Context, route string, err error, durati
 	}
 	instrument.ObserveWithExemplar(ctx, i.RequestDuration.WithLabelValues(method, route, respStatus, "false"), duration.Seconds())
 }
+
+type Tracer struct{}
+
+var _ queryrangebase.Middleware = Tracer{}
+
+// Wrap implements the queryrangebase.Middleware
+func (t Tracer) Wrap(next queryrangebase.Handler) queryrangebase.Handler {
+	return queryrangebase.HandlerFunc(func(ctx context.Context, r queryrangebase.Request) (queryrangebase.Response, error) {
+		route := DefaultCodec.Path(r)
+		route = middleware.MakeLabelValue(route)
+		span, ctx := opentracing.StartSpanFromContext(ctx, route)
+		defer span.Finish()
+		return next.Do(ctx, r)
+	})
+}