packer/rpc: implement Communicator

This commit is contained in:
Mitchell Hashimoto 2013-12-10 11:43:02 -08:00
parent 72fcb566a6
commit db06fc7501
6 changed files with 77 additions and 95 deletions

View File

@ -51,6 +51,13 @@ func (c *Client) Cache() packer.Cache {
}
}
func (c *Client) Communicator() packer.Communicator {
return &communicator{
client: c.client,
mux: c.mux,
}
}
func (c *Client) PostProcessor() packer.PostProcessor {
return &postProcessor{
client: c.client,

View File

@ -2,11 +2,9 @@ package rpc
import (
"encoding/gob"
"errors"
"github.com/mitchellh/packer/packer"
"io"
"log"
"net"
"net/rpc"
)
@ -14,12 +12,14 @@ import (
// executed over an RPC connection.
type communicator struct {
client *rpc.Client
mux *MuxConn
}
// CommunicatorServer wraps a packer.Communicator implementation and makes
// it exportable as part of a Golang RPC server.
type CommunicatorServer struct {
c packer.Communicator
mux *MuxConn
}
type CommandFinished struct {
@ -28,20 +28,20 @@ type CommandFinished struct {
type CommunicatorStartArgs struct {
Command string
StdinAddress string
StdoutAddress string
StderrAddress string
ResponseAddress string
StdinStreamId uint32
StdoutStreamId uint32
StderrStreamId uint32
ResponseStreamId uint32
}
type CommunicatorDownloadArgs struct {
Path string
WriterAddress string
WriterStreamId uint32
}
type CommunicatorUploadArgs struct {
Path string
ReaderAddress string
ReaderStreamId uint32
}
type CommunicatorUploadDirArgs struct {
@ -51,7 +51,7 @@ type CommunicatorUploadDirArgs struct {
}
func Communicator(client *rpc.Client) *communicator {
return &communicator{client}
return &communicator{client: client}
}
func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
@ -59,41 +59,38 @@ func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
args.Command = cmd.Command
if cmd.Stdin != nil {
stdinL := netListenerInRange(portRangeMin, portRangeMax)
args.StdinAddress = stdinL.Addr().String()
go serveSingleCopy("stdin", stdinL, nil, cmd.Stdin)
args.StdinStreamId = c.mux.NextId()
go serveSingleCopy("stdin", c.mux, args.StdinStreamId, nil, cmd.Stdin)
}
if cmd.Stdout != nil {
stdoutL := netListenerInRange(portRangeMin, portRangeMax)
args.StdoutAddress = stdoutL.Addr().String()
go serveSingleCopy("stdout", stdoutL, cmd.Stdout, nil)
args.StdoutStreamId = c.mux.NextId()
go serveSingleCopy("stdout", c.mux, args.StdoutStreamId, cmd.Stdout, nil)
}
if cmd.Stderr != nil {
stderrL := netListenerInRange(portRangeMin, portRangeMax)
args.StderrAddress = stderrL.Addr().String()
go serveSingleCopy("stderr", stderrL, cmd.Stderr, nil)
args.StderrStreamId = c.mux.NextId()
go serveSingleCopy("stderr", c.mux, args.StderrStreamId, cmd.Stderr, nil)
}
responseL := netListenerInRange(portRangeMin, portRangeMax)
args.ResponseAddress = responseL.Addr().String()
responseStreamId := c.mux.NextId()
args.ResponseStreamId = responseStreamId
go func() {
defer responseL.Close()
conn, err := responseL.Accept()
conn, err := c.mux.Accept(responseStreamId)
if err != nil {
log.Printf("[ERR] Error accepting response stream %d: %s",
responseStreamId, err)
cmd.SetExited(123)
return
}
defer conn.Close()
decoder := gob.NewDecoder(conn)
var finished CommandFinished
decoder := gob.NewDecoder(conn)
if err := decoder.Decode(&finished); err != nil {
log.Printf("[ERR] Error decoding response stream %d: %s",
responseStreamId, err)
cmd.SetExited(123)
return
}
@ -106,23 +103,13 @@ func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
}
func (c *communicator) Upload(path string, r io.Reader) (err error) {
// We need to create a server that can proxy the reader data
// over because we can't simply gob encode an io.Reader
readerL := netListenerInRange(portRangeMin, portRangeMax)
if readerL == nil {
err = errors.New("couldn't allocate listener for upload reader")
return
}
// Make sure at the end of this call, we close the listener
defer readerL.Close()
// Pipe the reader through to the connection
go serveSingleCopy("uploadReader", readerL, nil, r)
streamId := c.mux.NextId()
go serveSingleCopy("uploadReader", c.mux, streamId, nil, r)
args := CommunicatorUploadArgs{
path,
readerL.Addr().String(),
Path: path,
ReaderStreamId: streamId,
}
err = c.client.Call("Communicator.Upload", &args, new(interface{}))
@ -146,23 +133,13 @@ func (c *communicator) UploadDir(dst string, src string, exclude []string) error
}
func (c *communicator) Download(path string, w io.Writer) (err error) {
// We need to create a server that can proxy that data downloaded
// into the writer because we can't gob encode a writer directly.
writerL := netListenerInRange(portRangeMin, portRangeMax)
if writerL == nil {
err = errors.New("couldn't allocate listener for download writer")
return
}
// Make sure we close the listener once we're done because we'll be done
defer writerL.Close()
// Serve a single connection and a single copy
go serveSingleCopy("downloadWriter", writerL, w, nil)
streamId := c.mux.NextId()
go serveSingleCopy("downloadWriter", c.mux, streamId, w, nil)
args := CommunicatorDownloadArgs{
path,
writerL.Addr().String(),
Path: path,
WriterStreamId: streamId,
}
err = c.client.Call("Communicator.Download", &args, new(interface{}))
@ -175,40 +152,40 @@ func (c *CommunicatorServer) Start(args *CommunicatorStartArgs, reply *interface
var cmd packer.RemoteCmd
cmd.Command = args.Command
toClose := make([]net.Conn, 0)
if args.StdinAddress != "" {
stdinC, err := tcpDial(args.StdinAddress)
toClose := make([]io.Closer, 0)
if args.StdinStreamId > 0 {
conn, err := c.mux.Dial(args.StdinStreamId)
if err != nil {
return err
}
toClose = append(toClose, stdinC)
cmd.Stdin = stdinC
toClose = append(toClose, conn)
cmd.Stdin = conn
}
if args.StdoutAddress != "" {
stdoutC, err := tcpDial(args.StdoutAddress)
if args.StdoutStreamId > 0 {
conn, err := c.mux.Dial(args.StdoutStreamId)
if err != nil {
return err
}
toClose = append(toClose, stdoutC)
cmd.Stdout = stdoutC
toClose = append(toClose, conn)
cmd.Stdout = conn
}
if args.StderrAddress != "" {
stderrC, err := tcpDial(args.StderrAddress)
if args.StderrStreamId > 0 {
conn, err := c.mux.Dial(args.StderrStreamId)
if err != nil {
return err
}
toClose = append(toClose, stderrC)
cmd.Stderr = stderrC
toClose = append(toClose, conn)
cmd.Stderr = conn
}
// Connect to the response address so we can write our result to it
// when ready.
responseC, err := tcpDial(args.ResponseAddress)
responseC, err := c.mux.Dial(args.ResponseStreamId)
if err != nil {
return err
}
@ -234,11 +211,10 @@ func (c *CommunicatorServer) Start(args *CommunicatorStartArgs, reply *interface
}
func (c *CommunicatorServer) Upload(args *CommunicatorUploadArgs, reply *interface{}) (err error) {
readerC, err := tcpDial(args.ReaderAddress)
readerC, err := c.mux.Dial(args.ReaderStreamId)
if err != nil {
return
}
defer readerC.Close()
err = c.c.Upload(args.Path, readerC)
@ -250,21 +226,18 @@ func (c *CommunicatorServer) UploadDir(args *CommunicatorUploadDirArgs, reply *e
}
func (c *CommunicatorServer) Download(args *CommunicatorDownloadArgs, reply *interface{}) (err error) {
writerC, err := tcpDial(args.WriterAddress)
writerC, err := c.mux.Dial(args.WriterStreamId)
if err != nil {
return
}
defer writerC.Close()
err = c.c.Download(args.Path, writerC)
return
}
func serveSingleCopy(name string, l net.Listener, dst io.Writer, src io.Reader) {
defer l.Close()
conn, err := l.Accept()
func serveSingleCopy(name string, mux *MuxConn, id uint32, dst io.Writer, src io.Reader) {
conn, err := mux.Accept(id)
if err != nil {
log.Printf("'%s' accept error: %s", name, err)
return

View File

@ -4,7 +4,6 @@ import (
"bufio"
"github.com/mitchellh/packer/packer"
"io"
"net/rpc"
"reflect"
"testing"
)
@ -14,16 +13,11 @@ func TestCommunicatorRPC(t *testing.T) {
c := new(packer.MockCommunicator)
// Start the server
server := rpc.NewServer()
RegisterCommunicator(server, c)
address := serveSingleConn(server)
// Create the client over RPC and run some methods to verify it works
client, err := rpc.Dial("tcp", address)
if err != nil {
t.Fatalf("err: %s", err)
}
remote := Communicator(client)
client, server := testClientServer(t)
defer client.Close()
defer server.Close()
server.RegisterCommunicator(c)
remote := client.Communicator()
// The remote command we'll use
stdin_r, stdin_w := io.Pipe()
@ -42,7 +36,7 @@ func TestCommunicatorRPC(t *testing.T) {
c.StartExitStatus = 42
// Test Start
err = remote.Start(&cmd)
err := remote.Start(&cmd)
if err != nil {
t.Fatalf("err: %s", err)
}
@ -74,7 +68,7 @@ func TestCommunicatorRPC(t *testing.T) {
stdin_w.Close()
cmd.Wait()
if c.StartStdin != "info\n" {
t.Fatalf("bad data: %s", data)
t.Fatalf("bad data: %s", c.StartStdin)
}
// Test that we can get the exit status properly

View File

@ -266,7 +266,7 @@ func (m *MuxConn) loop() {
return
}
log.Printf("[DEBUG] Stream %d received packet %d", id, packetType)
//log.Printf("[DEBUG] Stream %d received packet %d", id, packetType)
switch packetType {
case muxPacketAck:
stream.mu.Lock()

View File

@ -38,7 +38,7 @@ func RegisterCommand(s *rpc.Server, c packer.Command) {
// Registers the appropriate endpoint on an RPC server to serve a
// Packer Communicator.
func RegisterCommunicator(s *rpc.Server, c packer.Communicator) {
registerComponent(s, "Communicator", &CommunicatorServer{c}, false)
registerComponent(s, "Communicator", &CommunicatorServer{c: c}, false)
}
// Registers the appropriate endpoint on an RPC server to serve a

View File

@ -14,6 +14,7 @@ var endpointId uint64
const (
DefaultArtifactEndpoint string = "Artifact"
DefaultCacheEndpoint = "Cache"
DefaultCommunicatorEndpoint = "Communicator"
DefaultPostProcessorEndpoint = "PostProcessor"
DefaultUiEndpoint = "Ui"
)
@ -55,6 +56,13 @@ func (s *Server) RegisterCache(c packer.Cache) {
})
}
func (s *Server) RegisterCommunicator(c packer.Communicator) {
s.server.RegisterName(DefaultCommunicatorEndpoint, &CommunicatorServer{
c: c,
mux: s.mux,
})
}
func (s *Server) RegisterPostProcessor(p packer.PostProcessor) {
s.server.RegisterName(DefaultPostProcessorEndpoint, &PostProcessorServer{
mux: s.mux,