-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.go
77 lines (63 loc) · 1.81 KB
/
client.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
package grpcutils
import (
"context"
"github.com/hashicorp/go-multierror"
"go.uber.org/zap"
"google.golang.org/grpc"
)
type DialOption struct {
grpc.EmptyDialOption
function func(conn *ClientConnection)
}
func NewDialOption(f func(conn *ClientConnection)) DialOption {
return DialOption{EmptyDialOption: grpc.EmptyDialOption{}, function: f}
}
type ClientConnection struct {
Connection *grpc.ClientConn
addr string
logger *zap.Logger
dialOpts []grpc.DialOption
unaryInterceptors []grpc.UnaryClientInterceptor
streamInterceptors []grpc.StreamClientInterceptor
cleanFuncs []func(c *ClientConnection) error
}
func NewClientConnection(addr string) *ClientConnection {
return &ClientConnection{addr: addr}
}
func (c *ClientConnection) Dial(opts ...grpc.DialOption) error {
return c.DialContext(context.Background(), opts...)
}
func (c *ClientConnection) DialContext(ctx context.Context, opts ...grpc.DialOption) error {
for _, opt := range opts {
if o, ok := opt.(DialOption); ok {
o.function(c)
} else {
c.dialOpts = append(c.dialOpts, opt)
}
}
if c.logger == nil {
c.logger = zap.L()
}
c.logger = c.logger.With(zap.String("addr", c.addr))
c.dialOpts = append(c.dialOpts,
grpc.WithChainUnaryInterceptor(c.unaryInterceptors...),
grpc.WithChainStreamInterceptor(c.streamInterceptors...),
)
c.logger.Info("dialing gRPC client connection...")
conn, err := grpc.NewClient(c.addr, c.dialOpts...)
if err != nil {
return err
}
c.Connection = conn
c.cleanFuncs = append(c.cleanFuncs, func(c *ClientConnection) error {
return c.Connection.Close()
})
return nil
}
func (c *ClientConnection) CleanUp() error {
var err *multierror.Error
for _, clean := range c.cleanFuncs {
err = multierror.Append(err, clean(c))
}
return err.ErrorOrNil()
}