vendor/github.com/go-kit/kit/transport/grpc/client.go

   1 package grpc
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "reflect"
   7
   8         "google.golang.org/grpc"
   9         "google.golang.org/grpc/metadata"
  10
  11         "github.com/go-kit/kit/endpoint"
  12 )
  13
  14 // Client wraps a gRPC connection and provides a method that implements
  15 // endpoint.Endpoint.
  16 type Client struct {
  17         client      *grpc.ClientConn
  18         serviceName string
  19         method      string
  20         enc         EncodeRequestFunc
  21         dec         DecodeResponseFunc
  22         grpcReply   reflect.Type
  23         before      []ClientRequestFunc
  24         after       []ClientResponseFunc
  25 }
  26
  27 // NewClient constructs a usable Client for a single remote endpoint.
  28 // Pass an zero-value protobuf message of the RPC response type as
  29 // the grpcReply argument.
  30 func NewClient(
  31         cc *grpc.ClientConn,
  32         serviceName string,
  33         method string,
  34         enc EncodeRequestFunc,
  35         dec DecodeResponseFunc,
  36         grpcReply interface{},
  37         options ...ClientOption,
  38 ) *Client {
  39         c := &Client{
  40                 client: cc,
  41                 method: fmt.Sprintf("/%s/%s", serviceName, method),
  42                 enc:    enc,
  43                 dec:    dec,
  44                 // We are using reflect.Indirect here to allow both reply structs and
  45                 // pointers to these reply structs. New consumers of the client should
  46                 // use structs directly, while existing consumers will not break if they
  47                 // remain to use pointers to structs.
  48                 grpcReply: reflect.TypeOf(
  49                         reflect.Indirect(
  50                                 reflect.ValueOf(grpcReply),
  51                         ).Interface(),
  52                 ),
  53                 before: []ClientRequestFunc{},
  54                 after:  []ClientResponseFunc{},
  55         }
  56         for _, option := range options {
  57                 option(c)
  58         }
  59         return c
  60 }
  61
  62 // ClientOption sets an optional parameter for clients.
  63 type ClientOption func(*Client)
  64
  65 // ClientBefore sets the RequestFuncs that are applied to the outgoing gRPC
  66 // request before it's invoked.
  67 func ClientBefore(before ...ClientRequestFunc) ClientOption {
  68         return func(c *Client) { c.before = append(c.before, before...) }
  69 }
  70
  71 // ClientAfter sets the ClientResponseFuncs that are applied to the incoming
  72 // gRPC response prior to it being decoded. This is useful for obtaining
  73 // response metadata and adding onto the context prior to decoding.
  74 func ClientAfter(after ...ClientResponseFunc) ClientOption {
  75         return func(c *Client) { c.after = append(c.after, after...) }
  76 }
  77
  78 // Endpoint returns a usable endpoint that will invoke the gRPC specified by the
  79 // client.
  80 func (c Client) Endpoint() endpoint.Endpoint {
  81         return func(ctx context.Context, request interface{}) (interface{}, error) {
  82                 ctx, cancel := context.WithCancel(ctx)
  83                 defer cancel()
  84
  85                 req, err := c.enc(ctx, request)
  86                 if err != nil {
  87                         return nil, err
  88                 }
  89
  90                 md := &metadata.MD{}
  91                 for _, f := range c.before {
  92                         ctx = f(ctx, md)
  93                 }
  94                 ctx = metadata.NewOutgoingContext(ctx, *md)
  95
  96                 var header, trailer metadata.MD
  97                 grpcReply := reflect.New(c.grpcReply).Interface()
  98                 if err = grpc.Invoke(
  99                         ctx, c.method, req, grpcReply, c.client,
 100                         grpc.Header(&header), grpc.Trailer(&trailer),
 101                 ); err != nil {
 102                         return nil, err
 103                 }
 104
 105                 for _, f := range c.after {
 106                         ctx = f(ctx, header, trailer)
 107                 }
 108
 109                 response, err := c.dec(ctx, grpcReply)
 110                 if err != nil {
 111                         return nil, err
 112                 }
 113                 return response, nil
 114         }
 115 }