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