vendor/github.com/hashicorp/go-plugin/rpc_client.go

   1 package plugin
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "net"
   7         "net/rpc"
   8
   9         "github.com/hashicorp/yamux"
  10 )
  11
  12 // RPCClient connects to an RPCServer over net/rpc to dispense plugin types.
  13 type RPCClient struct {
  14         broker  *MuxBroker
  15         control *rpc.Client
  16         plugins map[string]Plugin
  17
  18         // These are the streams used for the various stdout/err overrides
  19         stdout, stderr net.Conn
  20 }
  21
  22 // NewRPCClient creates a client from an already-open connection-like value.
  23 // Dial is typically used instead.
  24 func NewRPCClient(conn io.ReadWriteCloser, plugins map[string]Plugin) (*RPCClient, error) {
  25         // Create the yamux client so we can multiplex
  26         mux, err := yamux.Client(conn, nil)
  27         if err != nil {
  28                 conn.Close()
  29                 return nil, err
  30         }
  31
  32         // Connect to the control stream.
  33         control, err := mux.Open()
  34         if err != nil {
  35                 mux.Close()
  36                 return nil, err
  37         }
  38
  39         // Connect stdout, stderr streams
  40         stdstream := make([]net.Conn, 2)
  41         for i, _ := range stdstream {
  42                 stdstream[i], err = mux.Open()
  43                 if err != nil {
  44                         mux.Close()
  45                         return nil, err
  46                 }
  47         }
  48
  49         // Create the broker and start it up
  50         broker := newMuxBroker(mux)
  51         go broker.Run()
  52
  53         // Build the client using our broker and control channel.
  54         return &RPCClient{
  55                 broker:  broker,
  56                 control: rpc.NewClient(control),
  57                 plugins: plugins,
  58                 stdout:  stdstream[0],
  59                 stderr:  stdstream[1],
  60         }, nil
  61 }
  62
  63 // SyncStreams should be called to enable syncing of stdout,
  64 // stderr with the plugin.
  65 //
  66 // This will return immediately and the syncing will continue to happen
  67 // in the background. You do not need to launch this in a goroutine itself.
  68 //
  69 // This should never be called multiple times.
  70 func (c *RPCClient) SyncStreams(stdout io.Writer, stderr io.Writer) error {
  71         go copyStream("stdout", stdout, c.stdout)
  72         go copyStream("stderr", stderr, c.stderr)
  73         return nil
  74 }
  75
  76 // Close closes the connection. The client is no longer usable after this
  77 // is called.
  78 func (c *RPCClient) Close() error {
  79         // Call the control channel and ask it to gracefully exit. If this
  80         // errors, then we save it so that we always return an error but we
  81         // want to try to close the other channels anyways.
  82         var empty struct{}
  83         returnErr := c.control.Call("Control.Quit", true, &empty)
  84
  85         // Close the other streams we have
  86         if err := c.control.Close(); err != nil {
  87                 return err
  88         }
  89         if err := c.stdout.Close(); err != nil {
  90                 return err
  91         }
  92         if err := c.stderr.Close(); err != nil {
  93                 return err
  94         }
  95         if err := c.broker.Close(); err != nil {
  96                 return err
  97         }
  98
  99         // Return back the error we got from Control.Quit. This is very important
 100         // since we MUST return non-nil error if this fails so that Client.Kill
 101         // will properly try a process.Kill.
 102         return returnErr
 103 }
 104
 105 func (c *RPCClient) Dispense(name string) (interface{}, error) {
 106         p, ok := c.plugins[name]
 107         if !ok {
 108                 return nil, fmt.Errorf("unknown plugin type: %s", name)
 109         }
 110
 111         var id uint32
 112         if err := c.control.Call(
 113                 "Dispenser.Dispense", name, &id); err != nil {
 114                 return nil, err
 115         }
 116
 117         conn, err := c.broker.Dial(id)
 118         if err != nil {
 119                 return nil, err
 120         }
 121
 122         return p.Client(c.broker, rpc.NewClient(conn))
 123 }