Support graceful cancellation

CONTRIBUTING.md

@@ -28,7 +28,7 @@ func (tf *Terraform) Refresh(ctx context.Context, opts ...RefreshCmdOption) erro
 
 func (tf *Terraform) refreshCmd(ctx context.Context, opts ...RefreshCmdOption) (*exec.Cmd, error) {
 	...
-  	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
+  	return tf.buildTerraformCmd( mergeEnv, args...), nil
 }
 ```
 

tfexec/apply.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os/exec"
 	"strconv"
+	"time"
 )
 
 type applyConfig struct {
@@ -13,24 +14,26 @@ type applyConfig struct {
 	lock      bool
 
 	// LockTimeout must be a string with time unit, e.g. '10s'
-	lockTimeout  string
-	parallelism  int
-	reattachInfo ReattachInfo
-	refresh      bool
-	replaceAddrs []string
-	state        string
-	stateOut     string
-	targets      []string
+	lockTimeout             string
+	parallelism             int
+	reattachInfo            ReattachInfo
+	refresh                 bool
+	replaceAddrs            []string
+	state                   string
+	stateOut                string
+	targets                 []string
+	gracefulShutdownTimeout time.Duration
 
 	// Vars: each var must be supplied as a single string, e.g. 'foo=bar'
 	vars     []string
 	varFiles []string
 }
 
 var defaultApplyOptions = applyConfig{
-	lock:        true,
-	parallelism: 10,
-	refresh:     true,
+	lock:                    true,
+	parallelism:             10,
+	refresh:                 true,
+	gracefulShutdownTimeout: 0,
 }
 
 // ApplyOption represents options used in the Apply method.
@@ -42,6 +45,10 @@ func (opt *ParallelismOption) configureApply(conf *applyConfig) {
 	conf.parallelism = opt.parallelism
 }
 
+func (opt *GracefulShutdownTimeoutOption) configureApply(conf *applyConfig) {
+	conf.gracefulShutdownTimeout = opt.timeout
+}
+
 func (opt *BackupOption) configureApply(conf *applyConfig) {
 	conf.backup = opt.path
 }
@@ -92,14 +99,15 @@ func (opt *ReattachOption) configureApply(conf *applyConfig) {
 
 // Apply represents the terraform apply subcommand.
 func (tf *Terraform) Apply(ctx context.Context, opts ...ApplyOption) error {
-	cmd, err := tf.applyCmd(ctx, opts...)
+	cmd, applyOpts, err := tf.applyCmd(ctx, opts...)
 	if err != nil {
 		return err
 	}
-	return tf.runTerraformCmd(ctx, cmd)
+
+	return tf.runTerraformCmdWithGracefulshutdownTimeout(ctx, cmd, applyOpts.gracefulShutdownTimeout)
 }
 
-func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.Cmd, error) {
+func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.Cmd, *applyConfig, error) {
 	c := defaultApplyOptions
 
 	for _, o := range opts {
@@ -134,7 +142,7 @@ func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.C
 	if c.replaceAddrs != nil {
 		err := tf.compatible(ctx, tf0_15_2, nil)
 		if err != nil {
-			return nil, fmt.Errorf("replace option was introduced in Terraform 0.15.2: %w", err)
+			return nil, nil, fmt.Errorf("replace option was introduced in Terraform 0.15.2: %w", err)
 		}
 		for _, addr := range c.replaceAddrs {
 			args = append(args, "-replace="+addr)
@@ -160,10 +168,10 @@ func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.C
 	if c.reattachInfo != nil {
 		reattachStr, err := c.reattachInfo.marshalString()
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
-	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
+	return tf.buildTerraformCmd(mergeEnv, args...), &c, nil
 }

tfexec/apply_test.go

@@ -3,6 +3,7 @@ package tfexec
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/terraform-exec/tfexec/internal/testutil"
 )
@@ -19,7 +20,7 @@ func TestApplyCmd(t *testing.T) {
 	tf.SetEnv(map[string]string{})
 
 	t.Run("basic", func(t *testing.T) {
-		applyCmd, err := tf.applyCmd(context.Background(),
+		applyCmd, applyOpts, err := tf.applyCmd(context.Background(),
 			Backup("testbackup"),
 			LockTimeout("200s"),
 			State("teststate"),
@@ -36,6 +37,7 @@ func TestApplyCmd(t *testing.T) {
 			Var("var1=foo"),
 			Var("var2=bar"),
 			DirOrPlan("testfile"),
+			GracefulShutdownTimeout(10*time.Second),
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -63,5 +65,10 @@ func TestApplyCmd(t *testing.T) {
 			"-var", "var2=bar",
 			"testfile",
 		}, nil, applyCmd)
+
+		if applyOpts.gracefulShutdownTimeout != 10*time.Second {
+			t.Fatalf("graceful shutdown timeout mismatch\n\nexpected:\n%v\n\ngot:\n%v", 10*time.Second, applyOpts.gracefulShutdownTimeout)
+		}
 	})
+
 }

tfexec/cmd.go

@@ -179,8 +179,8 @@ func (tf *Terraform) buildEnv(mergeEnv map[string]string) []string {
 	return envSlice(env)
 }
 
-func (tf *Terraform) buildTerraformCmd(ctx context.Context, mergeEnv map[string]string, args ...string) *exec.Cmd {
-	cmd := exec.CommandContext(ctx, tf.execPath, args...)
+func (tf *Terraform) buildTerraformCmd(mergeEnv map[string]string, args ...string) *exec.Cmd {
+	cmd := exec.Command(tf.execPath, args...)
 
 	cmd.Env = tf.buildEnv(mergeEnv)
 	cmd.Dir = tf.workingDir

tfexec/cmd_default.go

@@ -5,12 +5,21 @@ package tfexec
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"os/exec"
 	"strings"
 	"sync"
+	"time"
 )
 
-func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
+const defaultGracefulShutdownTimeout = 0
+
+func (tf *Terraform) runTerraformCmd(parentCtx context.Context, cmd *exec.Cmd) error {
+	return tf.runTerraformCmdWithGracefulshutdownTimeout(parentCtx, cmd, defaultGracefulShutdownTimeout)
+}
+
+func (tf *Terraform) runTerraformCmdWithGracefulshutdownTimeout(ctx context.Context, cmd *exec.Cmd, _ time.Duration) error {
 	var errBuf strings.Builder
 
 	// check for early cancellation
@@ -38,7 +47,19 @@ func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
 	if err != nil {
 		return err
 	}
-
+	cmdDoneCh := make(chan error, 1)
+	returnCh := make(chan error, 1)
+	defer close(cmdDoneCh)
+	go func() {
+		defer close(returnCh)
+		select {
+		case <-ctx.Done(): // wait for context cancelled
+			cmd.Process.Signal(os.Kill)
+			returnCh <- fmt.Errorf("%w: terraform forcefully killed", ctx.Err())
+		case err := <-cmdDoneCh:
+			returnCh <- err
+		}
+	}()
 	err = cmd.Start()
 	if err == nil && ctx.Err() != nil {
 		err = ctx.Err()
@@ -65,7 +86,9 @@ func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
 	// can cause a race condition
 	wg.Wait()
 
-	err = cmd.Wait()
+	cmdDoneCh <- cmd.Wait()
+	err = <-returnCh
+
 	if err == nil && ctx.Err() != nil {
 		err = ctx.Err()
 	}

tfexec/cmd_default_test.go

@@ -6,6 +6,7 @@ package tfexec
 import (
 	"bytes"
 	"context"
+	"errors"
 	"log"
 	"strings"
 	"testing"
@@ -24,7 +25,7 @@ func Test_runTerraformCmd_default(t *testing.T) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	cmd := tf.buildTerraformCmd(ctx, nil, "hello tf-exec!")
+	cmd := tf.buildTerraformCmd(nil, "hello tf-exec!")
 	err := tf.runTerraformCmd(ctx, cmd)
 	if err != nil {
 		t.Fatal(err)
@@ -37,3 +38,24 @@ func Test_runTerraformCmd_default(t *testing.T) {
 		t.Fatal("canceling context should not lead to logging an error")
 	}
 }
+
+func Test_runTerraformCmdCancelCtx_default(t *testing.T) {
+	var buf bytes.Buffer
+
+	tf := &Terraform{
+		logger:   log.New(&buf, "", 0),
+		execPath: "sleep",
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cmd := tf.buildTerraformCmd(nil, "3")
+	go func() {
+		<-time.After(1 * time.Second)
+		cancel()
+	}()
+	err := tf.runTerraformCmd(ctx, cmd)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context.Canceled, got %T %s", err, err)
+	}
+}

tfexec/cmd_linux.go

@@ -2,13 +2,22 @@ package tfexec
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"os/exec"
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 )
 
-func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
+const defaultGracefulShutdownTimeout = 0
+
+func (tf *Terraform) runTerraformCmd(parentCtx context.Context, cmd *exec.Cmd) error {
+	return tf.runTerraformCmdWithGracefulshutdownTimeout(parentCtx, cmd, defaultGracefulShutdownTimeout)
+}
+
+func (tf *Terraform) runTerraformCmdWithGracefulshutdownTimeout(parentCtx context.Context, cmd *exec.Cmd, gracefulShutdownTimeout time.Duration) error {
 	var errBuf strings.Builder
 
 	cmd.SysProcAttr = &syscall.SysProcAttr{
@@ -20,10 +29,14 @@ func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
 
 	// check for early cancellation
 	select {
-	case <-ctx.Done():
-		return ctx.Err()
+	case <-parentCtx.Done():
+		return parentCtx.Err()
 	default:
 	}
+	// Context for the stdout and stderr writers so that they are not closed on parentCtx cancellation
+	// and avoiding "broken pipe"
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	// Read stdout / stderr logs from pipe instead of setting cmd.Stdout and
 	// cmd.Stderr because it can cause hanging when killing the command
@@ -43,8 +56,51 @@ func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
 	if err != nil {
 		return err
 	}
-
-	err = cmd.Start()
+	cmdMu := sync.Mutex{}
+	withCmdLock := func(fn func() error) error {
+		cmdMu.Lock()
+		defer cmdMu.Unlock()
+		return fn()
+	}
+	cmdDoneCh := make(chan error, 1)
+	returnCh := make(chan error, 1)
+	defer close(returnCh)
+	go func() {
+		select {
+		case <-parentCtx.Done(): // wait for context cancelled
+			if gracefulShutdownTimeout == 0 {
+				if err := withCmdLock(func() error { return cmd.Process.Signal(os.Kill) }); err != nil {
+					tf.logger.Printf("[ERROR] Error sending SIGKILL to terraform: %v", err)
+				}
+				cancel() // to cancel stdout/stderr writers
+				tf.logger.Printf("[ERROR] terraform forcefully killed")
+				returnCh <- fmt.Errorf("%w: terraform forcefully killed", parentCtx.Err())
+				return
+			}
+			tf.logger.Printf("[WARN] The context was cancelled, we'll let Terraform finish by sending SIGINT signal")
+			if err := withCmdLock(func() error { return cmd.Process.Signal(os.Interrupt) }); err != nil {
+				tf.logger.Printf("[ERROR] Error sending SIGINT to terraform: %v", err)
+			}
+			// give some time to the process before forcefully killing it
+			select {
+			case <-time.After(gracefulShutdownTimeout):
+				// Forcefully kill the process
+				if err := withCmdLock(func() error { return cmd.Process.Signal(os.Kill) }); err != nil {
+					tf.logger.Printf("[ERROR] Error sending SIGKILL to terraform: %v", err)
+				}
+				cancel() // to cancel stdout/stderr writers
+				tf.logger.Printf("[ERROR] terraform forcefully killed after graceful shutdown timeout")
+				returnCh <- fmt.Errorf("%w: terraform forcefully killed after graceful shutdown timeout", parentCtx.Err())
+			case err := <-cmdDoneCh:
+				returnCh <- fmt.Errorf("%w: %v", parentCtx.Err(), err)
+				tf.logger.Printf("[INFO] terraform successfully interrupted")
+			}
+		case err := <-cmdDoneCh:
+			tf.logger.Printf("[DEBUG] terraform run finished")
+			returnCh <- err
+		}
+	}()
+	err = withCmdLock(func() error { return cmd.Start() })
 	if err == nil && ctx.Err() != nil {
 		err = ctx.Err()
 	}
@@ -70,12 +126,14 @@ func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
 	// can cause a race condition
 	wg.Wait()
 
-	err = cmd.Wait()
-	if err == nil && ctx.Err() != nil {
-		err = ctx.Err()
+	cmdDoneCh <- cmd.Wait()
+	err = <-returnCh
+
+	if err == nil && parentCtx.Err() != nil {
+		err = parentCtx.Err()
 	}
 	if err != nil {
-		return tf.wrapExitError(ctx, err, errBuf.String())
+		return tf.wrapExitError(parentCtx, err, errBuf.String())
 	}
 
 	// Return error if there was an issue reading the std out/err

tfexec/cmd_linux_test.go

@@ -21,7 +21,7 @@ func Test_runTerraformCmd_linux(t *testing.T) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	cmd := tf.buildTerraformCmd(ctx, nil, "hello tf-exec!")
+	cmd := tf.buildTerraformCmd(nil, "hello tf-exec!")
 	err := tf.runTerraformCmd(ctx, cmd)
 	if err != nil {
 		t.Fatal(err)