packer: Waiting for a remote command and setting the exit status

packer/communicator.go

 package packer
 
-import "io"
+import (
+	"io"
+	"time"
+)
 
 // A Communicator is the interface used to communicate with the machine
 // that exists that will eventually be packaged into an image. Communicators
@@ -16,9 +19,28 @@ type Communicator interface {
 
 // This struct contains some information about the remote command being
 // executed and can be used to wait for it to complete.
+//
+// Stdin, Stdout, Stderr are readers and writers to varios IO streams for
+// the remote command.
+//
+// Exited is false until Wait is called. It can be used to check if Wait
+// has already been called.
+//
+// ExitStatus is the exit code of the remote process. It is only available
+// once Wait is called.
 type RemoteCommand struct {
 	Stdin  io.Writer
 	Stdout io.Reader
 	Stderr io.Reader
+	Exited bool
 	ExitStatus int
 }
+
+// Wait waits for the command to exit.
+func (r *RemoteCommand) Wait() {
+	// Busy wait on being exited. We put a sleep to be kind to the
+	// Go scheduler, and because we don't really need smaller granularity.
+	for !r.Exited {
+		time.Sleep(10 * time.Millisecond)
+	}
+}

packer/rpc/communicator.go

@@ -31,7 +31,7 @@ type CommunicatorStartResponse struct {
 	StdinAddress string
 	StdoutAddress string
 	StderrAddress string
-	ExitStatusAddress string
+	RemoteCommandAddress string
 }
 
 func Communicator(client *rpc.Client) *communicator {
@@ -62,14 +62,28 @@ func (c *communicator) Start(cmd string) (rc *packer.RemoteCommand, err error) {
 		return
 	}
 
+	// Connect to the RPC server for the remote command
+	client, err := rpc.Dial("tcp", response.RemoteCommandAddress)
+	if err != nil {
+		return
+	}
+
 	// Build the response object using the streams we created
 	rc = &packer.RemoteCommand{
 		stdinC,
 		stdoutC,
 		stderrC,
-		0,
+		false,
+		-1,
 	}
 
+	// In a goroutine, we wait for the process to exit, then we set
+	// that it has exited.
+	go func() {
+		client.Call("RemoteCommand.Wait", new(interface{}), &rc.ExitStatus)
+		rc.Exited = true
+	}()
+
 	return
 }
 
@@ -99,7 +113,7 @@ func (c *CommunicatorServer) Start(cmd *string, reply *CommunicatorStartResponse
 	// For the exit status, we use a simple RPC Server that serves
 	// some of the RemoteComand methods.
 	server := rpc.NewServer()
-	//server.RegisterName("RemoteCommand", &RemoteCommandServer{command})
+	server.RegisterName("RemoteCommand", &RemoteCommandServer{command})
 
 	*reply = CommunicatorStartResponse{
 		stdinL.Addr().String(),
@@ -111,6 +125,12 @@ func (c *CommunicatorServer) Start(cmd *string, reply *CommunicatorStartResponse
 	return
 }
 
+func (rc *RemoteCommandServer) Wait(args *interface{}, reply *int) error {
+	rc.rc.Wait()
+	*reply = rc.rc.ExitStatus
+	return nil
+}
+
 func serveSingleCopy(name string, l net.Listener, dst io.Writer, src io.Reader) {
 	defer l.Close()
 

packer/rpc/communicator_test.go

@@ -16,6 +16,8 @@ type testCommunicator struct {
 	startIn *io.PipeReader
 	startOut *io.PipeWriter
 	startErr *io.PipeWriter
+	startExited *bool
+	startExitStatus *int
 }
 
 func (t *testCommunicator) Start(cmd string) (*packer.RemoteCommand, error) {
@@ -33,9 +35,13 @@ func (t *testCommunicator) Start(cmd string) (*packer.RemoteCommand, error) {
 		stdin,
 		stdout,
 		stderr,
+		false,
 		0,
 	}
 
+	t.startExited = &rc.Exited
+	t.startExitStatus = &rc.ExitStatus
+
 	return rc, nil
 }
 
@@ -83,6 +89,12 @@ func TestCommunicatorRPC(t *testing.T) {
 	data, err = bufIn.ReadString('\n')
 	assert.Nil(err, "should have no problem reading stdin")
 	assert.Equal(data, "infoo\n", "should be correct stdin")
+
+	// Test that we can get the exit status properly
+	*c.startExitStatus = 42
+	*c.startExited = true
+	rc.Wait()
+	assert.Equal(rc.ExitStatus, 42, "should have proper exit status")
 }
 
 func TestCommunicator_ImplementsCommunicator(t *testing.T) {