github.com/replit/upm@v0.0.0-20240423230255-9ce4fc3ea24c/internal/trace/trace.go (about)

     1  package trace
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
     8  	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
     9  )
    10  
    11  var (
    12  	globalDDTraceID string
    13  	globalDDSpanID  string
    14  )
    15  
    16  func MaybeTrace(serviceVersion string) func() {
    17  	if os.Getenv("UPM_TRACE") != "1" {
    18  		return nil
    19  	}
    20  
    21  	replid := os.Getenv("REPL_ID")
    22  	if replid == "" {
    23  		return nil
    24  	}
    25  
    26  	globalDDTraceID = os.Getenv("DD_TRACE_ID")
    27  	globalDDSpanID = os.Getenv("DD_SPAN_ID")
    28  	os.Unsetenv("DD_TRACE_ID")
    29  	os.Unsetenv("DD_SPAN_ID")
    30  
    31  	logger, err := NewDatadogLogger()
    32  	if err != nil {
    33  		return nil
    34  	}
    35  
    36  	rules := []tracer.SamplingRule{
    37  		// send 100.00% of traces
    38  		tracer.ServiceRule("upm", 1.0),
    39  	}
    40  
    41  	tracer.Start(
    42  		tracer.WithService("upm"),
    43  		tracer.WithGlobalTag("replid", replid),
    44  		tracer.WithServiceVersion(serviceVersion),
    45  		tracer.WithLogger(logger),
    46  		tracer.WithSamplingRules(rules),
    47  	)
    48  	return func() {
    49  		tracer.Stop()
    50  		logger.Close()
    51  	}
    52  }
    53  
    54  func StartSpanFromExistingContext(name string) (ddtrace.Span, context.Context) {
    55  	ctx := context.Background()
    56  	parentContext, _ := GetParentContext()
    57  	if parentContext == nil {
    58  		return tracer.StartSpanFromContext(ctx, name)
    59  	}
    60  	return tracer.StartSpanFromContext(ctx, name, WithParentContext(parentContext))
    61  }
    62  
    63  func GetParentContext() (*SpanContext, error) {
    64  	traceID := globalDDTraceID
    65  	spanID := globalDDSpanID
    66  	if traceID == "" || spanID == "" {
    67  		return nil, nil
    68  	}
    69  	parentContext := &SpanContext{}
    70  	err := parentContext.ParseTraceID(traceID)
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  	err = parentContext.ParseSpanID(spanID)
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  	return parentContext, nil
    79  }
    80  
    81  func WithParentContext(c *SpanContext) ddtrace.StartSpanOption {
    82  	return func(cfg *ddtrace.StartSpanConfig) {
    83  		cfg.Parent = c
    84  	}
    85  }