From 0345d960b2068cdddc6546df2b06a9b22411b8b4 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Thu, 15 Feb 2018 15:01:55 -0500
Subject: [PATCH] simplify remote-exec runScripts

There no reason to retry around the execution of remote scripts. We've
already established a connection, so the only that could happen here is
to continually retry uploading or executing a script that can't succeed.

This also simplifies the streaming output from the command, which
doesn't need such explicit synchronization. Closing the output pipes is
sufficient to stop the copyOutput functions, and they don't close around
any values that are accessed again after the command executes.
---
 .../provisioners/file/resource_provisioner.go | 11 +++-
 .../remote-exec/resource_provisioner.go       | 63 ++++++-------------
 2 files changed, 28 insertions(+), 46 deletions(-)

diff --git a/builtin/provisioners/file/resource_provisioner.go b/builtin/provisioners/file/resource_provisioner.go
index 30ed5e359..5514250d7 100644
--- a/builtin/provisioners/file/resource_provisioner.go
+++ b/builtin/provisioners/file/resource_provisioner.go
@@ -101,13 +101,18 @@ func getSrc(data *schema.ResourceData) (string, bool, error) {
 func copyFiles(ctx context.Context, comm communicator.Communicator, src, dst string) error {
 	// Wait and retry until we establish the connection
 	err := communicator.Retry(ctx, func() error {
-		err := comm.Connect(nil)
-		return err
+		return comm.Connect(nil)
 	})
 	if err != nil {
 		return err
 	}
-	defer comm.Disconnect()
+
+	// disconnect when the context is canceled, which will close this after
+	// Apply as well.
+	go func() {
+		<-ctx.Done()
+		comm.Disconnect()
+	}()
 
 	info, err := os.Stat(src)
 	if err != nil {
diff --git a/builtin/provisioners/remote-exec/resource_provisioner.go b/builtin/provisioners/remote-exec/resource_provisioner.go
index 082675ce1..378a282ed 100644
--- a/builtin/provisioners/remote-exec/resource_provisioner.go
+++ b/builtin/provisioners/remote-exec/resource_provisioner.go
@@ -171,57 +171,40 @@ func runScripts(
 	err := communicator.Retry(ctx, func() error {
 		return comm.Connect(o)
 	})
-
 	if err != nil {
 		return err
 	}
 
 	for _, script := range scripts {
 		var cmd *remote.Cmd
+
 		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)
+		defer outW.Close()
+		defer errW.Close()
+
+		go copyOutput(o, outR)
+		go copyOutput(o, errR)
 
 		remotePath := comm.ScriptPath()
 
-		err = communicator.Retry(ctx, func() error {
-			if err := comm.UploadScript(remotePath, script); err != nil {
-				return fmt.Errorf("Failed to upload script: %v", err)
-			}
-
-			cmd = &remote.Cmd{
-				Command: remotePath,
-				Stdout:  outW,
-				Stderr:  errW,
-			}
-			if err := comm.Start(cmd); err != nil {
-				return fmt.Errorf("Error starting script: %v", err)
-			}
-
-			return nil
-		})
-		if err == nil {
-			cmd.Wait()
-			if cmd.ExitStatus != 0 {
-				err = fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
-			}
+		if err := comm.UploadScript(remotePath, script); err != nil {
+			return fmt.Errorf("Failed to upload script: %v", err)
 		}
 
-		// If we have an error, end our context so the disconnect happens.
-		// This has to happen before the output cleanup below since during
-		// an interrupt this will cause the outputs to end.
-		if err != nil {
-			cancelFunc()
+		cmd = &remote.Cmd{
+			Command: remotePath,
+			Stdout:  outW,
+			Stderr:  errW,
+		}
+		if err := comm.Start(cmd); err != nil {
+			return fmt.Errorf("Error starting script: %v", err)
 		}
 
-		// Wait for output to clean up
-		outW.Close()
-		errW.Close()
-		<-outDoneCh
-		<-errDoneCh
+		cmd.Wait()
+		if cmd.ExitStatus != 0 {
+			err = fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
+		}
 
 		// Upload a blank follow up file in the same path to prevent residual
 		// script contents from remaining on remote machine
@@ -230,19 +213,13 @@ func runScripts(
 			// This feature is best-effort.
 			log.Printf("[WARN] Failed to upload empty follow up script: %v", err)
 		}
-
-		// If we have an error, return it out now that we've cleaned up
-		if err != nil {
-			return err
-		}
 	}
 
 	return nil
 }
 
 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)