orchestrion: setup GLS-stored context

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

ddtrace/tracer/context.go

@@ -11,17 +11,19 @@ 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
 // 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) {
+	ctx = orchestrion.CtxOrGLS(ctx)
 	if ctx == nil {
 		return &traceinternal.NoopSpan{}, false
 	}
@@ -41,6 +43,7 @@ func StartSpanFromContext(ctx context.Context, operationName string, opts ...Sta
 	optsLocal := make([]StartSpanOption, len(opts), len(opts)+2)
 	copy(optsLocal, opts)
 
+	ctx = orchestrion.CtxOrGLS(ctx)
 	if ctx == nil {
 		// default to context.Background() to avoid panics on Go >= 1.15
 		ctx = context.Background()

internal/appsec/emitter/graphqlsec/context.go

@@ -10,6 +10,7 @@ import (
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 type typed[T dyngo.Operation] struct{}
@@ -29,5 +30,5 @@ func FromContext[T dyngo.Operation](ctx context.Context) T {
 
 // contextWithValue creates a new context with the specified operation stored in it.
 func contextWithValue[T dyngo.Operation](ctx context.Context, value T) context.Context {
-	return context.WithValue(ctx, typed[T]{}, value)
+	return orchestrion.CtxWithValue(ctx, typed[T]{}, value)
 }

internal/appsec/emitter/grpcsec/grpc.go

@@ -16,6 +16,7 @@ import (
 	"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"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 // StartHandlerOperation starts an gRPC server handler operation, along with the
@@ -27,7 +28,7 @@ func StartHandlerOperation(ctx context.Context, args types.HandlerOperationArgs,
 		Operation:  dyngo.NewOperation(parent),
 		TagsHolder: trace.NewTagsHolder(),
 	}
-	newCtx := context.WithValue(ctx, listener.ContextKey{}, op)
+	newCtx := orchestrion.CtxWithValue(ctx, listener.ContextKey{}, op)
 	for _, cb := range setup {
 		cb(op)
 	}

internal/appsec/emitter/httpsec/http.go

@@ -12,13 +12,13 @@ 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"
@@ -27,6 +27,7 @@ import (
 	"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"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/stacktrace"
 
 	"github.com/DataDog/appsec-internal-go/netip"
@@ -36,7 +37,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, _ := orchestrion.CtxOrGLS(ctx).Value(listener.ContextKey{}).(*types.Operation)
 	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,7 +196,7 @@ func StartOperation(ctx context.Context, args types.HandlerOperationArgs, setup
 		Operation:  dyngo.NewOperation(nil),
 		TagsHolder: trace.NewTagsHolder(),
 	}
-	newCtx := context.WithValue(ctx, listener.ContextKey{}, op)
+	newCtx := orchestrion.CtxWithValue(ctx, listener.ContextKey{}, op)
 	for _, cb := range setup {
 		cb(op)
 	}

internal/appsec/emitter/httpsec/roundtripper.go

@@ -14,6 +14,7 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/httpsec/types"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 var badInputContextOnce sync.Once
@@ -23,7 +24,7 @@ func ProtectRoundTrip(ctx context.Context, url string) error {
 		URL: url,
 	}
 
-	parent, _ := ctx.Value(listener.ContextKey{}).(dyngo.Operation)
+	parent, _ := orchestrion.CtxOrGLS(ctx).Value(listener.ContextKey{}).(dyngo.Operation)
 	if parent == nil { // No parent operation => we can't monitor the request
 		badInputContextOnce.Do(func() {
 			log.Debug("appsec: outgoing http request monitoring ignored: could not find the handler " +

internal/appsec/emitter/sharedsec/shared.go

@@ -7,12 +7,13 @@ package sharedsec
 
 import (
 	"context"
-	"gopkg.in/DataDog/dd-trace-go.v1/appsec/events"
 	"reflect"
 
+	"gopkg.in/DataDog/dd-trace-go.v1/appsec/events"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 type (
@@ -60,7 +61,7 @@ func (f OnUserIDOperationStart) Call(op dyngo.Operation, v interface{}) {
 // A call to the WAF is made to check the user ID and an error is returned if the
 // user should be blocked. The return value is nil otherwise.
 func MonitorUser(ctx context.Context, userID string) error {
-	if parent, ok := ctx.Value(listener.ContextKey{}).(dyngo.Operation); ok {
+	if parent, ok := orchestrion.CtxOrGLS(ctx).Value(listener.ContextKey{}).(dyngo.Operation); ok {
 		return ExecuteUserIDOperation(parent, UserIDOperationArgs{UserID: userID})
 	}
 	log.Error("appsec: user ID 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")

internal/appsec/emitter/sqlsec/sql.go

@@ -14,6 +14,7 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/sqlsec/types"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 var badInputContextOnce sync.Once
@@ -24,7 +25,7 @@ func ProtectSQLOperation(ctx context.Context, query, driver string) error {
 		Driver: driver,
 	}
 
-	parent, _ := ctx.Value(listener.ContextKey{}).(dyngo.Operation)
+	parent, _ := orchestrion.CtxOrGLS(ctx).Value(listener.ContextKey{}).(dyngo.Operation)
 	if parent == nil { // No parent operation => we can't monitor the request
 		badInputContextOnce.Do(func() {
 			log.Debug("appsec: outgoing SQL operation monitoring ignored: could not find the handler " +

internal/orchestrion/context.go

@@ -0,0 +1,37 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package orchestrion
+
+import "context"
+
+// CtxGLS returns the context stored in the compilation-injected field of the runtime.g struct by orchestrion
+func CtxGLS() context.Context {
+	return getDDGLS().(context.Context)
+}
+
+// CtxOrGLS implements the logic to choose between a user given context.Context or the context.Context inserted in
+// the Goroutine Global Storage (GLS) by orchestrion. If Orchestrion is not enabled it's only returning the given context.
+func CtxOrGLS(ctx context.Context) context.Context {
+	if !Enabled() || (ctx != nil && ctx != context.Background()) {
+		return ctx
+	}
+
+	gls := CtxGLS()
+	if gls != nil {
+		return gls
+	}
+
+	ctx = context.Background()
+	setDDGLS(ctx)
+	return ctx
+}
+
+// CtxWithValue runs context.WithValue, adds the result to the GLS slot of orchestrion, and returns it.
+func CtxWithValue(ctx context.Context, key, val any) context.Context {
+	ctx = context.WithValue(CtxOrGLS(ctx), key, val)
+	setDDGLS(ctx)
+	return ctx
+}

internal/orchestrion/gls.go

@@ -0,0 +1,33 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package orchestrion
+
+import _ "unsafe" // for go:linkname
+
+var (
+	// getDDGLS returns the current value from the field inserted in runtime.g by orchestrion.
+	// Or nil if orchestrion is not enabled.
+	getDDGLS = func() any { return nil }
+	// setDDGLS sets the value in the field inserted in runtime.g by orchestrion.
+	// Or does nothing if orchestrion is not enabled.
+	setDDGLS = func(any) {}
+)
+
+//go:linkname _DdOrchestrionGlsGet __dd_orchestrion_gls_get
+var _DdOrchestrionGlsGet func() any
+
+//go:linkname _DdOrchestrionGlsSet __dd_orchestrion_gls_set
+var _DdOrchestrionGlsSet func(any)
+
+// Check at Go init time that the two function variable values created by the
+// orchestrion are present, and set the get/set variables to their
+// values.
+func init() {
+	if _DdOrchestrionGlsGet != nil && _DdOrchestrionGlsSet != nil {
+		getDDGLS = _DdOrchestrionGlsGet
+		setDDGLS = _DdOrchestrionGlsSet
+	}
+}

internal/orchestrion/orchestrion.go

@@ -0,0 +1,12 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package orchestrion
+
+// Enabled returns whether the current build was compiled with orchestrion or not.
+// This function is modified at compile time by orchestrion to be true if orchestrion is used.
+func Enabled() bool {
+	return false
+}

internal/trace_context.go

@@ -7,6 +7,8 @@ package internal
 
 import (
 	"context"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/orchestrion"
 )
 
 type executionTracedKey struct{}
@@ -23,25 +25,25 @@ type executionTracedKey struct{}
 // general, APM instrumentation should prefer creating tasks around the
 // operation rather than after the fact, if possible.
 func WithExecutionTraced(ctx context.Context) context.Context {
-	return context.WithValue(ctx, executionTracedKey{}, true)
+	return orchestrion.CtxWithValue(ctx, executionTracedKey{}, true)
 }
 
 // WithExecutionNotTraced marks that the context is *not* covered by an
 // execution trace task.  This is intended to prevent child spans (which inherit
 // information from ctx) from being considered covered by a task, when an
 // integration may create its own child span with its own execution trace task.
 func WithExecutionNotTraced(ctx context.Context) context.Context {
-	if ctx.Value(executionTracedKey{}) == nil {
+	if orchestrion.CtxOrGLS(ctx).Value(executionTracedKey{}) == nil {
 		// Fast path: if it wasn't marked before, we don't need to wrap
 		// the context
 		return ctx
 	}
-	return context.WithValue(ctx, executionTracedKey{}, false)
+	return orchestrion.CtxWithValue(ctx, executionTracedKey{}, false)
 }
 
 // IsExecutionTraced returns whether ctx is associated with an execution trace
 // task, as indicated via WithExecutionTraced
 func IsExecutionTraced(ctx context.Context) bool {
-	v := ctx.Value(executionTracedKey{})
+	v := orchestrion.CtxOrGLS(ctx).Value(executionTracedKey{})
 	return v != nil && v.(bool)
 }