use correct context for communicator.Retry

The timeout for a provisioner is expected to only apply to the initial
connection. Keep the context for the communicator.Retry separate from
the global cancellation context.

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
-	go func() {
-		<-ctx.Done()
-		comm.Disconnect()
-	}()
+	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 {
 		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/salt-masterless/resource_provisioner.go

@@ -131,17 +131,11 @@ func applyFn(ctx context.Context) error {
 		return err
 	}
 
-	ctx, cancelFunc := context.WithTimeout(ctx, comm.Timeout())
-	defer cancelFunc()
-
-	// Wait for the context to end and then disconnect
-	go func() {
-		<-ctx.Done()
-		comm.Disconnect()
-	}()
+	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)
 	})
 
@@ -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...")