/
context.go
179 lines (160 loc) · 4.73 KB
/
context.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
174
175
176
177
178
179
package zlog
import (
"context"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
type ctxKey int
const (
fieldsKey ctxKey = iota
loggerKey
)
// CtxHandler customizes a logger's behavior at runtime dynamically.
type CtxHandler struct {
// ChangeLevel returns a non-nil Level if it wants to change
// the logger's logging level according to ctx.
// It returns nil to keep the logger's logging level as-is.
ChangeLevel func(ctx context.Context) *Level
// WithCtx is called by Logger.Ctx, SugaredLogger.Ctx and the global
// function WithCtx to check ctx for context-aware information.
// It returns CtxResult to customize the logger's behavior.
WithCtx func(ctx context.Context) CtxResult
}
// CtxResult holds information get from a context.
type CtxResult struct {
// Non-nil Level changes the logger's logging level.
Level *Level
// Fields will be added to the logger as additional fields.
Fields []zap.Field
}
// AddFields add logging fields to ctx which can be retrieved by
// GetFields, WithCtx.
// Duplicate fields override the old ones in ctx.
func AddFields(ctx context.Context, fields ...zap.Field) context.Context {
if len(fields) == 0 {
return ctx
}
old, ok := ctx.Value(fieldsKey).([]zap.Field)
if ok {
fields = appendFields(old, fields)
}
return context.WithValue(ctx, fieldsKey, fields)
}
// GetFields returns the logging fields added to ctx by AddFields.
func GetFields(ctx context.Context) []zap.Field {
fs, ok := ctx.Value(fieldsKey).([]zap.Field)
if ok {
return fs[:len(fs):len(fs)] // clip
}
return nil
}
// WithLogger returns a new context.Context with logger attached,
// which can be retrieved by WithCtx.
func WithLogger[T Logger | SugaredLogger | *zap.Logger | *zap.SugaredLogger](
ctx context.Context, logger T) context.Context {
ctx = context.WithValue(ctx, loggerKey, logger)
return ctx
}
func getLoggerFromCtx(ctx context.Context) any {
return ctx.Value(loggerKey)
}
// Ctx creates a child logger, it calls CtxHandler.WithCtx to get CtxResult
// from ctx, adds CtxResult.Fields to the child logger and changes
// the logger's level to CtxResult.Level, if it is not nil.
func (l Logger) Ctx(ctx context.Context, extra ...zap.Field) Logger {
if ctx == nil {
return l.With(extra...)
}
var fields []zap.Field
logger := l
ctxFunc := globals.Props.cfg.CtxHandler.WithCtx
if ctxFunc != nil {
ctxResult := ctxFunc(ctx)
fields = ctxResult.Fields
if ctxResult.Level != nil {
logger = logger.WithOptions(zap.WrapCore(changeLevel(*ctxResult.Level)))
}
}
fields = appendFields(fields, GetFields(ctx))
fields = appendFields(fields, extra)
if len(fields) > 0 {
logger = logger.With(fields...)
}
return logger
}
// Ctx creates a child logger, it calls CtxHandler.WithCtx to get CtxResult
// from ctx, adds CtxResult.Fields to the child logger and changes
// the logger's level to CtxResult.Level, if it is not nil.
func (s SugaredLogger) Ctx(ctx context.Context, extra ...zap.Field) SugaredLogger {
ctxFunc := globals.Props.cfg.CtxHandler.WithCtx
if (ctx == nil || ctxFunc == nil) && len(extra) == 0 {
return s
}
return s.Desugar().Ctx(ctx, extra...).Sugar()
}
// WithCtx creates a child logger and customizes its behavior using context
// data (e.g. adding fields, dynamically changing level, etc.)
//
// If ctx is created by WithLogger, it carries a logger instance,
// this function uses that logger as a base to create the child logger,
// else it calls Logger.Ctx to build the child logger with contextual fields
// and optional dynamic level from ctx.
//
// Also see WithLogger, CtxHandler and CtxResult for more details.
func WithCtx(ctx context.Context, extra ...zap.Field) Logger {
if ctx == nil {
return With(extra...)
}
if lg := getLoggerFromCtx(ctx); lg != nil {
var logger Logger
switch x := lg.(type) {
case Logger:
logger = x
case SugaredLogger:
logger = x.Desugar()
case *zap.Logger:
logger = Logger{Logger: x}
case *zap.SugaredLogger:
logger = Logger{Logger: x.Desugar()}
}
if logger.Logger != nil {
return logger.With(extra...)
}
}
return L().Ctx(ctx, extra...)
}
//nolint:predeclared
func appendFields(old []zap.Field, new []zap.Field) []zap.Field {
if len(new) == 0 {
return old
}
result := make([]zap.Field, len(old), len(old)+len(new))
copy(result, old)
// check namespace
nsIdx := 0
for i := len(result) - 1; i >= 0; i-- {
if result[i].Type == zapcore.NamespaceType {
nsIdx = i + 1
break
}
}
var hasNewNamespace bool
loop:
for _, f := range new {
if !hasNewNamespace {
if f.Type == zapcore.NamespaceType {
hasNewNamespace = true
result = append(result, f)
continue loop
}
for i := nsIdx; i < len(result); i++ {
if result[i].Key == f.Key {
result[i] = f
continue loop
}
}
}
result = append(result, f)
}
return result
}