forked from bojand/ghz
-
Notifications
You must be signed in to change notification settings - Fork 0
/
stats_handler.go
48 lines (39 loc) · 1.06 KB
/
stats_handler.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
package ghz
import (
"context"
"time"
"google.golang.org/grpc/stats"
"google.golang.org/grpc/status"
)
// StatsHandler is for gRPC stats
type statsHandler struct {
results chan *callResult
}
// HandleConn handle the connection
func (c *statsHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
// no-op
}
// TagConn exists to satisfy gRPC stats.Handler.
func (c *statsHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) context.Context {
// no-op
return ctx
}
// HandleRPC implements per-RPC tracing and stats instrumentation.
func (c *statsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
switch rs.(type) {
case *stats.End:
rpcStats := rs.(*stats.End)
end := time.Now()
duration := end.Sub(rpcStats.BeginTime)
var st string
s, ok := status.FromError(rpcStats.Error)
if ok {
st = s.Code().String()
}
c.results <- &callResult{rpcStats.Error, st, duration}
}
}
// TagRPC implements per-RPC context management.
func (c *statsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
return ctx
}