Preserve tracing context

Signed-off-by: Kemal Akkoyun <kakkoyun@gmail.com>

pkg/query/querier.go

@@ -197,7 +197,9 @@ func (q *querier) Select(_ bool, hints *storage.SelectHints, ms ...*labels.Match
 	}
 
 	// The querier has a context but it gets cancelled, as soon as query evaluation is completed, by the engine.
-	ctx, cancel := context.WithTimeout(context.Background(), q.selectTimeout)
+	// We want to prevent this from happening for the async storea API calls we make while preserving tracing context.
+	ctx := tracing.CopyTraceContext(context.Background(), q.ctx)
+	ctx, cancel := context.WithTimeout(ctx, q.selectTimeout)
 	span, ctx := tracing.StartSpan(ctx, "querier_select", opentracing.Tags{
 		"minTime":  hints.Start,
 		"maxTime":  hints.End,
@@ -209,7 +211,7 @@ func (q *querier) Select(_ bool, hints *storage.SelectHints, ms ...*labels.Match
 		defer close(promise)
 
 		var err error
-		tracing.DoInSpan(ctx, "querier_select_ismyturn", func(ctx context.Context) {
+		tracing.DoInSpan(ctx, "querier_select_gate_ismyturn", func(ctx context.Context) {
 			err = q.selectGate.Start(ctx)
 		})
 		if err != nil {

pkg/tracing/tracing.go

@@ -29,6 +29,7 @@ func ContextWithTracer(ctx context.Context, tracer opentracing.Tracer) context.C
 	return context.WithValue(ctx, tracerKey, tracer)
 }
 
+// tracerFromContext extracts opentracing.Tracer from the given context.
 func tracerFromContext(ctx context.Context) opentracing.Tracer {
 	val := ctx.Value(tracerKey)
 	if sp, ok := val.(opentracing.Tracer); ok {
@@ -37,6 +38,15 @@ func tracerFromContext(ctx context.Context) opentracing.Tracer {
 	return nil
 }
 
+// CopyTraceContext copies the necessary trace context from given source context to target context.
+func CopyTraceContext(trgt, src context.Context) context.Context {
+	ctx := ContextWithTracer(trgt, tracerFromContext(src))
+	if parentSpan := opentracing.SpanFromContext(src); parentSpan != nil {
+		ctx = opentracing.ContextWithSpan(ctx, parentSpan)
+	}
+	return ctx
+}
+
 // StartSpan starts and returns span with `operationName` and hooking as child to a span found within given context if any.
 // It uses opentracing.Tracer propagated in context. If no found, it uses noop tracer without notification.
 func StartSpan(ctx context.Context, operationName string, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {