/
logger.go
112 lines (89 loc) · 2.47 KB
/
logger.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
package twirpzap
import (
"context"
"sync"
"time"
"github.com/twitchtv/twirp"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
type contextKey struct{}
var (
logKey = contextKey{}
nullLogger = zap.NewNop()
)
// ServerHooks creates twirp server hooks for logging
// using zap
func ServerHooks(logger *zap.Logger) *twirp.ServerHooks {
return &twirp.ServerHooks{
RequestReceived: func(ctx context.Context) (context.Context, error) {
return requestReceived(ctx, logger)
},
RequestRouted: responseRouted,
ResponseSent: responseSent,
}
}
type requestLogger struct {
startTime time.Time
logger *zap.Logger
fields []zap.Field
}
var requestLoggerPool = sync.Pool{
New: func() interface{} {
return &requestLogger{
fields: make([]zap.Field, 0, 10),
}
},
}
func requestReceived(ctx context.Context, logger *zap.Logger) (context.Context, error) {
r := requestLoggerPool.Get().(*requestLogger)
r.startTime = time.Now()
r.logger = logger
r.fields = r.fields[:]
if pkg, ok := twirp.PackageName(ctx); ok {
r.fields = append(r.fields, zap.String("twirp_package", pkg))
}
if svc, ok := twirp.ServiceName(ctx); ok {
r.fields = append(r.fields, zap.String("twirp_service", svc))
}
ctx = context.WithValue(ctx, logKey, r)
return ctx, nil
}
func responseRouted(ctx context.Context) (context.Context, error) {
if meth, ok := twirp.MethodName(ctx); ok {
AddFields(ctx, zap.String("twirp_method", meth))
}
return ctx, nil
}
func responseSent(ctx context.Context) {
r, ok := ctx.Value(logKey).(*requestLogger)
if !ok || r == nil {
return
}
duration := time.Since(r.startTime)
r.fields = append(r.fields, zap.Duration("twirp_duration", duration))
if status, ok := twirp.StatusCode(ctx); ok {
r.fields = append(r.fields, zap.String("twirp_status", status))
}
r.logger.With(r.fields...).Info("response sent")
r.logger = nullLogger
r.fields = r.fields[:]
requestLoggerPool.Put(r)
}
// based on https://github.com/grpc-ecosystem/go-grpc-middleware/blob/master/logging/zap/ctxzap/context.go
// AddFields adds zap fields to the logger.
func AddFields(ctx context.Context, fields ...zapcore.Field) {
l, ok := ctx.Value(logKey).(*requestLogger)
if !ok || l == nil {
return
}
l.fields = append(l.fields, fields...)
}
// FromContext returns the request scoped logger.
func FromContext(ctx context.Context) *zap.Logger {
l, ok := ctx.Value(logKey).(*requestLogger)
if !ok || l == nil || l.logger == nil {
return nullLogger
}
return l.logger.With(l.fields...)
}