forked from goadesign/goa
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tracer.go
134 lines (117 loc) · 3.82 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
package middleware
import (
rd "math/rand"
"net/http"
"context"
"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
// tracedDoer is a goa client Doer that inserts the tracing headers for
// each request it makes.
tracedDoer struct {
client.Doer
}
)
// Tracer returns a middleware that initializes the trace information in the
// context. The information can be retrieved using any of the ContextXXX
// functions.
//
// sampleRate must be a value between 0 and 100. It represents the percentage of
// requests that should be traced.
//
// 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.
//
// If the incoming request has a TraceIDHeader header then the sample rate is
// disregarded and the tracing is enabled.
func Tracer(sampleRate int, spanIDFunc, traceIDFunc IDFunc) goa.Middleware {
if sampleRate < 0 || sampleRate > 100 {
panic("tracing: sample rate must be between 0 and 100")
}
return func(h goa.Handler) goa.Handler {
return func(ctx context.Context, rw http.ResponseWriter, req *http.Request) error {
// Compute trace info.
var (
traceID = req.Header.Get(TraceIDHeader)
parentID = req.Header.Get(ParentSpanIDHeader)
spanID = spanIDFunc()
)
if traceID == "" {
// Avoid computing a random value if unnecessary.
if sampleRate == 0 || rd.Intn(100) > sampleRate {
return h(ctx, rw, req)
}
traceID = traceIDFunc()
}
// Setup context.
ctx = WithTrace(ctx, traceID, spanID, parentID)
// Call next handler.
return h(ctx, rw, req)
}
}
}
// 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)
}