forked from honeycombio/beeline-go
/
context.go
74 lines (66 loc) · 2.46 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
package trace
import (
"context"
"errors"
)
const (
honeySpanContextKey = "honeycombSpanContextKey"
honeyTraceContextKey = "honeycombTraceContextKey"
)
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 {
return context.WithValue(ctx, honeyTraceContextKey, trace)
}
// 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
}