diff --git a/go/trace/fake.go b/go/trace/fake.go index 27145509ad2..715ea3e7ee2 100644 --- a/go/trace/fake.go +++ b/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)) { } diff --git a/go/trace/fake_test.go b/go/trace/fake_test.go index d7d01333202..66f46925224 100644 --- a/go/trace/fake_test.go +++ b/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) } diff --git a/go/trace/opentracing.go b/go/trace/opentracing.go index a0d1b9881a6..ea4af756643 100644 --- a/go/trace/opentracing.go +++ b/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 diff --git a/go/trace/opentracing_test.go b/go/trace/opentracing_test.go index 95559280d2a..0067704cd4e 100644 --- a/go/trace/opentracing_test.go +++ b/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() diff --git a/go/trace/trace.go b/go/trace/trace.go index 6363a8c59de..798233cb250 100644 --- a/go/trace/trace.go +++ b/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) diff --git a/go/trace/trace_test.go b/go/trace/trace_test.go index 1c7b8205fc2..db0e0d360ae 100644 --- a/go/trace/trace_test.go +++ b/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) {