-
Notifications
You must be signed in to change notification settings - Fork 0
/
grpc.go
109 lines (91 loc) · 3.11 KB
/
grpc.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
package bootstrap
import (
"context"
"time"
"google.golang.org/grpc"
"github.com/go-kratos/aegis/ratelimit"
"github.com/go-kratos/aegis/ratelimit/bbr"
"github.com/go-kratos/kratos/v2/log"
"github.com/go-kratos/kratos/v2/registry"
"github.com/go-kratos/kratos/v2/middleware"
midRateLimit "github.com/go-kratos/kratos/v2/middleware/ratelimit"
"github.com/go-kratos/kratos/v2/middleware/recovery"
"github.com/go-kratos/kratos/v2/middleware/tracing"
"github.com/go-kratos/kratos/v2/middleware/validate"
kratosGrpc "github.com/go-kratos/kratos/v2/transport/grpc"
conf "github.com/bugsmo/cy/contrib/kratos/bootstrap/gen/api/go/conf/v1"
)
const defaultTimeout = 5 * time.Second
// CreateGrpcClient 创建GRPC客户端
func CreateGrpcClient(ctx context.Context, r registry.Discovery, serviceName string, cfg *conf.Bootstrap, m ...middleware.Middleware) grpc.ClientConnInterface {
endpoint := "discovery:///" + serviceName
var ms []middleware.Middleware
timeout := defaultTimeout
if cfg.Client != nil && cfg.Client.Grpc != nil {
if cfg.Client.Grpc.Timeout != nil {
timeout = cfg.Client.Grpc.Timeout.AsDuration()
}
if cfg.Client.Grpc.Middleware != nil {
if cfg.Client.Grpc.Middleware.GetEnableRecovery() {
ms = append(ms, recovery.Recovery())
}
if cfg.Client.Grpc.Middleware.GetEnableTracing() {
ms = append(ms, tracing.Client())
}
if cfg.Client.Grpc.Middleware.GetEnableValidate() {
ms = append(ms, validate.Validator())
}
}
}
ms = append(ms, m...)
conn, err := kratosGrpc.DialInsecure(
ctx,
kratosGrpc.WithEndpoint(endpoint),
kratosGrpc.WithDiscovery(r),
kratosGrpc.WithTimeout(timeout),
kratosGrpc.WithMiddleware(ms...),
)
if err != nil {
log.Fatalf("dial grpc client [%s] failed: %s", serviceName, err.Error())
}
return conn
}
// CreateGrpcServer 创建GRPC服务端
func CreateGrpcServer(cfg *conf.Bootstrap, m ...middleware.Middleware) *kratosGrpc.Server {
var opts []kratosGrpc.ServerOption
var ms []middleware.Middleware
if cfg.Server != nil && cfg.Server.Grpc != nil && cfg.Server.Grpc.Middleware != nil {
if cfg.Server.Grpc.Middleware.GetEnableRecovery() {
ms = append(ms, recovery.Recovery())
}
if cfg.Server.Grpc.Middleware.GetEnableTracing() {
ms = append(ms, tracing.Server())
}
if cfg.Server.Grpc.Middleware.GetEnableValidate() {
ms = append(ms, validate.Validator())
}
if cfg.Server.Grpc.Middleware.GetEnableCircuitBreaker() {
}
if cfg.Server.Grpc.Middleware.Limiter != nil {
var limiter ratelimit.Limiter
switch cfg.Server.Grpc.Middleware.Limiter.GetName() {
case "bbr":
limiter = bbr.NewLimiter()
}
ms = append(ms, midRateLimit.Server(midRateLimit.WithLimiter(limiter)))
}
}
ms = append(ms, m...)
opts = append(opts, kratosGrpc.Middleware(ms...))
if cfg.Server.Grpc.Network != "" {
opts = append(opts, kratosGrpc.Network(cfg.Server.Grpc.Network))
}
if cfg.Server.Grpc.Addr != "" {
opts = append(opts, kratosGrpc.Address(cfg.Server.Grpc.Addr))
}
if cfg.Server.Grpc.Timeout != nil {
opts = append(opts, kratosGrpc.Timeout(cfg.Server.Grpc.Timeout.AsDuration()))
}
srv := kratosGrpc.NewServer(opts...)
return srv
}