Merge pull request #17609 from hashicorp/jbardin/remote-command

clean up remote.Cmd api
hashicorp · Mar 23, 2018 · e9e4ee4 · e9e4ee4
2 parents fe84adb + ad8642e
commit e9e4ee4
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 128 deletions.
diff --git a/builtin/provisioners/chef/resource_provisioner.go b/builtin/provisioners/chef/resource_provisioner.go
@@ -686,17 +686,10 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 
 	outR, outW := io.Pipe()
 	errR, errW := io.Pipe()
-	outDoneCh := make(chan struct{})
-	errDoneCh := make(chan struct{})
-	go p.copyOutput(o, outR, outDoneCh)
-	go p.copyOutput(o, errR, errDoneCh)
-	go func() {
-		// Wait for output to clean up
-		outW.Close()
-		errW.Close()
-		<-outDoneCh
-		<-errDoneCh
-	}()
+	go p.copyOutput(o, outR)
+	go p.copyOutput(o, errR)
+	defer outW.Close()
+	defer errW.Close()
 
 	cmd := &remote.Cmd{
 		Command: command,
@@ -709,20 +702,14 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 		return fmt.Errorf("Error executing command %q: %v", cmd.Command, err)
 	}
 
-	cmd.Wait()
-	if cmd.Err() != nil {
-		return cmd.Err()
-	}
-
-	if cmd.ExitStatus() != 0 {
-		return fmt.Errorf("Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus())
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
 
 	return nil
 }
 
-func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader, doneCh chan<- struct{}) {
-	defer close(doneCh)
+func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader) {
 	lr := linereader.New(r)
 	for line := range lr.Ch {
 		o.Output(line)

diff --git a/builtin/provisioners/habitat/resource_provisioner.go b/builtin/provisioners/habitat/resource_provisioner.go
@@ -729,8 +729,7 @@ func (p *provisioner) uploadUserTOML(o terraform.UIOutput, comm communicator.Com
 
 }
 
-func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader, doneCh chan<- struct{}) {
-	defer close(doneCh)
+func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader) {
 	lr := linereader.New(r)
 	for line := range lr.Ch {
 		o.Output(line)
@@ -741,16 +740,10 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 	outR, outW := io.Pipe()
 	errR, errW := io.Pipe()
 
-	outDoneCh := make(chan struct{})
-	errDoneCh := make(chan struct{})
-	go p.copyOutput(o, outR, outDoneCh)
-	go p.copyOutput(o, errR, errDoneCh)
-	defer func() {
-		outW.Close()
-		errW.Close()
-		<-outDoneCh
-		<-errDoneCh
-	}()
+	go p.copyOutput(o, outR)
+	go p.copyOutput(o, errR)
+	defer outW.Close()
+	defer errW.Close()
 
 	cmd := &remote.Cmd{
 		Command: command,
@@ -762,13 +755,8 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 		return fmt.Errorf("Error executing command %q: %v", cmd.Command, err)
 	}
 
-	cmd.Wait()
-	if cmd.Err() != nil {
-		return cmd.Err()
-	}
-
-	if cmd.ExitStatus() != 0 {
-		return fmt.Errorf("Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus())
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
 
 	return nil

diff --git a/builtin/provisioners/remote-exec/resource_provisioner.go b/builtin/provisioners/remote-exec/resource_provisioner.go
@@ -199,14 +199,9 @@ func runScripts(
 		if err := comm.Start(cmd); err != nil {
 			return fmt.Errorf("Error starting script: %v", err)
 		}
-		cmd.Wait()
 
-		if err := cmd.Err(); err != nil {
-			return fmt.Errorf("Remote command exited with error: %s", err)
-		}
-
-		if cmd.ExitStatus() != 0 {
-			err = fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus())
+		if err := cmd.Wait(); err != nil {
+			return err
 		}
 
 		// Upload a blank follow up file in the same path to prevent residual

diff --git a/builtin/provisioners/salt-masterless/resource_provisioner.go b/builtin/provisioners/salt-masterless/resource_provisioner.go
@@ -114,8 +114,6 @@ func Provisioner() terraform.ResourceProvisioner {
 // Apply executes the file provisioner
 func applyFn(ctx context.Context) error {
 	// Decode the raw config for this provisioner
-	var err error
-
 	o := ctx.Value(schema.ProvOutputKey).(terraform.UIOutput)
 	d := ctx.Value(schema.ProvConfigDataKey).(*schema.ResourceData)
 	connState := ctx.Value(schema.ProvRawStateKey).(*terraform.InstanceState)
@@ -162,46 +160,29 @@ func applyFn(ctx context.Context) error {
 			err = fmt.Errorf("Unable to download Salt: %s", err)
 		}
 
-		if err == nil {
-			cmd.Wait()
-			if cmd.Err() != nil {
-				err = cmd.Err()
-			} else if cmd.ExitStatus() != 0 {
-				err = fmt.Errorf("Curl exited with non-zero exit status: %d", cmd.ExitStatus())
-			}
+		if err := cmd.Wait(); err != nil {
+			return err
 		}
 
 		outR, outW := io.Pipe()
 		errR, errW := io.Pipe()
-		outDoneCh := make(chan struct{})
-		errDoneCh := make(chan struct{})
-		go copyOutput(o, outR, outDoneCh)
-		go copyOutput(o, errR, errDoneCh)
+		go copyOutput(o, outR)
+		go copyOutput(o, errR)
+		defer outW.Close()
+		defer errW.Close()
+
 		cmd = &remote.Cmd{
 			Command: fmt.Sprintf("%s /tmp/install_salt.sh %s", p.sudo("sh"), p.BootstrapArgs),
 			Stdout:  outW,
 			Stderr:  errW,
 		}
 
 		o.Output(fmt.Sprintf("Installing Salt with command %s", cmd.Command))
-		if err = comm.Start(cmd); err != nil {
-			err = fmt.Errorf("Unable to install Salt: %s", err)
+		if err := comm.Start(cmd); err != nil {
+			return fmt.Errorf("Unable to install Salt: %s", err)
 		}
 
-		if err == nil {
-			cmd.Wait()
-			if cmd.Err() != nil {
-				err = cmd.Err()
-			} else if cmd.ExitStatus() != 0 {
-				err = fmt.Errorf("install_salt.sh exited with non-zero exit status: %d", cmd.ExitStatus())
-			}
-		}
-		// Wait for output to clean up
-		outW.Close()
-		errW.Close()
-		<-outDoneCh
-		<-errDoneCh
-		if err != nil {
+		if err := cmd.Wait(); err != nil {
 			return err
 		}
 	}
@@ -270,11 +251,11 @@ func applyFn(ctx context.Context) error {
 
 	outR, outW := io.Pipe()
 	errR, errW := io.Pipe()
-	outDoneCh := make(chan struct{})
-	errDoneCh := make(chan struct{})
+	go copyOutput(o, outR)
+	go copyOutput(o, errR)
+	defer outW.Close()
+	defer errW.Close()
 
-	go copyOutput(o, outR, outDoneCh)
-	go copyOutput(o, errR, errDoneCh)
 	o.Output(fmt.Sprintf("Running: salt-call --local %s", p.CmdArgs))
 	cmd := &remote.Cmd{
 		Command: p.sudo(fmt.Sprintf("salt-call --local %s", p.CmdArgs)),
@@ -285,21 +266,10 @@ func applyFn(ctx context.Context) error {
 		err = fmt.Errorf("Error executing salt-call: %s", err)
 	}
 
-	if err == nil {
-		cmd.Wait()
-		if cmd.Err() != nil {
-			err = cmd.Err()
-		} else if cmd.ExitStatus() != 0 {
-			err = fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus())
-		}
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
-	// Wait for output to clean up
-	outW.Close()
-	errW.Close()
-	<-outDoneCh
-	<-errDoneCh
-
-	return err
+	return nil
 }
 
 // Prepends sudo to supplied command if config says to
@@ -360,12 +330,10 @@ func (p *provisioner) moveFile(o terraform.UIOutput, comm communicator.Communica
 	if err := comm.Start(cmd); err != nil {
 		return fmt.Errorf("Unable to move %s to %s: %s", src, dst, err)
 	}
-	cmd.Wait()
-	if cmd.ExitStatus() != 0 {
-		return fmt.Errorf("Unable to move %s to %s: exit status: %d", src, dst, cmd.ExitStatus())
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
-
-	return cmd.Err()
+	return nil
 }
 
 func (p *provisioner) createDir(o terraform.UIOutput, comm communicator.Communicator, dir string) error {
@@ -377,11 +345,10 @@ func (p *provisioner) createDir(o terraform.UIOutput, comm communicator.Communic
 		return err
 	}
 
-	cmd.Wait()
-	if cmd.ExitStatus() != 0 {
-		return fmt.Errorf("Non-zero exit status.")
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
-	return cmd.Err()
+	return nil
 }
 
 func (p *provisioner) removeDir(o terraform.UIOutput, comm communicator.Communicator, dir string) error {
@@ -392,11 +359,10 @@ func (p *provisioner) removeDir(o terraform.UIOutput, comm communicator.Communic
 	if err := comm.Start(cmd); err != nil {
 		return err
 	}
-	cmd.Wait()
-	if cmd.ExitStatus() != 0 {
-		return fmt.Errorf("Non-zero exit status.")
+	if err := cmd.Wait(); err != nil {
+		return err
 	}
-	return cmd.Err()
+	return nil
 }
 
 func (p *provisioner) uploadDir(o terraform.UIOutput, comm communicator.Communicator, dst, src string, ignore []string) error {
@@ -549,8 +515,7 @@ func decodeConfig(d *schema.ResourceData) (*provisioner, error) {
 }
 
 func copyOutput(
-	o terraform.UIOutput, r io.Reader, doneCh chan<- struct{}) {
-	defer close(doneCh)
+	o terraform.UIOutput, r io.Reader) {
 	lr := linereader.New(r)
 	for line := range lr.Ch {
 		o.Output(line)

diff --git a/communicator/remote/command.go b/communicator/remote/command.go
@@ -1,6 +1,7 @@
 package remote
 
 import (
+	"fmt"
 	"io"
 	"sync"
 )
@@ -59,23 +60,37 @@ func (c *Cmd) SetExitStatus(status int, err error) {
 	close(c.exitCh)
 }
 
-// Err returns any communicator related error.
-func (c *Cmd) Err() error {
+// Wait waits for the remote command to complete.
+// Wait may return an error from the communicator, or an ExitError if the
+// process exits with a non-zero exit status.
+func (c *Cmd) Wait() error {
+	<-c.exitCh
+
 	c.Lock()
 	defer c.Unlock()
 
-	return c.err
-}
+	if c.err != nil || c.exitStatus != 0 {
+		return &ExitError{
+			Command:    c.Command,
+			ExitStatus: c.exitStatus,
+			Err:        c.err,
+		}
+	}
 
-// ExitStatus returns the exit status of the remote command
-func (c *Cmd) ExitStatus() int {
-	c.Lock()
-	defer c.Unlock()
+	return nil
+}
 
-	return c.exitStatus
+// ExitError is returned by Wait to indicate and error executing the remote
+// command, or a non-zero exit status.
+type ExitError struct {
+	Command    string
+	ExitStatus int
+	Err        error
 }
 
-// Wait waits for the remote command to complete.
-func (c *Cmd) Wait() {
-	<-c.exitCh
+func (e *ExitError) Error() string {
+	if e.Err != nil {
+		return fmt.Sprintf("error executing %q: %v", e.Command, e.Err)
+	}
+	return fmt.Sprintf("%q exit status: %d", e.Command, e.ExitStatus)
 }
diff --git a/communicator/ssh/communicator.go b/communicator/ssh/communicator.go
@@ -359,11 +359,11 @@ func (c *Communicator) UploadScript(path string, input io.Reader) error {
 			"Error chmodding script file to 0777 in remote "+
 				"machine: %s", err)
 	}
-	cmd.Wait()
-	if cmd.ExitStatus() != 0 {
+
+	if err := cmd.Wait(); err != nil {
 		return fmt.Errorf(
 			"Error chmodding script file to 0777 in remote "+
-				"machine %d: %s %s", cmd.ExitStatus(), stdout.String(), stderr.String())
+				"machine %v: %s %s", err, stdout.String(), stderr.String())
 	}
 
 	return nil

diff --git a/communicator/ssh/communicator_test.go b/communicator/ssh/communicator_test.go
@@ -219,13 +219,10 @@ func TestLostConnection(t *testing.T) {
 		c.Disconnect()
 	}()
 
-	cmd.Wait()
-	if cmd.Err() == nil {
+	err = cmd.Wait()
+	if err == nil {
 		t.Fatal("expected communicator error")
 	}
-	if cmd.ExitStatus() != 0 {
-		t.Fatal("command should not have returned an exit status")
-	}
 }
 
 func TestHostKey(t *testing.T) {