From 8af3bfcdc26dfd0296f6d93208fc5fcbf2cd7d50 Mon Sep 17 00:00:00 2001
From: rghetia <rahulpa@google.com>
Date: Thu, 19 Sep 2019 12:23:07 -0700
Subject: [PATCH] Use trace.SetCurrentSpan to store the span in the context in
 SDK. (#129)

---
 sdk/trace/trace.go  | 12 ------------
 sdk/trace/tracer.go | 10 ++++++----
 2 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/sdk/trace/trace.go b/sdk/trace/trace.go
index 802fc11a3..6d24df006 100644
--- a/sdk/trace/trace.go
+++ b/sdk/trace/trace.go
@@ -15,7 +15,6 @@
 package trace
 
 import (
-	"context"
 	crand "crypto/rand"
 	"encoding/binary"
 	"math/rand"
@@ -62,14 +61,3 @@ func Register() apitrace.Tracer {
 	})
 	return tr
 }
-
-type contextKey struct{}
-
-func fromContext(ctx context.Context) *span {
-	s, _ := ctx.Value(contextKey{}).(*span)
-	return s
-}
-
-func newContext(parent context.Context, s *span) context.Context {
-	return context.WithValue(parent, contextKey{}, s)
-}
diff --git a/sdk/trace/tracer.go b/sdk/trace/tracer.go
index 18c8da09b..a6307bd46 100644
--- a/sdk/trace/tracer.go
+++ b/sdk/trace/tracer.go
@@ -52,9 +52,11 @@ func (tr *tracer) Start(ctx context.Context, name string, o ...apitrace.SpanOpti
 		parent = opts.Reference.SpanContext
 		remoteParent = true
 	} else {
-		if p := fromContext(ctx); p != nil {
-			p.addChild()
-			parent = p.spanContext
+		if p := apitrace.CurrentSpan(ctx); p != nil {
+			if sdkSpan, ok := p.(*span); ok {
+				sdkSpan.addChild()
+				parent = sdkSpan.spanContext
+			}
 		}
 	}
 
@@ -70,7 +72,7 @@ func (tr *tracer) Start(ctx context.Context, name string, o ...apitrace.SpanOpti
 
 	ctx, end := startExecutionTracerTask(ctx, name)
 	span.executionTracerTaskEnd = end
-	return newContext(ctx, span), span
+	return apitrace.SetCurrentSpan(ctx, span), span
 }
 
 func (tr *tracer) WithSpan(ctx context.Context, name string, body func(ctx context.Context) error) error {