Change tracingService methods to use contexts

go/trace/fake.go

@@ -25,9 +25,13 @@ import (
 
 type noopTracingServer struct{}
 
-func (noopTracingServer) New(Span, string) Span                                     { return NoopSpan{} }
-func (noopTracingServer) FromContext(context.Context) (Span, bool)                  { return nil, false }
-func (noopTracingServer) NewFromString(parent, label string) (Span, error)          { return NoopSpan{}, nil }
+func (noopTracingServer) New(ctx context.Context, label string) (Span, context.Context) {
+	return NoopSpan{}, ctx
+}
+func (noopTracingServer) FromContext(context.Context) (Span, bool) { return nil, false }
+func (noopTracingServer) NewFromString(ctx context.Context, parent, label string) (Span, context.Context, error) {
+	return NoopSpan{}, ctx, nil
+}
 func (noopTracingServer) NewContext(parent context.Context, _ Span) context.Context { return parent }
 func (noopTracingServer) AddGrpcServerOptions(addInterceptors func(s grpc.StreamServerInterceptor, u grpc.UnaryServerInterceptor)) {
 }

go/trace/fake_test.go

@@ -27,7 +27,8 @@ func TestNoopTracingServer(t *testing.T) {
 	tracingSvc, closer, err := factoryFunc("value")
 	require.NoError(t, err)
 	require.NoError(t, closer.Close())
-	span, err := tracingSvc.NewFromString("parent", "label")
+	span, ctx, err := tracingSvc.NewFromString(t.Context(), "parent", "label")
 	require.NoError(t, err)
 	require.Empty(t, span)
+	require.Equal(t, t.Context(), ctx)
 }

go/trace/opentracing.go

@@ -67,7 +67,8 @@ func (jf openTracingService) AddGrpcClientOptions(addInterceptors func(s grpc.St
 }
 
 // New is part of an interface implementation
-func (jf openTracingService) New(parent Span, label string) Span {
+func (jf openTracingService) New(ctx context.Context, label string) (Span, context.Context) {
+	parent, _ := jf.FromContext(ctx)
 	var innerSpan opentracing.Span
 	if parent == nil {
 		innerSpan = jf.Tracer.GetOpenTracingTracer().StartSpan(label)
@@ -76,7 +77,8 @@ func (jf openTracingService) New(parent Span, label string) Span {
 		span := jaegerParent.otSpan
 		innerSpan = jf.Tracer.GetOpenTracingTracer().StartSpan(label, opentracing.ChildOf(span.Context()))
 	}
-	return openTracingSpan{otSpan: innerSpan}
+	span := openTracingSpan{otSpan: innerSpan}
+	return span, jf.NewContext(ctx, span)
 }
 
 func extractMapFromString(in string) (opentracing.TextMapCarrier, error) {
@@ -94,17 +96,18 @@ func extractMapFromString(in string) (opentracing.TextMapCarrier, error) {
 	return dat, nil
 }
 
-func (jf openTracingService) NewFromString(parent, label string) (Span, error) {
+func (jf openTracingService) NewFromString(ctx context.Context, parent, label string) (Span, context.Context, error) {
 	carrier, err := extractMapFromString(parent)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	spanContext, err := jf.Tracer.GetOpenTracingTracer().Extract(opentracing.TextMap, carrier)
 	if err != nil {
-		return nil, vterrors.Wrap(err, "failed to deserialize span context")
+		return nil, nil, vterrors.Wrap(err, "failed to deserialize span context")
 	}
 	innerSpan := jf.Tracer.GetOpenTracingTracer().StartSpan(label, opentracing.ChildOf(spanContext))
-	return openTracingSpan{otSpan: innerSpan}, nil
+	span := openTracingSpan{otSpan: innerSpan}
+	return span, jf.NewContext(ctx, span), nil
 }
 
 // FromContext is part of an interface implementation

go/trace/opentracing_test.go

@@ -54,17 +54,19 @@ func TestNewSpan(t *testing.T) {
 	svc := openTracingService{
 		Tracer: &fakeTracer{},
 	}
-	clientSpan := svc.New(nil, "test-label")
+	clientSpan, ctx := svc.New(t.Context(), "test-label")
 	require.NotEmpty(t, clientSpan)
+	require.NotNil(t, ctx)
 
-	clientSpan = svc.New(clientSpan, "client-span")
+	clientSpan, ctx = svc.New(ctx, "client-span")
 	require.NotEmpty(t, clientSpan)
+	require.NotNil(t, ctx)
 
 	spanFromCtx, ok := svc.FromContext(context.Background())
 	require.False(t, ok)
 	require.Nil(t, spanFromCtx)
 
-	ctx := svc.NewContext(context.TODO(), clientSpan)
+	ctx = svc.NewContext(t.Context(), clientSpan)
 	require.NotNil(t, ctx)
 	clientSpan.Finish()
 

go/trace/trace.go

@@ -46,23 +46,14 @@ type Span interface {
 
 // NewSpan creates a new Span with the currently installed tracing plugin.
 // If no tracing plugin is installed, it returns a fake Span that does nothing.
-func NewSpan(inCtx context.Context, label string) (Span, context.Context) {
-	parent, _ := currentTracer.FromContext(inCtx)
-	span := currentTracer.New(parent, label)
-	outCtx := currentTracer.NewContext(inCtx, span)
-
-	return span, outCtx
+func NewSpan(ctx context.Context, label string) (Span, context.Context) {
+	return currentTracer.New(ctx, label)
 }
 
 // NewFromString creates a new Span with the currently installed tracing plugin, extracting the span context from
 // the provided string.
 func NewFromString(inCtx context.Context, parent, label string) (Span, context.Context, error) {
-	span, err := currentTracer.NewFromString(parent, label)
-	if err != nil {
-		return nil, nil, err
-	}
-	outCtx := currentTracer.NewContext(inCtx, span)
-	return span, outCtx, nil
+	return currentTracer.NewFromString(inCtx, parent, label)
 }
 
 // AnnotateSQL annotates information about a sql query in the span. This is done in a way
@@ -104,10 +95,10 @@ func AddGrpcClientOptions(addInterceptors func(s grpc.StreamClientInterceptor, u
 // tracingService is an interface for creating spans or extracting them from Contexts.
 type tracingService interface {
 	// New creates a new span from an existing one, if provided. The parent can also be nil
-	New(parent Span, label string) Span
+	New(ctx context.Context, label string) (Span, context.Context)
 
 	// NewFromString creates a new span and uses the provided string to reconstitute the parent span
-	NewFromString(parent, label string) (Span, error)
+	NewFromString(ctx context.Context, parent, label string) (Span, context.Context, error)
 
 	// FromContext extracts a span from a context, making it possible to annotate the span with additional information
 	FromContext(ctx context.Context) (Span, bool)

go/trace/trace_test.go

@@ -134,17 +134,19 @@ func (f *fakeTracer) GetOpenTracingTracer() opentracing.Tracer {
 	return opentracing.GlobalTracer()
 }
 
-func (f *fakeTracer) NewFromString(parent, label string) (Span, error) {
+func (f *fakeTracer) NewFromString(ctx context.Context, parent, label string) (Span, context.Context, error) {
 	if parent == "" {
-		return &mockSpan{tracer: f}, errors.New("parent is empty")
+		return nil, nil, errors.New("parent is empty")
 	}
-	return &mockSpan{tracer: f}, nil
+	span := &mockSpan{tracer: f}
+	return span, f.NewContext(ctx, span), nil
 }
 
-func (f *fakeTracer) New(parent Span, label string) Span {
+func (f *fakeTracer) New(ctx context.Context, label string) (Span, context.Context) {
 	f.log = append(f.log, "span started")
 
-	return &mockSpan{tracer: f}
+	span := &mockSpan{tracer: f}
+	return span, f.NewContext(ctx, span)
 }
 
 func (f *fakeTracer) FromContext(ctx context.Context) (Span, bool) {