vendor/github.com/hashicorp/go-plugin/examples/bidirectional/shared/grpc.go

   1 package shared
   2
   3 import (
   4         hclog "github.com/hashicorp/go-hclog"
   5         plugin "github.com/hashicorp/go-plugin"
   6         "github.com/hashicorp/go-plugin/examples/bidirectional/proto"
   7         "golang.org/x/net/context"
   8         "google.golang.org/grpc"
   9 )
  10
  11 // GRPCClient is an implementation of KV that talks over RPC.
  12 type GRPCClient struct {
  13         broker *plugin.GRPCBroker
  14         client proto.CounterClient
  15 }
  16
  17 func (m *GRPCClient) Put(key string, value int64, a AddHelper) error {
  18         addHelperServer := &GRPCAddHelperServer{Impl: a}
  19
  20         var s *grpc.Server
  21         serverFunc := func(opts []grpc.ServerOption) *grpc.Server {
  22                 s = grpc.NewServer(opts...)
  23                 proto.RegisterAddHelperServer(s, addHelperServer)
  24
  25                 return s
  26         }
  27
  28         brokerID := m.broker.NextId()
  29         go m.broker.AcceptAndServe(brokerID, serverFunc)
  30
  31         _, err := m.client.Put(context.Background(), &proto.PutRequest{
  32                 AddServer: brokerID,
  33                 Key:       key,
  34                 Value:     value,
  35         })
  36
  37         s.Stop()
  38         return err
  39 }
  40
  41 func (m *GRPCClient) Get(key string) (int64, error) {
  42         resp, err := m.client.Get(context.Background(), &proto.GetRequest{
  43                 Key: key,
  44         })
  45         if err != nil {
  46                 return 0, err
  47         }
  48
  49         return resp.Value, nil
  50 }
  51
  52 // Here is the gRPC server that GRPCClient talks to.
  53 type GRPCServer struct {
  54         // This is the real implementation
  55         Impl Counter
  56
  57         broker *plugin.GRPCBroker
  58 }
  59
  60 func (m *GRPCServer) Put(ctx context.Context, req *proto.PutRequest) (*proto.Empty, error) {
  61         conn, err := m.broker.Dial(req.AddServer)
  62         if err != nil {
  63                 return nil, err
  64         }
  65         defer conn.Close()
  66
  67         a := &GRPCAddHelperClient{proto.NewAddHelperClient(conn)}
  68         return &proto.Empty{}, m.Impl.Put(req.Key, req.Value, a)
  69 }
  70
  71 func (m *GRPCServer) Get(ctx context.Context, req *proto.GetRequest) (*proto.GetResponse, error) {
  72         v, err := m.Impl.Get(req.Key)
  73         return &proto.GetResponse{Value: v}, err
  74 }
  75
  76 // GRPCClient is an implementation of KV that talks over RPC.
  77 type GRPCAddHelperClient struct{ client proto.AddHelperClient }
  78
  79 func (m *GRPCAddHelperClient) Sum(a, b int64) (int64, error) {
  80         resp, err := m.client.Sum(context.Background(), &proto.SumRequest{
  81                 A: a,
  82                 B: b,
  83         })
  84         if err != nil {
  85                 hclog.Default().Info("add.Sum", "client", "start", "err", err)
  86                 return 0, err
  87         }
  88         return resp.R, err
  89 }
  90
  91 // Here is the gRPC server that GRPCClient talks to.
  92 type GRPCAddHelperServer struct {
  93         // This is the real implementation
  94         Impl AddHelper
  95 }
  96
  97 func (m *GRPCAddHelperServer) Sum(ctx context.Context, req *proto.SumRequest) (resp *proto.SumResponse, err error) {
  98         r, err := m.Impl.Sum(req.A, req.B)
  99         if err != nil {
 100                 return nil, err
 101         }
 102         return &proto.SumResponse{R: r}, err
 103 }