client.go 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. package grpclient
  2. import (
  3. "context"
  4. "errors"
  5. "time"
  6. "google.golang.org/grpc"
  7. "google.golang.org/grpc/credentials"
  8. "google.golang.org/grpc/credentials/insecure"
  9. "google.golang.org/grpc/keepalive"
  10. )
  11. var (
  12. defaultDialTimeout = time.Second * 2
  13. )
  14. type Option func(*option)
  15. type option struct {
  16. credential credentials.TransportCredentials
  17. keepalive *keepalive.ClientParameters
  18. dialTimeout time.Duration
  19. unaryInterceptor grpc.UnaryClientInterceptor
  20. }
  21. // WithCredential setup credential for tls
  22. func WithCredential(credential credentials.TransportCredentials) Option {
  23. return func(opt *option) {
  24. opt.credential = credential
  25. }
  26. }
  27. // WithKeepAlive setup keepalive parameters
  28. func WithKeepAlive(keepalive *keepalive.ClientParameters) Option {
  29. return func(opt *option) {
  30. opt.keepalive = keepalive
  31. }
  32. }
  33. // WithDialTimeout set up the dial timeout
  34. func WithDialTimeout(timeout time.Duration) Option {
  35. return func(opt *option) {
  36. opt.dialTimeout = timeout
  37. }
  38. }
  39. func WithUnaryInterceptor(unaryInterceptor grpc.UnaryClientInterceptor) Option {
  40. return func(opt *option) {
  41. opt.unaryInterceptor = unaryInterceptor
  42. }
  43. }
  44. func New(target string, options ...Option) (*grpc.ClientConn, error) {
  45. if target == "" {
  46. return nil, errors.New("target required")
  47. }
  48. opt := new(option)
  49. for _, f := range options {
  50. f(opt)
  51. }
  52. kacp := defaultKeepAlive
  53. if opt.keepalive != nil {
  54. kacp = opt.keepalive
  55. }
  56. dialTimeout := defaultDialTimeout
  57. if opt.dialTimeout > 0 {
  58. dialTimeout = opt.dialTimeout
  59. }
  60. dialOptions := []grpc.DialOption{
  61. grpc.WithBlock(),
  62. grpc.WithKeepaliveParams(*kacp),
  63. }
  64. if opt.unaryInterceptor != nil {
  65. dialOptions = append(dialOptions, grpc.WithUnaryInterceptor(opt.unaryInterceptor))
  66. }
  67. if opt.credential == nil {
  68. dialOptions = append(dialOptions, grpc.WithTransportCredentials(insecure.NewCredentials()))
  69. } else {
  70. dialOptions = append(dialOptions, grpc.WithTransportCredentials(opt.credential))
  71. }
  72. ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
  73. defer cancel()
  74. conn, err := grpc.DialContext(ctx, target, dialOptions...)
  75. if err != nil {
  76. return nil, err
  77. }
  78. return conn, nil
  79. }