ddtrace/tracer: setup GLS-stored context (#2761)

Signed-off-by: Eliott Bouhana <[email protected]>

contrib/99designs/gqlgen/tracer.go

@@ -104,12 +104,12 @@ func (t *gqlTracer) Validate(_ graphql.ExecutableSchema) error {
 func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.OperationHandler) graphql.ResponseHandler {
 	opCtx := graphql.GetOperationContext(ctx)
 	span, ctx := t.createRootSpan(ctx, opCtx)
-	ctx, req := graphqlsec.StartRequestOperation(ctx, nil /* root */, span, types.RequestOperationArgs{
+	ctx, req := graphqlsec.StartRequestOperation(ctx, span, types.RequestOperationArgs{
 		RawQuery:      opCtx.RawQuery,
 		OperationName: opCtx.OperationName,
 		Variables:     opCtx.Variables,
 	})
-	ctx, query := graphqlsec.StartExecutionOperation(ctx, req, span, types.ExecutionOperationArgs{
+	ctx, query := graphqlsec.StartExecutionOperation(ctx, span, types.ExecutionOperationArgs{
 		Query:         opCtx.RawQuery,
 		OperationName: opCtx.OperationName,
 		Variables:     opCtx.Variables,
@@ -167,8 +167,7 @@ func (t *gqlTracer) InterceptField(ctx context.Context, next graphql.Resolver) (
 
 	span, ctx := tracer.StartSpanFromContext(ctx, fieldOp, opts...)
 	defer func() { span.Finish(tracer.WithError(err)) }()
-
-	ctx, op := graphqlsec.StartResolveOperation(ctx, graphqlsec.FromContext[*types.ExecutionOperation](ctx), span, types.ResolveOperationArgs{
+	ctx, op := graphqlsec.StartResolveOperation(ctx, span, types.ResolveOperationArgs{
 		Arguments: fieldCtx.Args,
 		TypeName:  fieldCtx.Object,
 		FieldName: fieldCtx.Field.Name,

contrib/graph-gophers/graphql-go/graphql.go

@@ -71,12 +71,12 @@ func (t *Tracer) TraceQuery(ctx context.Context, queryString, operationName stri
 	}
 	span, ctx := ddtracer.StartSpanFromContext(ctx, t.cfg.querySpanName, opts...)
 
-	ctx, request := graphqlsec.StartRequestOperation(ctx, nil, span, types.RequestOperationArgs{
+	ctx, request := graphqlsec.StartRequestOperation(ctx, span, types.RequestOperationArgs{
 		RawQuery:      queryString,
 		OperationName: operationName,
 		Variables:     variables,
 	})
-	ctx, query := graphqlsec.StartExecutionOperation(ctx, request, span, types.ExecutionOperationArgs{
+	ctx, query := graphqlsec.StartExecutionOperation(ctx, span, types.ExecutionOperationArgs{
 		Query:         queryString,
 		OperationName: operationName,
 		Variables:     variables,
@@ -120,7 +120,7 @@ func (t *Tracer) TraceField(ctx context.Context, _, typeName, fieldName string,
 	}
 	span, ctx := ddtracer.StartSpanFromContext(ctx, "graphql.field", opts...)
 
-	ctx, field := graphqlsec.StartResolveOperation(ctx, graphqlsec.FromContext[*types.ExecutionOperation](ctx), span, types.ResolveOperationArgs{
+	ctx, field := graphqlsec.StartResolveOperation(ctx, span, types.ResolveOperationArgs{
 		TypeName:  typeName,
 		FieldName: fieldName,
 		Arguments: arguments,

contrib/graphql-go/graphql/graphql.go

@@ -98,7 +98,7 @@ func (i datadogExtension) Init(ctx context.Context, params *graphql.Params) cont
 		tracer.Tag(ext.Component, componentName),
 		tracer.Measured(),
 	)
-	ctx, request := graphqlsec.StartRequestOperation(ctx, nil, span, types.RequestOperationArgs{
+	ctx, request := graphqlsec.StartRequestOperation(ctx, span, types.RequestOperationArgs{
 		RawQuery:      params.RequestString,
 		Variables:     params.VariableValues,
 		OperationName: params.OperationName,
@@ -193,7 +193,7 @@ func (i datadogExtension) ExecutionDidStart(ctx context.Context) (context.Contex
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, i.config.analyticsRate))
 	}
 	span, ctx := tracer.StartSpanFromContext(ctx, spanExecute, opts...)
-	ctx, op := graphqlsec.StartExecutionOperation(ctx, graphqlsec.FromContext[*types.RequestOperation](ctx), span, types.ExecutionOperationArgs{
+	ctx, op := graphqlsec.StartExecutionOperation(ctx, span, types.ExecutionOperationArgs{
 		Query:         data.query,
 		OperationName: data.operationName,
 		Variables:     data.variables,
@@ -240,7 +240,7 @@ func (i datadogExtension) ResolveFieldDidStart(ctx context.Context, info *graphq
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, i.config.analyticsRate))
 	}
 	span, ctx := tracer.StartSpanFromContext(ctx, spanResolve, opts...)
-	ctx, op := graphqlsec.StartResolveOperation(ctx, graphqlsec.FromContext[*types.ExecutionOperation](ctx), span, types.ResolveOperationArgs{
+	ctx, op := graphqlsec.StartResolveOperation(ctx, span, types.ResolveOperationArgs{
 		TypeName:  info.ParentType.Name(),
 		FieldName: info.FieldName,
 		Arguments: collectArguments(info),

ddtrace/tracer/context.go

@@ -11,11 +11,12 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	traceinternal "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 // ContextWithSpan returns a copy of the given context which includes the span s.
 func ContextWithSpan(ctx context.Context, s Span) context.Context {
-	return context.WithValue(ctx, internal.ActiveSpanKey, s)
+	return orchestrion.CtxWithValue(ctx, internal.ActiveSpanKey, s)
 }
 
 // SpanFromContext returns the span contained in the given context. A second return
@@ -25,7 +26,7 @@ func SpanFromContext(ctx context.Context) (Span, bool) {
 	if ctx == nil {
 		return &traceinternal.NoopSpan{}, false
 	}
-	v := ctx.Value(internal.ActiveSpanKey)
+	v := orchestrion.WrapContext(ctx).Value(internal.ActiveSpanKey)
 	if s, ok := v.(ddtrace.Span); ok {
 		return s, true
 	}

ddtrace/tracer/span.go

@@ -29,6 +29,7 @@ import (
 	sharedinternal "gopkg.in/DataDog/dd-trace-go.v1/internal"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/traceprof"
 
@@ -503,6 +504,7 @@ func (s *span) Finish(opts ...ddtrace.FinishOption) {
 	}
 
 	s.finish(t)
+	orchestrion.GLSPopValue(sharedinternal.ActiveSpanKey)
 }
 
 // SetOperationName sets or changes the operation name.

internal/appsec/dyngo/operation.go

@@ -21,11 +21,12 @@
 package dyngo
 
 import (
+	"context"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 	"sync"
 
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
-
 	"go.uber.org/atomic"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 )
 
 // Operation interface type allowing to register event listeners to the
@@ -61,6 +62,9 @@ type ResultOf[O Operation] interface {
 // dispatch calls to the underlying event listener function.
 type EventListener[O Operation, T any] func(O, T)
 
+// contextKey is used to store in a context.Context the ongoing Operation
+type contextKey struct{}
+
 // Atomic *Operation so we can atomically read or swap it.
 var rootOperation atomic.Pointer[Operation]
 
@@ -86,6 +90,10 @@ type operation struct {
 
 	disabled bool
 	mu       sync.RWMutex
+
+	// inContext is used to determine if RegisterOperation was called to put the Operation in the context tree.
+	// If so we need to remove it from the context tree when the Operation is finished.
+	inContext bool
 }
 
 func (o *operation) Parent() Operation {
@@ -140,6 +148,17 @@ func NewOperation(parent Operation) Operation {
 	return &operation{parent: parentOp}
 }
 
+// FromContext looks into the given context (or the GLS if orchestrion is enabled) for a parent Operation and returns it.
+func FromContext(ctx context.Context) (Operation, bool) {
+	ctx = orchestrion.WrapContext(ctx)
+	if ctx == nil {
+		return nil, false
+	}
+
+	op, ok := ctx.Value(contextKey{}).(Operation)
+	return op, ok
+}
+
 // StartOperation starts a new operation along with its arguments and emits a
 // start event with the operation arguments.
 func StartOperation[O Operation, E ArgOf[O]](op O, args E) {
@@ -150,6 +169,19 @@ func StartOperation[O Operation, E ArgOf[O]](op O, args E) {
 	}
 }
 
+// StartAndRegisterOperation calls StartOperation and returns RegisterOperation result
+func StartAndRegisterOperation[O Operation, E ArgOf[O]](ctx context.Context, op O, args E) context.Context {
+	StartOperation(op, args)
+	return RegisterOperation(ctx, op)
+}
+
+// RegisterOperation registers the operation in the context tree. All operations that plan to have children operations
+// should call this function to ensure the operation is properly linked in the context tree.
+func RegisterOperation(ctx context.Context, op Operation) context.Context {
+	op.unwrap().inContext = true
+	return context.WithValue(ctx, contextKey{}, op)
+}
+
 // FinishOperation finishes the operation along with its results and emits a
 // finish event with the operation results.
 // The operation is then disabled and its event listeners removed.
@@ -160,6 +192,10 @@ func FinishOperation[O Operation, E ResultOf[O]](op O, results E) {
 	o.mu.RLock()
 	defer o.mu.RUnlock() // Deferred and stacked on top of the previously deferred call to o.disable()
 
+	if o.inContext {
+		orchestrion.GLSPopValue(contextKey{})
+	}
+
 	if o.disabled {
 		return
 	}

internal/appsec/emitter/graphqlsec/execution.go

@@ -11,6 +11,7 @@ package graphqlsec
 
 import (
 	"context"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/graphqlsec/types"
@@ -20,19 +21,22 @@ import (
 // StartExecutionOperation starts a new GraphQL query operation, along with the given arguments, and
 // emits a start event up in the operation stack. The operation is tracked on the returned context,
 // and can be extracted later on using FromContext.
-func StartExecutionOperation(ctx context.Context, parent *types.RequestOperation, span trace.TagSetter, args types.ExecutionOperationArgs) (context.Context, *types.ExecutionOperation) {
+func StartExecutionOperation(ctx context.Context, span trace.TagSetter, args types.ExecutionOperationArgs) (context.Context, *types.ExecutionOperation) {
 	if span == nil {
 		// The span may be nil (e.g: in case of GraphQL subscriptions with certian contribs). Child
 		// operations might have spans however... and these should be used then.
 		span = trace.NoopTagSetter{}
 	}
 
+	parent, ok := dyngo.FromContext(ctx)
+	if !ok {
+		log.Debug("appsec: StartExecutionOperation: no parent operation found in context")
+	}
+
 	op := &types.ExecutionOperation{
 		Operation: dyngo.NewOperation(parent),
 		TagSetter: span,
 	}
-	newCtx := contextWithValue(ctx, op)
-	dyngo.StartOperation(op, args)
 
-	return newCtx, op
+	return dyngo.StartAndRegisterOperation(ctx, op, args), op
 }

internal/appsec/emitter/graphqlsec/request.go

@@ -21,18 +21,16 @@ import (
 // emits a start event up in the operation stack. The operation is usually linked to tge global root
 // operation. The operation is tracked on the returned context, and can be extracted later on using
 // FromContext.
-func StartRequestOperation(ctx context.Context, parent dyngo.Operation, span trace.TagSetter, args types.RequestOperationArgs) (context.Context, *types.RequestOperation) {
+func StartRequestOperation(ctx context.Context, span trace.TagSetter, args types.RequestOperationArgs) (context.Context, *types.RequestOperation) {
 	if span == nil {
 		// The span may be nil (e.g: in case of GraphQL subscriptions with certian contribs)
 		span = trace.NoopTagSetter{}
 	}
 
 	op := &types.RequestOperation{
-		Operation: dyngo.NewOperation(parent),
+		Operation: dyngo.NewOperation(nil),
 		TagSetter: span,
 	}
-	newCtx := contextWithValue(ctx, op)
-	dyngo.StartOperation(op, args)
 
-	return newCtx, op
+	return dyngo.StartAndRegisterOperation(ctx, op, args), op
 }

internal/appsec/emitter/graphqlsec/resolve.go

@@ -7,6 +7,7 @@ package graphqlsec
 
 import (
 	"context"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/graphqlsec/types"
@@ -16,13 +17,15 @@ import (
 // StartResolveOperation starts a new GraphQL Resolve operation, along with the given arguments, and
 // emits a start event up in the operation stack. The operation is tracked on the returned context,
 // and can be extracted later on using FromContext.
-func StartResolveOperation(ctx context.Context, parent *types.ExecutionOperation, span trace.TagSetter, args types.ResolveOperationArgs) (context.Context, *types.ResolveOperation) {
+func StartResolveOperation(ctx context.Context, span trace.TagSetter, args types.ResolveOperationArgs) (context.Context, *types.ResolveOperation) {
+	parent, ok := dyngo.FromContext(ctx)
+	if !ok {
+		log.Debug("appsec: StartResolveOperation: no parent operation found in context")
+	}
+
 	op := &types.ResolveOperation{
 		Operation: dyngo.NewOperation(parent),
 		TagSetter: span,
 	}
-	newCtx := contextWithValue(ctx, op)
-	dyngo.StartOperation(op, args)
-
-	return newCtx, op
+	return dyngo.StartAndRegisterOperation(ctx, op, args), op
 }

internal/appsec/emitter/grpcsec/grpc.go

@@ -14,7 +14,6 @@ import (
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/grpcsec/types"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/trace"
 )
 
@@ -27,12 +26,10 @@ func StartHandlerOperation(ctx context.Context, args types.HandlerOperationArgs,
 		Operation:  dyngo.NewOperation(parent),
 		TagsHolder: trace.NewTagsHolder(),
 	}
-	newCtx := context.WithValue(ctx, listener.ContextKey{}, op)
 	for _, cb := range setup {
 		cb(op)
 	}
-	dyngo.StartOperation(op, args)
-	return newCtx, op
+	return dyngo.StartAndRegisterOperation(ctx, op, args), op
 }
 
 // StartReceiveOperation starts a receive operation of a gRPC handler, along

internal/appsec/emitter/httpsec/http.go

@@ -12,18 +12,17 @@ package httpsec
 
 import (
 	"context"
-	"gopkg.in/DataDog/dd-trace-go.v1/appsec/events"
 
 	// Blank import needed to use embed for the default blocked response payloads
 	_ "embed"
 	"net/http"
 	"strings"
 
+	"gopkg.in/DataDog/dd-trace-go.v1/appsec/events"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/httpsec/types"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/sharedsec"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/trace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/trace/httptrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
@@ -36,7 +35,7 @@ import (
 // This function should not be called when AppSec is disabled in order to
 // get preciser error logs.
 func MonitorParsedBody(ctx context.Context, body any) error {
-	parent, _ := ctx.Value(listener.ContextKey{}).(*types.Operation)
+	parent, _ := dyngo.FromContext(ctx)
 	if parent == nil {
 		log.Error("appsec: parsed http body monitoring ignored: could not find the http handler instrumentation metadata in the request context: the request handler is not being monitored by a middleware function or the provided context is not the expected request context")
 		return nil
@@ -195,10 +194,9 @@ func StartOperation(ctx context.Context, args types.HandlerOperationArgs, setup
 		Operation:  dyngo.NewOperation(nil),
 		TagsHolder: trace.NewTagsHolder(),
 	}
-	newCtx := context.WithValue(ctx, listener.ContextKey{}, op)
 	for _, cb := range setup {
 		cb(op)
 	}
-	dyngo.StartOperation(op, args)
-	return newCtx, op
+
+	return dyngo.StartAndRegisterOperation(ctx, op, args), op
 }