ddtrace: support v1 behaviour for mock spans

DataDog · Feb 22, 2024 · c2a6279 · c2a6279
1 parent a0bbf07
commit c2a6279
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 30 deletions.
diff --git a/ddtrace/mocktracer/mockspan.go b/ddtrace/mocktracer/mockspan.go
@@ -57,28 +57,30 @@ type MockspanV2Adapter struct {
 
 // BaggageItem implements ddtrace.Span.
 func (msa MockspanV2Adapter) BaggageItem(key string) string {
-	// no-op
-	return ""
+	return msa.Span.Unwrap().BaggageItem(key)
 }
 
 // Finish implements ddtrace.Span.
-func (MockspanV2Adapter) Finish(opts ...ddtrace.FinishOption) {
-	// no-op
+func (msa MockspanV2Adapter) Finish(opts ...ddtrace.FinishOption) {
+	t := internal.GetGlobalTracer().(internal.TracerV2Adapter)
+	sp := msa.Span.Unwrap()
+	t.Tracer.(v2.Tracer).FinishSpan(sp)
+	sp.Finish()
 }
 
 // SetBaggageItem implements ddtrace.Span.
-func (MockspanV2Adapter) SetBaggageItem(key string, val string) {
-	// no-op
+func (msa MockspanV2Adapter) SetBaggageItem(key string, val string) {
+	msa.Span.Unwrap().SetBaggageItem(key, val)
 }
 
 // SetOperationName implements ddtrace.Span.
-func (MockspanV2Adapter) SetOperationName(operationName string) {
-	// no-op
+func (msa MockspanV2Adapter) SetOperationName(operationName string) {
+	msa.Span.Unwrap().SetOperationName(operationName)
 }
 
 // SetTag implements ddtrace.Span.
-func (MockspanV2Adapter) SetTag(key string, value interface{}) {
-	// no-op
+func (msa MockspanV2Adapter) SetTag(key string, value interface{}) {
+	msa.Span.SetTag(key, value)
 }
 
 // Context implements Span.

diff --git a/ddtrace/tracer/context.go b/ddtrace/tracer/context.go
@@ -16,32 +16,30 @@ import (
 
 // ContextWithSpan returns a copy of the given context which includes the span s.
 func ContextWithSpan(ctx context.Context, s Span) context.Context {
-	if ctx == nil {
-		ctx = context.Background()
+	switch s := s.(type) {
+	case internal.SpanV2Adapter:
+		return v2.ContextWithSpan(ctx, s.Span)
+	case mocktracer.MockspanV2Adapter:
+		return v2.ContextWithSpan(ctx, s.Span.Unwrap())
+	case internal.NoopSpan:
+		return v2.ContextWithSpan(ctx, nil)
 	}
-	return context.WithValue(ctx, internal.ActiveSpanKey, s)
+	// TODO: remove this case once we remove the v1 tracer
+	return ctx
 }
 
 // SpanFromContext returns the span contained in the given context. A second return
 // value indicates if a span was found in the context. If no span is found, a no-op
 // span is returned.
 func SpanFromContext(ctx context.Context) (Span, bool) {
-	if ctx == nil {
+	s, ok := v2.SpanFromContext(ctx)
+	if !ok {
 		return internal.NoopSpan{}, false
 	}
-	v := ctx.Value(internal.ActiveSpanKey)
-	switch v.(type) {
-	case internal.SpanV2Adapter:
-		sa := v.(internal.SpanV2Adapter)
-		if mocktracer.IsActive() {
-			return mocktracer.MockspanV2Adapter{Span: v2mock.MockSpan(sa.Span)}, true
-		}
-		return sa, true
-	case Span:
-		return v.(Span), true
-	default:
-		return internal.NoopSpan{}, false
+	if mocktracer.IsActive() {
+		return mocktracer.MockspanV2Adapter{Span: v2mock.MockSpan(s)}, true
 	}
+	return internal.SpanV2Adapter{Span: s}, true
 }
 
 // StartSpanFromContext returns a new span with the given operation name and options. If a span
@@ -50,7 +48,12 @@ func SpanFromContext(ctx context.Context) (Span, bool) {
 func StartSpanFromContext(ctx context.Context, operationName string, opts ...StartSpanOption) (Span, context.Context) {
 	cfg := internal.BuildStartSpanConfigV2(opts...)
 	span, ctx := v2.StartSpanFromContext(ctx, operationName, v2.WithStartSpanConfig(cfg))
-	sa := internal.SpanV2Adapter{Span: span}
-	ctx = ContextWithSpan(ctx, sa)
-	return sa, ctx
+	var s Span
+	if mocktracer.IsActive() {
+		s = mocktracer.MockspanV2Adapter{Span: v2mock.MockSpan(span)}
+	} else {
+		s = internal.SpanV2Adapter{Span: span}
+	}
+	ctx = ContextWithSpan(ctx, s)
+	return s, ctx
 }
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	cloud.google.com/go/pubsub v1.33.0
 	github.com/99designs/gqlgen v0.17.36
 	github.com/DataDog/appsec-internal-go v1.4.1
-	github.com/DataDog/dd-trace-go/v2 v2.0.0-20240216140524-914f0c30cab1
+	github.com/DataDog/dd-trace-go/v2 v2.0.0-20240222143418-f576f22a3707
 	github.com/DataDog/dd-trace-go/v2/contrib/99designs/gqlgen v0.0.0-20240216140524-914f0c30cab1
 	github.com/DataDog/dd-trace-go/v2/contrib/emicklei/go-restful.v3 v0.0.0-20240216140524-914f0c30cab1
 	github.com/DataDog/dd-trace-go/v2/contrib/gin-gonic/gin v0.0.0-20240216140524-914f0c30cab1

diff --git a/go.sum b/go.sum
@@ -636,6 +636,10 @@ github.com/DataDog/datadog-go/v5 v5.4.0 h1:Ea3eXUVwrVV28F/fo3Dr3aa+TL/Z7Xi6SUPKW
 github.com/DataDog/datadog-go/v5 v5.4.0/go.mod h1:K9kcYBlxkcPP8tvvjZZKs/m1edNAUFzBbdpTUKfCsuw=
 github.com/DataDog/dd-trace-go/v2 v2.0.0-20240216140524-914f0c30cab1 h1:EAL7x/fqiHzbTGELq1iiadwQu+r7PQ01ipXj1d7k6Nc=
 github.com/DataDog/dd-trace-go/v2 v2.0.0-20240216140524-914f0c30cab1/go.mod h1:Xh0wvissbVfhJH5piMwhufBX6Os+EcjSilgnT9vRBmg=
+github.com/DataDog/dd-trace-go/v2 v2.0.0-20240222142412-f8d4db314c5e h1:k5fJ+rN5M3Y+I3dSQZrhtt3pqdFurEIEmlK8zRNyGz4=
+github.com/DataDog/dd-trace-go/v2 v2.0.0-20240222142412-f8d4db314c5e/go.mod h1:guG8Q4usCSwiYEsEMSyST8wosYKU9nNmpFy/4B/PyT0=
+github.com/DataDog/dd-trace-go/v2 v2.0.0-20240222143418-f576f22a3707 h1:/Vun9ABddQdngcSeDdEGUM/WYr2LYWpYQANO1NfdIj0=
+github.com/DataDog/dd-trace-go/v2 v2.0.0-20240222143418-f576f22a3707/go.mod h1:guG8Q4usCSwiYEsEMSyST8wosYKU9nNmpFy/4B/PyT0=
 github.com/DataDog/dd-trace-go/v2/contrib/99designs/gqlgen v0.0.0-20240216140524-914f0c30cab1 h1:tXKQV0O3XCFwZgintpKq3vnZqaMChY/vMCcXNaclHes=
 github.com/DataDog/dd-trace-go/v2/contrib/99designs/gqlgen v0.0.0-20240216140524-914f0c30cab1/go.mod h1:pmT1ExCMPtowsPWs493lMOYBu+nzSspKstJkN2gmgsk=
 github.com/DataDog/dd-trace-go/v2/contrib/emicklei/go-restful.v3 v0.0.0-20240216140524-914f0c30cab1 h1:KjUXtUuQR2kZ15BPDT5cu017abAVeRWP9kLcsMXZjS4=