packer: More efficient RemoteCommand.ExitChan

communicator/ssh/communicator.go

@@ -41,10 +41,10 @@ func (c *comm) Start(cmd string) (remote *packer.RemoteCommand, err error) {
 
 	// Setup the remote command
 	remote = &packer.RemoteCommand{
-		Stdin: stdin,
-		Stdout: stdout,
-		Stderr: stderr,
-		Exited: false,
+		Stdin:      stdin,
+		Stdout:     stdout,
+		Stderr:     stderr,
+		Exited:     false,
 		ExitStatus: -1,
 	}
 
@@ -153,7 +153,6 @@ func (c *comm) Upload(path string, input io.Reader) error {
 		return err
 	}
 
-
 	log.Printf("scp stdout (length %d): %#v", stdout.Len(), stdout.Bytes())
 	log.Printf("scp stderr (length %d): %s", stderr.Len(), stderr.String())
 

packer/communicator.go

@@ -2,6 +2,7 @@ package packer
 
 import (
 	"io"
+	"log"
 	"sync"
 	"time"
 )
@@ -36,34 +37,55 @@ type RemoteCommand struct {
 	Exited     bool
 	ExitStatus int
 
+	exitChans    []chan<- int
 	exitChanLock sync.Mutex
 }
 
 // StdoutStream returns a channel that will be sent all the output
 // of stdout as it comes. The output isn't guaranteed to be a full line.
 // When the channel is closed, the process is exited.
-func (r *RemoteCommand) StdoutChan() (<-chan string) {
+func (r *RemoteCommand) StdoutChan() <-chan string {
 	return nil
 }
 
 // ExitChan returns a channel that will be sent the exit status once
 // the process exits. This can be used in cases such a select statement
 // waiting on the process to end.
-func (r *RemoteCommand) ExitChan() (<-chan int) {
-	// TODO(mitchellh): Something more efficient than multiple Wait() calls
-
+func (r *RemoteCommand) ExitChan() <-chan int {
 	r.exitChanLock.Lock()
 	defer r.exitChanLock.Unlock()
 
-	// Make a single buffered channel so that the send doesn't block.
+	// If we haven't made any channels yet, make that slice
+	if r.exitChans == nil {
+		r.exitChans = make([]chan<- int, 0, 5)
+
+		go func() {
+			// Wait for the command to finish
+			r.Wait()
+
+			// Grab the exit chan lock so we can iterate over it and
+			// message to each channel.
+			r.exitChanLock.Lock()
+			defer r.exitChanLock.Unlock()
+
+			for _, ch := range r.exitChans {
+				// Use a select so the send never blocks
+				select {
+				case ch <- r.ExitStatus:
+				default:
+					log.Println("remote command exit channel wouldn't blocked. Weird.")
+				}
+
+				close(ch)
+			}
+
+			r.exitChans = nil
+		}()
+	}
+
+	// Append our new channel onto it and return it
 	exitChan := make(chan int, 1)
-
-	go func() {
-		defer close(exitChan)
-		r.Wait()
-		exitChan <- r.ExitStatus
-	}()
-
+	r.exitChans = append(r.exitChans, exitChan)
 	return exitChan
 }
 

packer/rpc/communicator.go

@@ -80,10 +80,10 @@ func (c *communicator) Start(cmd string) (rc *packer.RemoteCommand, err error) {
 
 	// Build the response object using the streams we created
 	rc = &packer.RemoteCommand{
-		Stdin: stdinC,
-		Stdout: stdoutC,
-		Stderr: stderrC,
-		Exited: false,
+		Stdin:      stdinC,
+		Stdout:     stdoutC,
+		Stderr:     stderrC,
+		Exited:     false,
 		ExitStatus: -1,
 	}
 

packer/rpc/communicator_test.go

@@ -39,10 +39,10 @@ func (t *testCommunicator) Start(cmd string) (*packer.RemoteCommand, error) {
 	stderr, t.startErr = io.Pipe()
 
 	rc := &packer.RemoteCommand{
-		Stdin: stdin,
-		Stdout: stdout,
-		Stderr: stderr,
-		Exited: false,
+		Stdin:      stdin,
+		Stdout:     stdout,
+		Stderr:     stderr,
+		Exited:     false,
 		ExitStatus: 0,
 	}
 

packer/rpc/ui_test.go

@@ -7,9 +7,9 @@ import (
 )
 
 type testUi struct {
-	errorCalled bool
+	errorCalled  bool
 	errorMessage string
-	sayCalled   bool
+	sayCalled    bool
 	sayMessage   string
 }
 

plugin/provisioner-shell/main.go

@@ -1,8 +1,8 @@
 package main
 
 import (
-	"github.com/mitchellh/packer/provisioner/shell"
 	"github.com/mitchellh/packer/packer/plugin"
+	"github.com/mitchellh/packer/provisioner/shell"
 )
 
 func main() {

provisioner/shell/provisioner_test.go

@@ -14,7 +14,7 @@ func TestProvisioner_Impl(t *testing.T) {
 }
 
 func TestProvisionerPrepare_Defaults(t *testing.T) {
-	raw := map[string]interface{} {}
+	raw := map[string]interface{}{}
 
 	p := &Provisioner{}
 	p.Prepare(raw, nil)