packer/rpc: implement Communicator
This commit is contained in:
Mitchell Hashimoto
parent
72fcb566a6
commit
db06fc7501
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
c packer.Communicator
|
||||
mux *MuxConn
|
||||
}
|
||||
|
||||
type CommandFinished struct {
|
||||
|
|
@ -27,21 +27,21 @@ type CommandFinished struct {
|
|||
}
|
||||
|
||||
type CommunicatorStartArgs struct {
|
||||
Command string
|
||||
StdinAddress string
|
||||
StdoutAddress string
|
||||
StderrAddress string
|
||||
ResponseAddress string
|
||||
Command string
|
||||
StdinStreamId uint32
|
||||
StdoutStreamId uint32
|
||||
StderrStreamId uint32
|
||||
ResponseStreamId uint32
|
||||
}
|
||||
|
||||
type CommunicatorDownloadArgs struct {
|
||||
Path string
|
||||
WriterAddress string
|
||||
Path string
|
||||
WriterStreamId uint32
|
||||
}
|
||||
|
||||
type CommunicatorUploadArgs struct {
|
||||
Path string
|
||||
ReaderAddress string
|
||||
Path 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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue