-
Notifications
You must be signed in to change notification settings - Fork 1
/
trace.go
43 lines (36 loc) · 1 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
// Package trace 基于 context 封装一个 trace ID 功能
package trace
import (
"context"
objectid "github.com/Andrew-M-C/go.objectid"
)
type traceIDKey struct{}
// GetTraceID 读取 trace ID
func GetTraceID(ctx context.Context) string {
if v := ctx.Value(traceIDKey{}); v != nil {
id, _ := v.(string)
return id
}
return ""
}
// EnsureTraceID 如果没有 trace id 的话, 那就填充一个并返回
func EnsureTraceID(ctx context.Context) context.Context {
if id := GetTraceID(ctx); id != "" {
return ctx
}
return SetTraceID(ctx)
}
// SetTraceID 设置一个 trace ID, 如果之前已经设置过, 则会覆盖
func SetTraceID(ctx context.Context, traceID ...string) context.Context {
id := ""
if len(traceID) > 0 && traceID[0] != "" {
id = traceID[0]
} else {
id = objectid.New16().String()
}
return context.WithValue(ctx, traceIDKey{}, id)
}
// WithTraceID 等价于 SetTraceID
func WithTraceID(ctx context.Context, traceID ...string) context.Context {
return SetTraceID(ctx, traceID...)
}