base-golang/pkg/grpclient/client.go
2024-07-23 10:23:43 +08:00

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
}