github.com/lingyao2333/mo-zero@v1.4.1/zrpc/client.go (about)

     1  package zrpc
     2  
     3  import (
     4  	"log"
     5  	"time"
     6  
     7  	"github.com/lingyao2333/mo-zero/zrpc/internal"
     8  	"github.com/lingyao2333/mo-zero/zrpc/internal/auth"
     9  	"github.com/lingyao2333/mo-zero/zrpc/internal/clientinterceptors"
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  var (
    14  	// WithDialOption is an alias of internal.WithDialOption.
    15  	WithDialOption = internal.WithDialOption
    16  	// WithNonBlock sets the dialing to be nonblock.
    17  	WithNonBlock = internal.WithNonBlock
    18  	// WithStreamClientInterceptor is an alias of internal.WithStreamClientInterceptor.
    19  	WithStreamClientInterceptor = internal.WithStreamClientInterceptor
    20  	// WithTimeout is an alias of internal.WithTimeout.
    21  	WithTimeout = internal.WithTimeout
    22  	// WithTransportCredentials return a func to make the gRPC calls secured with given credentials.
    23  	WithTransportCredentials = internal.WithTransportCredentials
    24  	// WithUnaryClientInterceptor is an alias of internal.WithUnaryClientInterceptor.
    25  	WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor
    26  )
    27  
    28  type (
    29  	// Client is an alias of internal.Client.
    30  	Client = internal.Client
    31  	// ClientOption is an alias of internal.ClientOption.
    32  	ClientOption = internal.ClientOption
    33  
    34  	// A RpcClient is a rpc client.
    35  	RpcClient struct {
    36  		client Client
    37  	}
    38  )
    39  
    40  // MustNewClient returns a Client, exits on any error.
    41  func MustNewClient(c RpcClientConf, options ...ClientOption) Client {
    42  	cli, err := NewClient(c, options...)
    43  	if err != nil {
    44  		log.Fatal(err)
    45  	}
    46  
    47  	return cli
    48  }
    49  
    50  // NewClient returns a Client.
    51  func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
    52  	var opts []ClientOption
    53  	if c.HasCredential() {
    54  		opts = append(opts, WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
    55  			App:   c.App,
    56  			Token: c.Token,
    57  		})))
    58  	}
    59  	if c.NonBlock {
    60  		opts = append(opts, WithNonBlock())
    61  	}
    62  	if c.Timeout > 0 {
    63  		opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
    64  	}
    65  
    66  	opts = append(opts, options...)
    67  
    68  	target, err := c.BuildTarget()
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  
    73  	client, err := internal.NewClient(target, opts...)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  
    78  	return &RpcClient{
    79  		client: client,
    80  	}, nil
    81  }
    82  
    83  // NewClientWithTarget returns a Client with connecting to given target.
    84  func NewClientWithTarget(target string, opts ...ClientOption) (Client, error) {
    85  	return internal.NewClient(target, opts...)
    86  }
    87  
    88  // Conn returns the underlying grpc.ClientConn.
    89  func (rc *RpcClient) Conn() *grpc.ClientConn {
    90  	return rc.client.Conn()
    91  }
    92  
    93  // SetClientSlowThreshold sets the slow threshold on client side.
    94  func SetClientSlowThreshold(threshold time.Duration) {
    95  	clientinterceptors.SetSlowThreshold(threshold)
    96  }