-
Notifications
You must be signed in to change notification settings - Fork 0
/
tracelog.go
115 lines (100 loc) · 2.54 KB
/
tracelog.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
// Package tracelog 提供带 tracing 功能的日志功能以及一些工具
package tracelog
import (
"context"
"github.com/Andrew-M-C/go.util/ids/snowflake"
"golang.org/x/exp/slices"
"trpc.group/trpc-go/trpc-go/codec"
"trpc.group/trpc-go/trpc-go/log"
)
// 保存 ctx 中的 trace ID 字段
type traceIDKey struct{}
type traceIDStackyValue []string
func (t traceIDStackyValue) id() string {
if len(t) == 0 {
return ""
}
return t[len(t)-1]
}
// WithTraceID 更新 trace ID
func WithTraceID(ctx context.Context, traceID string) context.Context {
if traceID == "" {
return ctx
}
stack := traceIDStack(ctx)
if stack.id() == traceID {
return ctx
}
stack = slices.Clone(stack)
stack = append(stack, traceID)
return cloneCtxAndGenerateLog(ctx, stack)
}
// WithTraceIDStack 完全替换整个 trace ID 栈
func WithTraceIDStack(ctx context.Context, traceIDStack []string) context.Context {
if len(traceIDStack) == 0 {
return ctx
}
stack := traceIDStackyValue(slices.Clone(traceIDStack))
return cloneCtxAndGenerateLog(ctx, stack)
}
func cloneCtxAndGenerateLog(ctx context.Context, stack traceIDStackyValue) context.Context {
ctx, _ = codec.WithCloneMessage(ctx)
fields := []log.Field{{
Key: "trace_id",
Value: stack.id(),
}}
if le := len(stack); le > 1 {
fields = append(fields, log.Field{
Key: "trace_id_stack",
Value: stack[:le-1],
})
}
// TODO: dyeing
// if msg.Dyeing() {
// fields = append(fields, log.Field{
// Key: "dyeing",
// Value: true,
// })
// }
logger := log.GetDefaultLogger().With(fields...)
ctx, msg := codec.EnsureMessage(ctx)
msg.WithLogger(logger)
return context.WithValue(ctx, traceIDKey{}, stack)
}
func traceIDStack(ctx context.Context) traceIDStackyValue {
v := ctx.Value(traceIDKey{})
if v == nil {
return traceIDStackyValue{}
}
st, _ := v.(traceIDStackyValue)
return st
}
// TraceID 从 context 中读取 trace ID
func TraceID(ctx context.Context) string {
v := ctx.Value(traceIDKey{})
if v == nil {
return ""
}
s, _ := v.(traceIDStackyValue)
return s.id()
}
// TraceIDStack 从 context 中读取历史 trace ID 栈
func TraceIDStack(ctx context.Context) []string {
v := ctx.Value(traceIDKey{})
if v == nil {
return nil
}
s, _ := v.(traceIDStackyValue)
return slices.Clone(s)
}
// EnsureTraceID 确保 context 中有一个 trace ID, 协程不安全
func EnsureTraceID(ctx context.Context) context.Context {
traceID := TraceID(ctx)
if traceID == "" {
return WithTraceID(ctx, generateTraceID())
}
return ctx
}
func generateTraceID() string {
return snowflake.New().Base36()
}