client.go 2.1 KB

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