You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/zrpc/internal/client.go

105 lines
2.4 KiB
Go

4 years ago
package internal
4 years ago
import (
"context"
"errors"
4 years ago
"fmt"
"strings"
4 years ago
"time"
"github.com/tal-tech/go-zero/zrpc/internal/balancer/p2c"
"github.com/tal-tech/go-zero/zrpc/internal/clientinterceptors"
"github.com/tal-tech/go-zero/zrpc/internal/resolver"
4 years ago
"google.golang.org/grpc"
)
const (
dialTimeout = time.Second * 3
separator = '/'
)
4 years ago
func init() {
resolver.RegisterResolver()
}
4 years ago
type (
ClientOptions struct {
Timeout time.Duration
DialOptions []grpc.DialOption
}
ClientOption func(options *ClientOptions)
client struct {
conn *grpc.ClientConn
}
4 years ago
)
func NewClient(target string, opts ...ClientOption) (*client, error) {
opts = append(opts, WithDialOption(grpc.WithBalancerName(p2c.Name)))
conn, err := dial(target, opts...)
if err != nil {
return nil, err
}
return &client{conn: conn}, nil
}
func (c *client) Conn() *grpc.ClientConn {
return c.conn
}
4 years ago
func WithDialOption(opt grpc.DialOption) ClientOption {
return func(options *ClientOptions) {
options.DialOptions = append(options.DialOptions, opt)
}
}
func WithTimeout(timeout time.Duration) ClientOption {
return func(options *ClientOptions) {
options.Timeout = timeout
}
}
func buildDialOptions(opts ...ClientOption) []grpc.DialOption {
var clientOptions ClientOptions
for _, opt := range opts {
opt(&clientOptions)
}
options := []grpc.DialOption{
grpc.WithInsecure(),
grpc.WithBlock(),
WithUnaryClientInterceptors(
clientinterceptors.TracingInterceptor,
4 years ago
clientinterceptors.DurationInterceptor,
clientinterceptors.BreakerInterceptor,
clientinterceptors.PrometheusInterceptor,
4 years ago
clientinterceptors.TimeoutInterceptor(clientOptions.Timeout),
),
}
return append(options, clientOptions.DialOptions...)
}
func dial(server string, opts ...ClientOption) (*grpc.ClientConn, error) {
options := buildDialOptions(opts...)
timeCtx, cancel := context.WithTimeout(context.Background(), dialTimeout)
defer cancel()
conn, err := grpc.DialContext(timeCtx, server, options...)
if err != nil {
service := server
if errors.Is(err, context.DeadlineExceeded) {
pos := strings.LastIndexByte(server, separator)
// len(server) - 1 is the index of last char
if 0 < pos && pos < len(server)-1 {
service = server[pos+1:]
}
}
return nil, fmt.Errorf("rpc dial: %s, error: %s, make sure rpc service %q is alread started",
server, err.Error(), service)
4 years ago
}
return conn, nil
}