forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
rpcserver.go
141 lines (118 loc) · 3.79 KB
/
rpcserver.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
package internal
import (
"fmt"
"net"
"github.com/zeromicro/go-zero/core/proc"
"github.com/zeromicro/go-zero/core/stat"
"github.com/zeromicro/go-zero/internal/health"
"github.com/zeromicro/go-zero/zrpc/internal/serverinterceptors"
"google.golang.org/grpc"
"google.golang.org/grpc/health/grpc_health_v1"
)
const probeNamePrefix = "zrpc"
type (
// ServerOption defines the method to customize a rpcServerOptions.
ServerOption func(options *rpcServerOptions)
rpcServerOptions struct {
metrics *stat.Metrics
health bool
}
rpcServer struct {
*baseRpcServer
name string
middlewares ServerMiddlewaresConf
healthManager health.Probe
}
)
// NewRpcServer returns a Server.
func NewRpcServer(addr string, middlewares ServerMiddlewaresConf, opts ...ServerOption) Server {
var options rpcServerOptions
for _, opt := range opts {
opt(&options)
}
if options.metrics == nil {
options.metrics = stat.NewMetrics(addr)
}
return &rpcServer{
baseRpcServer: newBaseRpcServer(addr, &options),
middlewares: middlewares,
healthManager: health.NewHealthManager(fmt.Sprintf("%s-%s", probeNamePrefix, addr)),
}
}
func (s *rpcServer) SetName(name string) {
s.name = name
s.baseRpcServer.SetName(name)
}
func (s *rpcServer) Start(register RegisterFn) error {
lis, err := net.Listen("tcp", s.address)
if err != nil {
return err
}
unaryInterceptorOption := grpc.ChainUnaryInterceptor(s.buildUnaryInterceptors()...)
streamInterceptorOption := grpc.ChainStreamInterceptor(s.buildStreamInterceptors()...)
options := append(s.options, unaryInterceptorOption, streamInterceptorOption)
server := grpc.NewServer(options...)
register(server)
// register the health check service
if s.health != nil {
grpc_health_v1.RegisterHealthServer(server, s.health)
s.health.Resume()
}
s.healthManager.MarkReady()
health.AddProbe(s.healthManager)
// we need to make sure all others are wrapped up,
// so we do graceful stop at shutdown phase instead of wrap up phase
waitForCalled := proc.AddWrapUpListener(func() {
if s.health != nil {
s.health.Shutdown()
}
server.GracefulStop()
})
defer waitForCalled()
return server.Serve(lis)
}
func (s *rpcServer) buildStreamInterceptors() []grpc.StreamServerInterceptor {
var interceptors []grpc.StreamServerInterceptor
if s.middlewares.Trace {
interceptors = append(interceptors, serverinterceptors.StreamTracingInterceptor)
}
if s.middlewares.Recover {
interceptors = append(interceptors, serverinterceptors.StreamRecoverInterceptor)
}
if s.middlewares.Breaker {
interceptors = append(interceptors, serverinterceptors.StreamBreakerInterceptor)
}
return append(interceptors, s.streamInterceptors...)
}
func (s *rpcServer) buildUnaryInterceptors() []grpc.UnaryServerInterceptor {
var interceptors []grpc.UnaryServerInterceptor
if s.middlewares.Trace {
interceptors = append(interceptors, serverinterceptors.UnaryTracingInterceptor)
}
if s.middlewares.Recover {
interceptors = append(interceptors, serverinterceptors.UnaryRecoverInterceptor)
}
if s.middlewares.Stat {
interceptors = append(interceptors,
serverinterceptors.UnaryStatInterceptor(s.metrics, s.middlewares.StatConf))
}
if s.middlewares.Prometheus {
interceptors = append(interceptors, serverinterceptors.UnaryPrometheusInterceptor)
}
if s.middlewares.Breaker {
interceptors = append(interceptors, serverinterceptors.UnaryBreakerInterceptor)
}
return append(interceptors, s.unaryInterceptors...)
}
// WithMetrics returns a func that sets metrics to a Server.
func WithMetrics(metrics *stat.Metrics) ServerOption {
return func(options *rpcServerOptions) {
options.metrics = metrics
}
}
// WithRpcHealth returns a func that sets rpc health switch to a Server.
func WithRpcHealth(health bool) ServerOption {
return func(options *rpcServerOptions) {
options.health = health
}
}