-
Notifications
You must be signed in to change notification settings - Fork 19
/
trace.go
173 lines (149 loc) · 4.38 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
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
package cocaine12
import (
"context"
"fmt"
"math/rand"
"sync"
"time"
)
const (
TraceInfoValue = "trace.traceinfo"
TraceStartTimeValue = "trace.starttime"
)
var (
initTraceLogger sync.Once
traceLogger Logger
closeDummySpan CloseSpan = func() {}
)
func GetTraceInfo(ctx context.Context) *TraceInfo {
if val, ok := ctx.Value(TraceInfoValue).(TraceInfo); ok {
return &val
}
return nil
}
// CloseSpan closes attached span. It should be call after
// the rpc ends.
type CloseSpan func()
type TraceInfo struct {
Trace, Span, Parent uint64
logger Logger
}
func (traceInfo *TraceInfo) getLog() Logger {
if traceInfo.logger != nil {
return traceInfo.logger
}
initTraceLogger.Do(func() {
var err error
traceLogger, err = NewLogger(context.Background())
// there must be no error
if err != nil {
panic(fmt.Sprintf("unable to create trace logger: %v", err))
}
})
return traceLogger
}
type traced struct {
context.Context
traceInfo TraceInfo
startTime time.Time
}
func (t *traced) Value(key interface{}) interface{} {
switch key {
case TraceInfoValue:
return t.traceInfo
case TraceStartTimeValue:
return t.startTime
default:
return t.Context.Value(key)
}
}
// It might be used in client applications.
func BeginNewTraceContext(ctx context.Context) context.Context {
return BeginNewTraceContextWithLogger(ctx, nil)
}
func BeginNewTraceContextWithLogger(ctx context.Context, logger Logger) context.Context {
ts := uint64(rand.Int63())
return AttachTraceInfo(ctx, TraceInfo{
Trace: ts,
Span: ts,
Parent: 0,
logger: logger,
})
}
// AttachTraceInfo binds given TraceInfo to the context.
// If ctx is nil, then TraceInfo will be attached to context.Background()
func AttachTraceInfo(ctx context.Context, traceInfo TraceInfo) context.Context {
if ctx == nil {
ctx = context.Background()
}
return &traced{
Context: ctx,
traceInfo: traceInfo,
startTime: time.Now(),
}
}
// CleanTraceInfo might be used to clear context instance from trace info
// to disable tracing in some RPC calls to get rid of overhead
func CleanTraceInfo(ctx context.Context) context.Context {
return context.WithValue(ctx, TraceInfoValue, nil)
}
// NewSpan starts new span and returns a context with attached TraceInfo and Done.
// If ctx is nil or has no TraceInfo new span won't start to support sampling,
// so it's user responsibility to make sure that the context has TraceInfo.
// Anyway it safe to call CloseSpan function even in this case, it actually does nothing.
func NewSpan(ctx context.Context, rpcNameFormat string, args ...interface{}) (context.Context, CloseSpan) {
if ctx == nil {
// I'm not sure it is a valid action.
// According to the rule "no trace info, no new span"
// to support sampling, nil Context has no TraceInfo, so
// it cannot start new Span.
return context.Background(), closeDummySpan
}
traceInfo := GetTraceInfo(ctx)
if traceInfo == nil {
// given context has no TraceInfo
// so we can't start new trace to support sampling.
// closeDummySpan does nohing
return ctx, closeDummySpan
}
var rpcName string
if len(args) > 0 {
rpcName = fmt.Sprintf(rpcNameFormat, args...)
} else {
rpcName = rpcNameFormat
}
// startTime is not used only to log the start of an RPC
// It's stored in Context to calculate the RPC call duration.
// A user can get it via Context.Value(TraceStartTimeValue)
startTime := time.Now()
// Tracing magic:
// * the previous span becomes our parent
// * new span is set as random number
// * trace still stays the same
traceInfo.Parent = traceInfo.Span
traceInfo.Span = uint64(rand.Int63())
traceInfo.getLog().WithFields(Fields{
"trace_id": fmt.Sprintf("%x", traceInfo.Trace),
"span_id": fmt.Sprintf("%x", traceInfo.Span),
"parent_id": fmt.Sprintf("%x", traceInfo.Parent),
"real_timestamp": startTime.UnixNano() / 1000,
"rpc_name": rpcName,
}).Infof("start")
ctx = &traced{
Context: ctx,
traceInfo: *traceInfo,
startTime: startTime,
}
return ctx, func() {
now := time.Now()
duration := now.Sub(startTime)
traceInfo.getLog().WithFields(Fields{
"trace_id": fmt.Sprintf("%x", traceInfo.Trace),
"span_id": fmt.Sprintf("%x", traceInfo.Span),
"parent_id": fmt.Sprintf("%x", traceInfo.Parent),
"real_timestamp": now.UnixNano() / 1000,
"duration": duration.Nanoseconds() / 1000,
"rpc_name": rpcName,
}).Infof("finish")
}
}