github.com/dahs81/otto@v0.2.1-0.20160126165905-6400716cf085/rpc/client.go (about)

     1  package rpc
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"net/rpc"
     7  
     8  	"github.com/hashicorp/otto/app"
     9  	"github.com/hashicorp/yamux"
    10  )
    11  
    12  // Client connects to a Server in order to request plugin implementations
    13  // for Otto.
    14  type Client struct {
    15  	broker  *muxBroker
    16  	control *rpc.Client
    17  
    18  	// These are the streams used for the various stdout/err overrides
    19  	stdout, stderr net.Conn
    20  }
    21  
    22  // Dial opens a connection to an RPC server and returns a client.
    23  func Dial(network, address string) (*Client, error) {
    24  	conn, err := net.Dial(network, address)
    25  	if err != nil {
    26  		return nil, err
    27  	}
    28  
    29  	if tcpConn, ok := conn.(*net.TCPConn); ok {
    30  		// Make sure to set keep alive so that the connection doesn't die
    31  		tcpConn.SetKeepAlive(true)
    32  	}
    33  
    34  	return NewClient(conn)
    35  }
    36  
    37  // NewClient creates a client from an already-open connection-like value.
    38  // Dial is typically used instead.
    39  func NewClient(conn io.ReadWriteCloser) (*Client, error) {
    40  	// Create the yamux client so we can multiplex
    41  	mux, err := yamux.Client(conn, nil)
    42  	if err != nil {
    43  		conn.Close()
    44  		return nil, err
    45  	}
    46  
    47  	// Connect to the control stream.
    48  	control, err := mux.Open()
    49  	if err != nil {
    50  		mux.Close()
    51  		return nil, err
    52  	}
    53  
    54  	// Connect stdout, stderr streams
    55  	stdstream := make([]net.Conn, 2)
    56  	for i, _ := range stdstream {
    57  		stdstream[i], err = mux.Open()
    58  		if err != nil {
    59  			mux.Close()
    60  			return nil, err
    61  		}
    62  	}
    63  
    64  	// Create the broker and start it up
    65  	broker := newMuxBroker(mux)
    66  	go broker.Run()
    67  
    68  	// Build the client using our broker and control channel.
    69  	return &Client{
    70  		broker:  broker,
    71  		control: rpc.NewClient(control),
    72  		stdout:  stdstream[0],
    73  		stderr:  stdstream[1],
    74  	}, nil
    75  }
    76  
    77  // SyncStreams should be called to enable syncing of stdout,
    78  // stderr with the plugin.
    79  //
    80  // This will return immediately and the syncing will continue to happen
    81  // in the background. You do not need to launch this in a goroutine itself.
    82  //
    83  // This should never be called multiple times.
    84  func (c *Client) SyncStreams(stdout io.Writer, stderr io.Writer) error {
    85  	go copyStream("stdout", stdout, c.stdout)
    86  	go copyStream("stderr", stderr, c.stderr)
    87  	return nil
    88  }
    89  
    90  // Close closes the connection. The client is no longer usable after this
    91  // is called.
    92  func (c *Client) Close() error {
    93  	if err := c.control.Close(); err != nil {
    94  		return err
    95  	}
    96  	if err := c.stdout.Close(); err != nil {
    97  		return err
    98  	}
    99  	if err := c.stderr.Close(); err != nil {
   100  		return err
   101  	}
   102  
   103  	return c.broker.Close()
   104  }
   105  
   106  func (c *Client) App() (app.App, error) {
   107  	var id uint32
   108  	if err := c.control.Call(
   109  		"Dispenser.App", new(interface{}), &id); err != nil {
   110  		return nil, err
   111  	}
   112  
   113  	conn, err := c.broker.Dial(id)
   114  	if err != nil {
   115  		return nil, err
   116  	}
   117  
   118  	return &App{
   119  		Broker: c.broker,
   120  		Client: rpc.NewClient(conn),
   121  		Name:   "App",
   122  	}, nil
   123  }