Merge pull request #17646 from hashicorp/jbardin/GH-17638

use correct context for communicator.Retry

builtin/provisioners/chef/resource_provisioner.go

@@ -312,11 +312,11 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, comm.Timeout())
+	retryCtx, cancel := context.WithTimeout(ctx, comm.Timeout())
 	defer cancel()
 
 	// Wait and retry until we establish the connection
-	err = communicator.Retry(ctx, func() error {
+	err = communicator.Retry(retryCtx, func() error {
 		return comm.Connect(o)
 	})
 	if err != nil {

builtin/provisioners/file/resource_provisioner.go

@@ -48,9 +48,6 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, comm.Timeout())
-	defer cancel()
-
 	// Get the source
 	src, deleteSource, err := getSrc(data)
 	if err != nil {
@@ -99,8 +96,11 @@ func getSrc(data *schema.ResourceData) (string, bool, error) {
 
 // copyFiles is used to copy the files from a source to a destination
 func copyFiles(ctx context.Context, comm communicator.Communicator, src, dst string) error {
+	retryCtx, cancel := context.WithTimeout(ctx, comm.Timeout())
+	defer cancel()
+
 	// Wait and retry until we establish the connection
-	err := communicator.Retry(ctx, func() error {
+	err := communicator.Retry(retryCtx, func() error {
 		return comm.Connect(nil)
 	})
 	if err != nil {

builtin/provisioners/habitat/resource_provisioner.go

@@ -231,10 +231,10 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, comm.Timeout())
+	retryCtx, cancel := context.WithTimeout(ctx, comm.Timeout())
 	defer cancel()
 
-	err = communicator.Retry(ctx, func() error {
+	err = communicator.Retry(retryCtx, func() error {
 		return comm.Connect(o)
 	})
 

builtin/provisioners/remote-exec/resource_provisioner.go

@@ -157,20 +157,23 @@ func runScripts(
 	comm communicator.Communicator,
 	scripts []io.ReadCloser) error {
 
-	// Wait for the context to end and then disconnect
+	retryCtx, cancel := context.WithTimeout(ctx, comm.Timeout())
-	go func() {
+	defer cancel()
-		<-ctx.Done()
-		comm.Disconnect()
-	}()
 
 	// Wait and retry until we establish the connection
-	err := communicator.Retry(ctx, func() error {
+	err := communicator.Retry(retryCtx, func() error {
 		return comm.Connect(o)
 	})
 	if err != nil {
 		return err
 	}
 
+	// Wait for the context to end and then disconnect
+	go func() {
+		<-ctx.Done()
+		comm.Disconnect()
+	}()
+
 	for _, script := range scripts {
 		var cmd *remote.Cmd
 

builtin/provisioners/remote-exec/resource_provisioner_test.go

@@ -2,11 +2,15 @@ package remoteexec
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"testing"
+	"time"
 
 	"strings"
 
+	"github.com/hashicorp/terraform/communicator"
+	"github.com/hashicorp/terraform/communicator/remote"
 	"github.com/hashicorp/terraform/config"
 	"github.com/hashicorp/terraform/helper/schema"
 	"github.com/hashicorp/terraform/terraform"
@@ -206,6 +210,59 @@ func TestResourceProvider_CollectScripts_scriptsEmpty(t *testing.T) {
 	}
 }
 
+func TestProvisionerTimeout(t *testing.T) {
+	o := new(terraform.MockUIOutput)
+	c := new(communicator.MockCommunicator)
+
+	disconnected := make(chan struct{})
+	c.DisconnectFunc = func() error {
+		close(disconnected)
+		return nil
+	}
+
+	completed := make(chan struct{})
+	c.CommandFunc = func(cmd *remote.Cmd) error {
+		defer close(completed)
+		cmd.Init()
+		time.Sleep(2 * time.Second)
+		cmd.SetExitStatus(0, nil)
+		return nil
+	}
+	c.ConnTimeout = time.Second
+	c.UploadScripts = map[string]string{"hello": "echo hello"}
+	c.RemoteScriptPath = "hello"
+
+	p := Provisioner().(*schema.Provisioner)
+	conf := map[string]interface{}{
+		"inline": []interface{}{"echo hello"},
+	}
+
+	scripts, err := collectScripts(schema.TestResourceDataRaw(
+		t, p.Schema, conf))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		if err := runScripts(ctx, o, c, scripts); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	select {
+	case <-disconnected:
+		t.Fatal("communicator disconnected before command completed")
+	case <-completed:
+	}
+
+	<-done
+}
+
 func testConfig(t *testing.T, c map[string]interface{}) *terraform.ResourceConfig {
 	r, err := config.NewRawConfig(c)
 	if err != nil {

builtin/provisioners/salt-masterless/resource_provisioner.go

@@ -131,17 +131,11 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
-	ctx, cancelFunc := context.WithTimeout(ctx, comm.Timeout())
+	retryCtx, cancel := context.WithTimeout(ctx, comm.Timeout())
-	defer cancelFunc()
+	defer cancel()
-
-	// Wait for the context to end and then disconnect
-	go func() {
-		<-ctx.Done()
-		comm.Disconnect()
-	}()
 
 	// Wait and retry until we establish the connection
-	err = communicator.Retry(ctx, func() error {
+	err = communicator.Retry(retryCtx, func() error {
 		return comm.Connect(o)
 	})
 
@@ -149,6 +143,12 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
+	// Wait for the context to end and then disconnect
+	go func() {
+		<-ctx.Done()
+		comm.Disconnect()
+	}()
+
 	var src, dst string
 
 	o.Output("Provisioning with Salt...")

communicator/communicator_mock.go

@@ -18,6 +18,9 @@ type MockCommunicator struct {
 	Uploads          map[string]string
 	UploadScripts    map[string]string
 	UploadDirs       map[string]string
+	CommandFunc      func(*remote.Cmd) error
+	DisconnectFunc   func() error
+	ConnTimeout      time.Duration
 }
 
 // Connect implementation of communicator.Communicator interface
@@ -27,11 +30,17 @@ func (c *MockCommunicator) Connect(o terraform.UIOutput) error {
 
 // Disconnect implementation of communicator.Communicator interface
 func (c *MockCommunicator) Disconnect() error {
+	if c.DisconnectFunc != nil {
+		return c.DisconnectFunc()
+	}
 	return nil
 }
 
 // Timeout implementation of communicator.Communicator interface
 func (c *MockCommunicator) Timeout() time.Duration {
+	if c.ConnTimeout != 0 {
+		return c.ConnTimeout
+	}
 	return time.Duration(5 * time.Second)
 }
 
@@ -44,6 +53,10 @@ func (c *MockCommunicator) ScriptPath() string {
 func (c *MockCommunicator) Start(r *remote.Cmd) error {
 	r.Init()
 
+	if c.CommandFunc != nil {
+		return c.CommandFunc(r)
+	}
+
 	if !c.Commands[r.Command] {
 		return fmt.Errorf("Command not found!")
 	}