-
Notifications
You must be signed in to change notification settings - Fork 224
/
zipkin.go
116 lines (97 loc) · 2.65 KB
/
zipkin.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
package handlers
import (
"net/http"
"github.com/openzipkin/zipkin-go/propagation/b3"
"github.com/uber-go/zap"
"github.com/urfave/negroni/v3"
"code.cloudfoundry.org/gorouter/logger"
)
// Zipkin is a handler that sets Zipkin headers on requests
type Zipkin struct {
zipkinEnabled bool
logger logger.Logger
}
var _ negroni.Handler = new(Zipkin)
// NewZipkin creates a new handler that sets Zipkin headers on requests
func NewZipkin(enabled bool, logger logger.Logger) *Zipkin {
return &Zipkin{
zipkinEnabled: enabled,
logger: logger,
}
}
func (z *Zipkin) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
defer next(rw, r)
logger := LoggerWithTraceInfo(z.logger, r)
if !z.zipkinEnabled {
return
}
requestInfo, err := ContextRequestInfo(r)
if err != nil {
logger.Error("failed-to-get-request-info", zap.Error(err))
return
}
existingContext := r.Header.Get(b3.Context)
if existingContext != "" {
logger.Debug("b3-header-exists",
zap.String("b3", existingContext),
)
sc, err := b3.ParseSingleHeader(existingContext)
if err != nil {
logger.Error("failed-to-parse-single-header", zap.Error(err))
} else {
err = requestInfo.SetTraceInfo(sc.TraceID.String(), sc.ID.String())
if err != nil {
logger.Error("failed-to-set-trace-info", zap.Error(err))
} else {
return
}
}
}
existingTraceID := r.Header.Get(b3.TraceID)
existingSpanID := r.Header.Get(b3.SpanID)
if existingTraceID != "" && existingSpanID != "" {
sc, err := b3.ParseHeaders(
existingTraceID,
existingSpanID,
r.Header.Get(b3.ParentSpanID),
r.Header.Get(b3.Sampled),
r.Header.Get(b3.Flags),
)
if err != nil {
logger.Info("failed-to-parse-b3-trace-id", zap.Error(err))
return
}
r.Header.Set(b3.Context, b3.BuildSingleHeader(*sc))
logger.Debug("b3-trace-id-span-id-header-exists",
zap.String("trace-id", existingTraceID),
zap.String("span-id", existingSpanID),
)
err = requestInfo.SetTraceInfo(sc.TraceID.String(), sc.ID.String())
if err != nil {
logger.Error("failed-to-set-trace-info", zap.Error(err))
} else {
return
}
}
traceInfo, err := requestInfo.ProvideTraceInfo()
if err != nil {
logger.Error("failed-to-get-trace-info", zap.Error(err))
return
}
r.Header.Set(b3.TraceID, traceInfo.TraceID)
r.Header.Set(b3.SpanID, traceInfo.SpanID)
r.Header.Set(b3.Context, traceInfo.TraceID+"-"+traceInfo.SpanID)
}
// HeadersToLog specifies the headers which should be logged if Zipkin headers
// are enabled
func (z *Zipkin) HeadersToLog() []string {
if !z.zipkinEnabled {
return []string{}
}
return []string{
b3.TraceID,
b3.SpanID,
b3.ParentSpanID,
b3.Context,
}
}