// 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 2016 Datadog, Inc.

package tracer

import (

"context"

"github.com/DataDog/dd-trace-go/v2/instrumentation/options"

"github.com/DataDog/dd-trace-go/v2/internal"

illmobs "github.com/DataDog/dd-trace-go/v2/internal/llmobs"

"github.com/DataDog/dd-trace-go/v2/internal/log"

"github.com/DataDog/dd-trace-go/v2/internal/orchestrion"

)

// ContextWithSpan returns a copy of the given context which includes the span s.

// If ctx is nil, a new background context is created to avoid panicking.

func ContextWithSpan(ctx context.Context, s *Span) context.Context {

if ctx == nil {

log.Warn("ContextWithSpan: received nil context, falling back to context.Background()")

ctx = context.Background()

}

newCtx := orchestrion.CtxWithValue(ctx, internal.ActiveSpanKey, s)

return contextWithPropagatedLLMSpan(newCtx, s)

}

func contextWithPropagatedLLMSpan(ctx context.Context, s *Span) context.Context {

if s == nil {

return ctx

}

// if there is a propagated llm span already just skip

if _, ok := illmobs.PropagatedLLMSpanFromContext(ctx); ok {

return ctx

}

propagatedLLMObs := propagatedLLMSpanFromTags(s)

if propagatedLLMObs.SpanID == "" || propagatedLLMObs.TraceID == "" {

return ctx

}

return illmobs.ContextWithPropagatedLLMSpan(ctx, propagatedLLMObs)

}

// propagatedLLMSpanFromTags extracts LLMObs propagation information from the trace propagating tags.

// This is used during distributed tracing to set the correct parent span for the current span.

func propagatedLLMSpanFromTags(s *Span) *illmobs.PropagatedLLMSpan {

propagatedLLMObs := &illmobs.PropagatedLLMSpan{}

if s.context == nil || s.context.trace == nil {

return propagatedLLMObs

}

if parentID := s.context.trace.propagatingTag(keyPropagatedLLMObsParentID); parentID != "" {

propagatedLLMObs.SpanID = parentID

}

if mlApp := s.context.trace.propagatingTag(keyPropagatedLLMObsMLAPP); mlApp != "" {

propagatedLLMObs.MLApp = mlApp

}

if trID := s.context.trace.propagatingTag(keyPropagatedLLMObsTraceID); trID != "" {

propagatedLLMObs.TraceID = trID

}

return propagatedLLMObs

}

// 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 {

return nil, false

}

v := orchestrion.WrapContext(ctx).Value(internal.ActiveSpanKey)

if s, ok := v.(*Span); ok {

// We may have a nil *Span wrapped in an interface in the GLS context stack,

// in which case we need to act as if there was nothing (otherwise we'll

// forcefully un-do a [ChildOf] option if one was passed).

if s == nil {

return nil, false

}

return s, true

}

return nil, false

}

// StartSpanFromContext returns a new span with the given operation name and options. If a span

// is found in the context, it will be used as the parent of the resulting span. If the ChildOf

// option is passed, it will only be used as the parent if there is no span found in `ctx`.

// +checklocksignore — Initialization time, span just created by StartSpan, not yet shared.

func StartSpanFromContext(ctx context.Context, operationName string, opts ...StartSpanOption) (*Span, context.Context) {

// copy opts in case the caller reuses the slice in parallel

// we will add at least 1, at most 2 items

optsLocal := options.Expand(opts, 0, 2)

if ctx == nil {

// default to context.Background() to avoid panics on Go >= 1.15

ctx = context.Background()

} else if s, ok := SpanFromContext(ctx); ok {

optsLocal = append(optsLocal, ChildOf(s.Context()))

}

optsLocal = append(optsLocal, withContext(ctx))

s := StartSpan(operationName, optsLocal...)

if s != nil && s.pprofCtxActive != nil {

ctx = s.pprofCtxActive

}

return s, ContextWithSpan(ctx, s)

}