99 lines
2.1 KiB
Go
99 lines
2.1 KiB
Go
|
package grpclient
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
"google.golang.org/grpc/credentials"
|
||
|
|
"google.golang.org/grpc/credentials/insecure"
|
||
|
|
"google.golang.org/grpc/keepalive"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
defaultDialTimeout = time.Second * 2
|
||
|
|
)
|
||
|
|
|
||
|
|
type Option func(*option)
|
||
|
|
|
||
|
|
type option struct {
|
||
|
|
credential credentials.TransportCredentials
|
||
|
|
keepalive *keepalive.ClientParameters
|
||
|
|
dialTimeout time.Duration
|
||
|
|
unaryInterceptor grpc.UnaryClientInterceptor
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithCredential setup credential for tls
|
||
|
|
func WithCredential(credential credentials.TransportCredentials) Option {
|
||
|
|
return func(opt *option) {
|
||
|
|
opt.credential = credential
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithKeepAlive setup keepalive parameters
|
||
|
|
func WithKeepAlive(keepalive *keepalive.ClientParameters) Option {
|
||
|
|
return func(opt *option) {
|
||
|
|
opt.keepalive = keepalive
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithDialTimeout set up the dial timeout
|
||
|
|
func WithDialTimeout(timeout time.Duration) Option {
|
||
|
|
return func(opt *option) {
|
||
|
|
opt.dialTimeout = timeout
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func WithUnaryInterceptor(unaryInterceptor grpc.UnaryClientInterceptor) Option {
|
||
|
|
return func(opt *option) {
|
||
|
|
opt.unaryInterceptor = unaryInterceptor
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func New(target string, options ...Option) (*grpc.ClientConn, error) {
|
||
|
|
if target == "" {
|
||
|
|
return nil, errors.New("target required")
|
||
|
|
}
|
||
|
|
|
||
|
|
opt := new(option)
|
||
|
|
for _, f := range options {
|
||
|
|
f(opt)
|
||
|
|
}
|
||
|
|
|
||
|
|
kacp := defaultKeepAlive
|
||
|
|
if opt.keepalive != nil {
|
||
|
|
kacp = opt.keepalive
|
||
|
|
}
|
||
|
|
|
||
|
|
dialTimeout := defaultDialTimeout
|
||
|
|
if opt.dialTimeout > 0 {
|
||
|
|
dialTimeout = opt.dialTimeout
|
||
|
|
}
|
||
|
|
|
||
|
|
dialOptions := []grpc.DialOption{
|
||
|
|
grpc.WithBlock(),
|
||
|
|
grpc.WithKeepaliveParams(*kacp),
|
||
|
|
}
|
||
|
|
|
||
|
|
if opt.unaryInterceptor != nil {
|
||
|
|
dialOptions = append(dialOptions, grpc.WithUnaryInterceptor(opt.unaryInterceptor))
|
||
|
|
}
|
||
|
|
|
||
|
|
if opt.credential == nil {
|
||
|
|
dialOptions = append(dialOptions, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
||
|
|
} else {
|
||
|
|
dialOptions = append(dialOptions, grpc.WithTransportCredentials(opt.credential))
|
||
|
|
}
|
||
|
|
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
conn, err := grpc.DialContext(ctx, target, dialOptions...)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return conn, nil
|
||
|
|
}
|