-
Notifications
You must be signed in to change notification settings - Fork 48
/
context.go
86 lines (78 loc) · 2.95 KB
/
context.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
86
package trace
import (
"context"
"errors"
"runtime/pprof"
)
const (
honeySpanContextKey = "honeycombSpanContextKey"
honeyTraceContextKey = "honeycombTraceContextKey"
profileIDLabelName = "profile_id"
)
var (
ErrTraceNotFoundInContext = errors.New("beeline trace not found in source context")
)
// GetTraceFromContext retrieves a trace from the passed in context or returns
// nil if no trace exists.
func GetTraceFromContext(ctx context.Context) *Trace {
if ctx != nil {
if val := ctx.Value(honeyTraceContextKey); val != nil {
if trace, ok := val.(*Trace); ok {
return trace
}
}
}
return nil
}
// PutTraceInContext takes an existing context and a trace and pushes the trace
// into the context. It will replace any traces that already exist in the
// context. Traces put in context are retrieved using GetTraceFromContext.
func PutTraceInContext(ctx context.Context, trace *Trace) context.Context {
// Copy the old context object to preserve its pprof labels to restore later.
oldCtx := ctx
ctx = context.WithValue(ctx, honeyTraceContextKey, trace)
if GlobalConfig.PprofTagging && trace != nil && trace.GetRootSpan() != nil {
// This returns a pointer type, so it's safe to directly manipulate fields.
rootSpan := trace.GetRootSpan()
rootSpan.oldCtx = &oldCtx
ctx = pprof.WithLabels(ctx, pprof.Labels(profileIDLabelName, rootSpan.GetSpanID()))
pprof.SetGoroutineLabels(ctx)
}
return ctx
}
// GetSpanFromContext identifies the currently active span via the span context
// key. It returns that span, and access to the trace is available via the span
// or from the context directly. It will return nil if there is no span
// available.
func GetSpanFromContext(ctx context.Context) *Span {
if ctx != nil {
if val := ctx.Value(honeySpanContextKey); val != nil {
if span, ok := val.(*Span); ok {
return span
}
}
}
return nil
}
// PutSpanInContext takes an existing context and a span and pushes the span
// into the context. It will replace any spans that already exist in the
// context. Spans put in context are retrieved using GetSpanFromContext.
func PutSpanInContext(ctx context.Context, span *Span) context.Context {
return context.WithValue(ctx, honeySpanContextKey, span)
}
// CopyContext takes a context that has a beeline trace and one that doesn't. It
// copies all the bits necessary to continue the trace from one to the other.
// This is useful if you need to break context to launch a goroutine that
// shouldn't be cancelled by the parent's cancellation context. It returns the
// newly populated context. If it can't find a trace in the source context, it
// returns the unchanged dest context with an error.
func CopyContext(dest context.Context, src context.Context) (context.Context, error) {
trace := GetTraceFromContext(src)
span := GetSpanFromContext(src)
if trace == nil || span == nil {
return dest, ErrTraceNotFoundInContext
}
dest = PutTraceInContext(dest, trace)
dest = PutSpanInContext(dest, span)
return dest, nil
}