From 142df657c3597d1d1868ce840b1402a22a912def Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 26 Dec 2016 16:29:44 -0800
Subject: [PATCH] provisioners/remote-exec: listen to Stop

---
 .../remote-exec/resource_provisioner.go       | 86 +++++++++++++++----
 terraform/context.go                          |  7 ++
 2 files changed, 77 insertions(+), 16 deletions(-)

diff --git a/builtin/provisioners/remote-exec/resource_provisioner.go b/builtin/provisioners/remote-exec/resource_provisioner.go
index cda3c0b69..042e8544e 100644
--- a/builtin/provisioners/remote-exec/resource_provisioner.go
+++ b/builtin/provisioners/remote-exec/resource_provisioner.go
@@ -8,6 +8,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/terraform/communicator"
@@ -68,7 +69,7 @@ func applyFn(ctx context.Context) error {
 	}
 
 	// Copy and execute each script
-	if err := runScripts(o, comm, scripts); err != nil {
+	if err := runScripts(ctx, o, comm, scripts); err != nil {
 		return err
 	}
 
@@ -133,18 +134,29 @@ func collectScripts(d *schema.ResourceData) ([]io.ReadCloser, error) {
 
 // runScripts is used to copy and execute a set of scripts
 func runScripts(
+	ctx context.Context,
 	o terraform.UIOutput,
 	comm communicator.Communicator,
 	scripts []io.ReadCloser) error {
+	// Wrap out context in a cancelation function that we use to
+	// kill the connection.
+	ctx, cancelFunc := context.WithCancel(ctx)
+	defer cancelFunc()
+
+	// Wait for the context to end and then disconnect
+	go func() {
+		<-ctx.Done()
+		comm.Disconnect()
+	}()
+
 	// Wait and retry until we establish the connection
-	err := retryFunc(comm.Timeout(), func() error {
+	err := retryFunc(ctx, comm.Timeout(), func() error {
 		err := comm.Connect(o)
 		return err
 	})
 	if err != nil {
 		return err
 	}
-	defer comm.Disconnect()
 
 	for _, script := range scripts {
 		var cmd *remote.Cmd
@@ -156,7 +168,7 @@ func runScripts(
 		go copyOutput(o, errR, errDoneCh)
 
 		remotePath := comm.ScriptPath()
-		err = retryFunc(comm.Timeout(), func() error {
+		err = retryFunc(ctx, comm.Timeout(), func() error {
 			if err := comm.UploadScript(remotePath, script); err != nil {
 				return fmt.Errorf("Failed to upload script: %v", err)
 			}
@@ -179,6 +191,13 @@ func runScripts(
 			}
 		}
 
+		// 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()
+		}
+
 		// Wait for output to clean up
 		outW.Close()
 		errW.Close()
@@ -212,19 +231,54 @@ func copyOutput(
 }
 
 // retryFunc is used to retry a function for a given duration
-func retryFunc(timeout time.Duration, f func() error) error {
-	finish := time.After(timeout)
-	for {
-		err := f()
-		if err == nil {
-			return nil
-		}
-		log.Printf("Retryable error: %v", err)
+func retryFunc(ctx context.Context, timeout time.Duration, f func() error) error {
+	// Build a new context with the timeout
+	ctx, done := context.WithTimeout(ctx, timeout)
+	defer done()
 
-		select {
-		case <-finish:
-			return err
-		case <-time.After(3 * time.Second):
+	// Try the function in a goroutine
+	var errVal atomic.Value
+	doneCh := make(chan struct{})
+	go func() {
+		defer close(doneCh)
+
+		for {
+			// If our context ended, we want to exit right away.
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+			// Try the function call
+			err := f()
+			if err == nil {
+				return
+			}
+
+			log.Printf("Retryable error: %v", err)
+			errVal.Store(err)
 		}
+	}()
+
+	// Wait for completion
+	select {
+	case <-doneCh:
+	case <-ctx.Done():
 	}
+
+	// Check if we have a context error to check if we're interrupted or timeout
+	switch ctx.Err() {
+	case context.Canceled:
+		return fmt.Errorf("interrupted")
+	case context.DeadlineExceeded:
+		return fmt.Errorf("timeout")
+	}
+
+	// Check if we got an error executing
+	if err, ok := errVal.Load().(error); ok {
+		return err
+	}
+
+	return nil
 }
diff --git a/terraform/context.go b/terraform/context.go
index 52b95f058..82e517a86 100644
--- a/terraform/context.go
+++ b/terraform/context.go
@@ -632,10 +632,14 @@ func (c *Context) Refresh() (*State, error) {
 //
 // Stop will block until the task completes.
 func (c *Context) Stop() {
+	log.Printf("[WARN] terraform: Stop called, initiating interrupt sequence")
+
 	c.l.Lock()
 
 	// If we're running, then stop
 	if c.runContextCancel != nil {
+		log.Printf("[WARN] terraform: run context exists, stopping")
+
 		// Tell the hook we want to stop
 		c.sh.Stop()
 
@@ -652,8 +656,11 @@ func (c *Context) Stop() {
 
 	// Wait if we have a context
 	if ctx != nil {
+		log.Printf("[WARN] terraform: stop waiting for context completion")
 		<-ctx.Done()
 	}
+
+	log.Printf("[WARN] terraform: stop complete")
 }
 
 // Validate validates the configuration and returns any warnings or errors.