-
Notifications
You must be signed in to change notification settings - Fork 0
/
trace.go
123 lines (104 loc) · 2.73 KB
/
trace.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
package trace
import (
"log"
"net/http"
"github.com/opentracing/opentracing-go"
"github.com/opentracing/opentracing-go/ext"
"github.com/rs/xid"
"github.com/opengoats/goat/http/response"
"github.com/opengoats/goat/logger"
)
const (
// RequestIDHeaderKey 补充的RequestID Header Key
RequestIDHeaderKey = "X-Request-Id"
// DefaultComponentName 默认组件名称
DefaultComponentName = "goat.server.trace"
)
// New 实例
func New(tr opentracing.Tracer, peer PeerInfo) *Tracer {
opentracing.SetGlobalTracer(tr)
return &Tracer{
tr: tr,
debug: true,
peer: peer,
}
}
// PeerInfo 端的信息
type PeerInfo struct {
Service string
Address string
Port uint16
Hostname string
}
// Tracer 链路追踪器
type Tracer struct {
tr opentracing.Tracer
log logger.Logger
debug bool
peer PeerInfo
}
// Debug 是否开启debug
func (h *Tracer) Debug(on bool) {
h.debug = on
}
// SetLogger 设置Logger
func (h *Tracer) SetLogger(l logger.Logger) {
h.log = l
}
// Handler todo
func (h *Tracer) Handler(next http.Handler) http.Handler {
return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
var span opentracing.Span
// Attempt to join a trace by getting trace context from the headers.
carrier := opentracing.HTTPHeadersCarrier(req.Header)
ctx, err := h.tr.Extract(opentracing.HTTPHeaders, carrier)
if err != nil {
span = opentracing.StartSpan(req.URL.String(), ext.RPCServerOption(nil))
} else {
span = opentracing.StartSpan(req.URL.String(), ext.RPCServerOption(ctx))
}
// 添加Request ID
requestID := req.Header.Get(RequestIDHeaderKey)
if requestID == "" {
requestID = xid.New().String()
req.Header.Set(RequestIDHeaderKey, requestID)
}
span.SetTag("request.id", requestID)
if err := span.Tracer().Inject(span.Context(), opentracing.HTTPHeaders, carrier); err != nil {
ext.Error.Set(span, true)
} else {
req = req.WithContext(opentracing.ContextWithSpan(req.Context(), span))
}
// 补充基本信息
if span != nil {
ext.HTTPMethod.Set(span, req.Method)
ext.Component.Set(span, DefaultComponentName)
ext.PeerService.Set(span, h.peer.Service)
ext.PeerAddress.Set(span, h.peer.Address)
ext.PeerHostname.Set(span, h.peer.Hostname)
ext.PeerPort.Set(span, h.peer.Port)
}
// 结束时补充status code
defer func() {
if span != nil {
res := rw.(response.Response)
res.Header().Set(RequestIDHeaderKey, requestID)
ext.HTTPStatusCode.Set(span, uint16(res.Status()))
span.Finish()
}
h.logf("%s", span)
}()
next.ServeHTTP(rw, req)
})
}
func (h *Tracer) logf(format string, args ...interface{}) {
if !h.debug {
return
}
if h.log != nil {
h.log.Debugf(format, args...)
return
}
log.Printf(format+"\n", args...)
return
}