forked from goadesign/goa
/
tracer.go
220 lines (197 loc) · 6.31 KB
/
tracer.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
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
package middleware
import (
"context"
"net/http"
"github.com/goadesign/goa"
"github.com/goadesign/goa/client"
)
var (
// TraceIDHeader is the name of the HTTP request header containing the
// current TraceID if any.
TraceIDHeader = "TraceID"
// ParentSpanIDHeader is the name of the HTTP request header containing
// the parent span ID if any.
ParentSpanIDHeader = "ParentSpanID"
)
type (
// IDFunc is a function that produces span and trace IDs for cosumption by
// tracing systems such as Zipkin or AWS X-Ray.
IDFunc func() string
// TracerOption is a constructor option that makes it possible to customize
// the middleware.
TracerOption func(*tracerOptions) *tracerOptions
// tracerOptions is the struct storing all the options.
tracerOptions struct {
traceIDFunc IDFunc
spanIDFunc IDFunc
samplingPercent int
maxSamplingRate int
sampleSize int
}
// tracedDoer is a goa client Doer that inserts the tracing headers for
// each request it makes.
tracedDoer struct {
client.Doer
}
)
// TraceIDFunc is a constructor option that overrides the function used to
// compute trace IDs.
func TraceIDFunc(f IDFunc) TracerOption {
return func(o *tracerOptions) *tracerOptions {
if f == nil {
panic("trace ID function cannot be nil")
}
o.traceIDFunc = f
return o
}
}
// SpanIDFunc is a constructor option that overrides the function used to
// compute span IDs.
func SpanIDFunc(f IDFunc) TracerOption {
return func(o *tracerOptions) *tracerOptions {
if f == nil {
panic("span ID function cannot be nil")
}
o.spanIDFunc = f
return o
}
}
// SamplingPercent sets the tracing sampling rate as a percentage value.
// It panics if p is less than 0 or more than 100.
// SamplingPercent and MaxSamplingRate are mutually exclusive.
func SamplingPercent(p int) TracerOption {
if p < 0 || p > 100 {
panic("sampling rate must be between 0 and 100")
}
return func(o *tracerOptions) *tracerOptions {
o.samplingPercent = p
return o
}
}
// MaxSamplingRate sets a target sampling rate in requests per second. Setting a
// max sampling rate causes the middleware to adjust the sampling percent
// dynamically.
// SamplingPercent and MaxSamplingRate are mutually exclusive.
func MaxSamplingRate(r int) TracerOption {
if r <= 0 {
panic("max sampling rate must be greater than 0")
}
return func(o *tracerOptions) *tracerOptions {
o.maxSamplingRate = r
return o
}
}
// SampleSize sets the number of requests between two adjustments of the sampling
// rate when MaxSamplingRate is set. Defaults to 1,000.
func SampleSize(s int) TracerOption {
if s <= 0 {
panic("sample size must be greater than 0")
}
return func(o *tracerOptions) *tracerOptions {
o.sampleSize = s
return o
}
}
// NewTracer returns a trace middleware that initializes the trace information
// in the request context. The information can be retrieved using any of the
// ContextXXX functions.
//
// samplingPercent must be a value between 0 and 100. It represents the percentage
// of requests that should be traced. If the incoming request has a Trace ID
// header then the sampling rate is disregarded and the tracing is enabled.
//
// spanIDFunc and traceIDFunc are the functions used to create Span and Trace
// IDs respectively. This is configurable so that the created IDs are compatible
// with the various backend tracing systems. The xray package provides
// implementations that produce AWS X-Ray compatible IDs.
func NewTracer(opts ...TracerOption) goa.Middleware {
o := &tracerOptions{
traceIDFunc: shortID,
spanIDFunc: shortID,
samplingPercent: 100,
sampleSize: 1000, // only applies if maxSamplingRate is set
}
for _, opt := range opts {
o = opt(o)
}
var sampler Sampler
if o.maxSamplingRate > 0 {
sampler = NewAdaptiveSampler(o.maxSamplingRate, o.sampleSize)
} else {
sampler = NewFixedSampler(o.samplingPercent)
}
return func(h goa.Handler) goa.Handler {
return func(ctx context.Context, rw http.ResponseWriter, req *http.Request) error {
// insert a new trace ID only if not already being traced.
traceID := req.Header.Get(TraceIDHeader)
if traceID == "" {
// insert tracing only within sample.
if sampler.Sample() {
traceID = o.traceIDFunc()
} else {
return h(ctx, rw, req)
}
}
// insert IDs into context to enable tracing.
spanID := o.spanIDFunc()
parentID := req.Header.Get(ParentSpanIDHeader)
ctx = WithTrace(ctx, traceID, spanID, parentID)
return h(ctx, rw, req)
}
}
}
// Tracer is deprecated in favor of NewTracer.
func Tracer(sampleRate int, spanIDFunc, traceIDFunc IDFunc) goa.Middleware {
return NewTracer(SamplingPercent(sampleRate), SpanIDFunc(spanIDFunc), TraceIDFunc(traceIDFunc))
}
// TraceDoer wraps a goa client Doer and sets the trace headers so that the
// downstream service may properly retrieve the parent span ID and trace ID.
func TraceDoer(doer client.Doer) client.Doer {
return &tracedDoer{doer}
}
// ContextTraceID returns the trace ID extracted from the given context if any,
// the empty string otherwise.
func ContextTraceID(ctx context.Context) string {
if t := ctx.Value(traceKey); t != nil {
return t.(string)
}
return ""
}
// ContextSpanID returns the span ID extracted from the given context if any,
// the empty string otherwise.
func ContextSpanID(ctx context.Context) string {
if s := ctx.Value(spanKey); s != nil {
return s.(string)
}
return ""
}
// ContextParentSpanID returns the parent span ID extracted from the given
// context if any, the empty string otherwise.
func ContextParentSpanID(ctx context.Context) string {
if p := ctx.Value(parentSpanKey); p != nil {
return p.(string)
}
return ""
}
// WithTrace returns a context containing the given trace, span and parent span
// IDs.
func WithTrace(ctx context.Context, traceID, spanID, parentID string) context.Context {
if parentID != "" {
ctx = context.WithValue(ctx, parentSpanKey, parentID)
}
ctx = context.WithValue(ctx, traceKey, traceID)
ctx = context.WithValue(ctx, spanKey, spanID)
return ctx
}
// Do adds the tracing headers to the requests before making it.
func (d *tracedDoer) Do(ctx context.Context, req *http.Request) (*http.Response, error) {
var (
traceID = ContextTraceID(ctx)
spanID = ContextSpanID(ctx)
)
if traceID != "" {
req.Header.Set(TraceIDHeader, traceID)
req.Header.Set(ParentSpanIDHeader, spanID)
}
return d.Doer.Do(ctx, req)
}