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.
91 lines
2.0 KiB
Go
91 lines
2.0 KiB
Go
package internal
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"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"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
const dialTimeout = time.Second * 3
|
|
|
|
func init() {
|
|
resolver.RegisterResolver()
|
|
}
|
|
|
|
type (
|
|
ClientOptions struct {
|
|
Timeout time.Duration
|
|
DialOptions []grpc.DialOption
|
|
}
|
|
|
|
ClientOption func(options *ClientOptions)
|
|
|
|
client struct {
|
|
conn *grpc.ClientConn
|
|
}
|
|
)
|
|
|
|
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
|
|
}
|
|
|
|
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.BreakerInterceptor,
|
|
clientinterceptors.DurationInterceptor,
|
|
clientinterceptors.PromMetricInterceptor,
|
|
clientinterceptors.TimeoutInterceptor(clientOptions.Timeout),
|
|
clientinterceptors.TracingInterceptor,
|
|
),
|
|
}
|
|
|
|
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 {
|
|
return nil, fmt.Errorf("rpc dial: %s, error: %s", server, err.Error())
|
|
}
|
|
|
|
return conn, nil
|
|
}
|