/
server.go
74 lines (60 loc) · 2.47 KB
/
server.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
// Copyright (c) The go-grpc-middleware Authors.
// Licensed under the Apache License 2.0.
// Go gRPC Middleware monitoring interceptors for server-side gRPC.
package interceptors
import (
"context"
"time"
"google.golang.org/grpc"
)
// UnaryServerInterceptor is a gRPC server-side interceptor that provides reporting for Unary RPCs.
func UnaryServerInterceptor(reportable ServerReportable) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
r := newReport(Unary, info.FullMethod)
reporter, newCtx := reportable.ServerReporter(ctx, CallMeta{ReqProtoOrNil: req, Typ: r.rpcType, Service: r.service, Method: r.method})
reporter.PostMsgReceive(req, nil, time.Since(r.startTime))
resp, err := handler(newCtx, req)
reporter.PostMsgSend(resp, err, time.Since(r.startTime))
reporter.PostCall(err, time.Since(r.startTime))
return resp, err
}
}
// StreamServerInterceptor is a gRPC server-side interceptor that provides reporting for Streaming RPCs.
func StreamServerInterceptor(reportable ServerReportable) grpc.StreamServerInterceptor {
return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
r := newReport(ServerStream, info.FullMethod)
reporter, newCtx := reportable.ServerReporter(ss.Context(), CallMeta{ReqProtoOrNil: nil, Typ: StreamRPCType(info), Service: r.service, Method: r.method})
err := handler(srv, &monitoredServerStream{ServerStream: ss, newCtx: newCtx, reporter: reporter})
reporter.PostCall(err, time.Since(r.startTime))
return err
}
}
func StreamRPCType(info *grpc.StreamServerInfo) GRPCType {
if info.IsClientStream && !info.IsServerStream {
return ClientStream
} else if !info.IsClientStream && info.IsServerStream {
return ServerStream
}
return BidiStream
}
// monitoredStream wraps grpc.ServerStream allowing each Sent/Recv of message to report.
type monitoredServerStream struct {
grpc.ServerStream
newCtx context.Context
reporter Reporter
}
func (s *monitoredServerStream) Context() context.Context {
return s.newCtx
}
func (s *monitoredServerStream) SendMsg(m interface{}) error {
start := time.Now()
err := s.ServerStream.SendMsg(m)
s.reporter.PostMsgSend(m, err, time.Since(start))
return err
}
func (s *monitoredServerStream) RecvMsg(m interface{}) error {
start := time.Now()
err := s.ServerStream.RecvMsg(m)
s.reporter.PostMsgReceive(m, err, time.Since(start))
return err
}