forked from DataDog/dd-trace-go
-
Notifications
You must be signed in to change notification settings - Fork 1
/
textmap.go
197 lines (173 loc) · 5.73 KB
/
textmap.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
package tracer
import (
"net/http"
"strconv"
"strings"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
)
// HTTPHeadersCarrier wraps an http.Header as a TextMapWriter and TextMapReader, allowing
// it to be used using the provided Propagator implementation.
type HTTPHeadersCarrier http.Header
var _ TextMapWriter = (*HTTPHeadersCarrier)(nil)
var _ TextMapReader = (*HTTPHeadersCarrier)(nil)
// Set implements TextMapWriter.
func (c HTTPHeadersCarrier) Set(key, val string) {
http.Header(c).Set(key, val)
}
// ForeachKey implements TextMapReader.
func (c HTTPHeadersCarrier) ForeachKey(handler func(key, val string) error) error {
for k, vals := range c {
for _, v := range vals {
if err := handler(k, v); err != nil {
return err
}
}
}
return nil
}
// TextMapCarrier allows the use of a regular map[string]string as both TextMapWriter
// and TextMapReader, making it compatible with the provided Propagator.
type TextMapCarrier map[string]string
var _ TextMapWriter = (*TextMapCarrier)(nil)
var _ TextMapReader = (*TextMapCarrier)(nil)
// Set implements TextMapWriter.
func (c TextMapCarrier) Set(key, val string) {
c[key] = val
}
// ForeachKey conforms to the TextMapReader interface.
func (c TextMapCarrier) ForeachKey(handler func(key, val string) error) error {
for k, v := range c {
if err := handler(k, v); err != nil {
return err
}
}
return nil
}
const (
// DefaultBaggageHeaderPrefix specifies the prefix that will be used in
// HTTP headers or text maps to prefix baggage keys.
DefaultBaggageHeaderPrefix = "ot-baggage-"
// DefaultTraceIDHeader specifies the key that will be used in HTTP headers
// or text maps to store the trace ID.
DefaultTraceIDHeader = "x-datadog-trace-id"
// DefaultParentIDHeader specifies the key that will be used in HTTP headers
// or text maps to store the parent ID.
DefaultParentIDHeader = "x-datadog-parent-id"
// DefaultPriorityHeader specifies the key that will be used in HTTP headers
// or text maps to store the sampling priority value.
DefaultPriorityHeader = "x-datadog-sampling-priority"
)
// PropagatorConfig defines the configuration for initializing a propagator.
type PropagatorConfig struct {
// BaggagePrefix specifies the prefix that will be used to store baggage
// items in a map. It defaults to DefaultBaggageHeaderPrefix.
BaggagePrefix string
// TraceHeader specifies the map key that will be used to store the trace ID.
// It defaults to DefaultTraceIDHeader.
TraceHeader string
// ParentHeader specifies the map key that will be used to store the parent ID.
// It defaults to DefaultParentIDHeader.
ParentHeader string
// PriorityHeader specifies the map key that will be used to store the sampling priority.
// It deafults to DefaultPriorityHeader.
PriorityHeader string
}
// NewPropagator returns a new propagator which uses TextMap to inject
// and extract values. It propagates trace and span IDs and baggage.
// To use the defaults, nil may be provided in place of the config.
func NewPropagator(cfg *PropagatorConfig) Propagator {
if cfg == nil {
cfg = new(PropagatorConfig)
}
if cfg.BaggagePrefix == "" {
cfg.BaggagePrefix = DefaultBaggageHeaderPrefix
}
if cfg.TraceHeader == "" {
cfg.TraceHeader = DefaultTraceIDHeader
}
if cfg.ParentHeader == "" {
cfg.ParentHeader = DefaultParentIDHeader
}
if cfg.PriorityHeader == "" {
cfg.PriorityHeader = DefaultPriorityHeader
}
return &propagator{cfg}
}
// propagator implements a propagator which uses TextMap internally.
// It propagates the trace and span IDs, as well as the baggage from the
// context.
type propagator struct{ cfg *PropagatorConfig }
// Inject defines the Propagator to propagate SpanContext data
// out of the current process. The implementation propagates the
// TraceID and the current active SpanID, as well as the Span baggage.
func (p *propagator) Inject(spanCtx ddtrace.SpanContext, carrier interface{}) error {
switch v := carrier.(type) {
case TextMapWriter:
return p.injectTextMap(spanCtx, v)
default:
return ErrInvalidCarrier
}
}
func (p *propagator) injectTextMap(spanCtx ddtrace.SpanContext, writer TextMapWriter) error {
ctx, ok := spanCtx.(*spanContext)
if !ok || ctx.traceID == 0 || ctx.spanID == 0 {
return ErrInvalidSpanContext
}
// propagate the TraceID and the current active SpanID
writer.Set(p.cfg.TraceHeader, strconv.FormatUint(ctx.traceID, 10))
writer.Set(p.cfg.ParentHeader, strconv.FormatUint(ctx.spanID, 10))
if ctx.hasSamplingPriority() {
writer.Set(p.cfg.PriorityHeader, strconv.Itoa(ctx.samplingPriority()))
}
// propagate OpenTracing baggage
for k, v := range ctx.baggage {
writer.Set(p.cfg.BaggagePrefix+k, v)
}
return nil
}
// Extract implements Propagator.
func (p *propagator) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
switch v := carrier.(type) {
case TextMapReader:
return p.extractTextMap(v)
default:
return nil, ErrInvalidCarrier
}
}
func (p *propagator) extractTextMap(reader TextMapReader) (ddtrace.SpanContext, error) {
var ctx spanContext
err := reader.ForeachKey(func(k, v string) error {
var err error
key := strings.ToLower(k)
switch key {
case p.cfg.TraceHeader:
ctx.traceID, err = parseUint64(v)
if err != nil {
return ErrSpanContextCorrupted
}
case p.cfg.ParentHeader:
ctx.spanID, err = parseUint64(v)
if err != nil {
return ErrSpanContextCorrupted
}
case p.cfg.PriorityHeader:
ctx.priority, err = strconv.Atoi(v)
if err != nil {
return ErrSpanContextCorrupted
}
ctx.hasPriority = true
default:
if strings.HasPrefix(key, p.cfg.BaggagePrefix) {
ctx.setBaggageItem(strings.TrimPrefix(key, p.cfg.BaggagePrefix), v)
}
}
return nil
})
if err != nil {
return nil, err
}
if ctx.traceID == 0 || ctx.spanID == 0 {
return nil, ErrSpanContextNotFound
}
return &ctx, nil
}