From f6be3ce7130990ae68eae2eaf1e1e448ae7f16b0 Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Wed, 19 Jun 2019 23:33:16 -0700
Subject: [PATCH] Allow setting tracer in the context. (#672)

Allows people to avoid setting a global if they don't want to.
---
 trace/trace.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/trace/trace.go b/trace/trace.go
index 80a204209..e094377e7 100644
--- a/trace/trace.go
+++ b/trace/trace.go
@@ -38,9 +38,25 @@ var (
 	T Tracer = nopTracer{}
 )
 
+type tracerKey struct{}
+
+// WithTracer sets the Tracer which will be used by `StartSpan` in the context.
+func WithTracer(ctx context.Context, t Tracer) context.Context {
+	return context.WithValue(ctx, tracerKey{}, t)
+}
+
 // StartSpan starts a span from the configured default tracer
 func StartSpan(ctx context.Context, name string) (context.Context, Span) {
-	ctx, span := T.StartSpan(ctx, name)
+	t := ctx.Value(tracerKey{})
+
+	var tracer Tracer
+	if t != nil {
+		tracer = t.(Tracer)
+	} else {
+		tracer = T
+	}
+
+	ctx, span := tracer.StartSpan(ctx, name)
 	ctx = log.WithLogger(ctx, span.Logger())
 	return ctx, span
 }