feat: add graceful shutdown option for apply&destroy (#5)

go.mod

@@ -1,6 +1,6 @@
 module github.com/hashicorp/terraform-exec
 
-go 1.18
+go 1.20
 
 require (
 	github.com/google/go-cmp v0.6.0

tfexec/apply.go

@@ -32,6 +32,8 @@ type applyConfig struct {
 	// Vars: each var must be supplied as a single string, e.g. 'foo=bar'
 	vars     []string
 	varFiles []string
+
+	gracefulShutdownConfig GracefulShutdownConfig
 }
 
 var defaultApplyOptions = applyConfig{
@@ -110,12 +112,31 @@ func (opt *AllowDeferralOption) configureApply(conf *applyConfig) {
 	conf.allowDeferral = opt.allowDeferral
 }
 
+func (opt *GracefulShutdownOption) configureApply(conf *applyConfig) {
+	conf.gracefulShutdownConfig = opt.config
+}
+
+func newApplyConfig(opts ...ApplyOption) applyConfig {
+	c := defaultApplyOptions
+
+	for _, o := range opts {
+		o.configureApply(&c)
+	}
+	return c
+}
+
 // Apply represents the terraform apply subcommand.
 func (tf *Terraform) Apply(ctx context.Context, opts ...ApplyOption) error {
-	cmd, err := tf.applyCmd(ctx, opts...)
+	c := newApplyConfig(opts...)
+
+	cmd, err := tf.applyCmd(ctx, c)
 	if err != nil {
 		return err
 	}
+
+	if c.gracefulShutdownConfig.Enable {
+		return tf.runTerraformCmdWithGracefulShutdown(ctx, cmd)
+	}
 	return tf.runTerraformCmd(ctx, cmd)
 }
 
@@ -132,21 +153,20 @@ func (tf *Terraform) ApplyJSON(ctx context.Context, w io.Writer, opts ...ApplyOp
 
 	tf.SetStdout(w)
 
-	cmd, err := tf.applyJSONCmd(ctx, opts...)
+	c := newApplyConfig(opts...)
+
+	cmd, err := tf.applyJSONCmd(ctx, c)
 	if err != nil {
 		return err
 	}
 
+	if c.gracefulShutdownConfig.Enable {
+		return tf.runTerraformCmdWithGracefulShutdown(ctx, cmd)
+	}
 	return tf.runTerraformCmd(ctx, cmd)
 }
 
-func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.Cmd, error) {
-	c := defaultApplyOptions
-
-	for _, o := range opts {
-		o.configureApply(&c)
-	}
-
+func (tf *Terraform) applyCmd(ctx context.Context, c applyConfig) (*exec.Cmd, error) {
 	args, err := tf.buildApplyArgs(ctx, c)
 	if err != nil {
 		return nil, err
@@ -155,13 +175,7 @@ func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.C
 	return tf.buildApplyCmd(ctx, c, args)
 }
 
-func (tf *Terraform) applyJSONCmd(ctx context.Context, opts ...ApplyOption) (*exec.Cmd, error) {
-	c := defaultApplyOptions
-
-	for _, o := range opts {
-		o.configureApply(&c)
-	}
-
+func (tf *Terraform) applyJSONCmd(ctx context.Context, c applyConfig) (*exec.Cmd, error) {
 	args, err := tf.buildApplyArgs(ctx, c)
 	if err != nil {
 		return nil, err
@@ -271,5 +285,8 @@ func (tf *Terraform) buildApplyCmd(ctx context.Context, c applyConfig, args []st
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	if c.gracefulShutdownConfig.Enable {
+		return tf.buildTerraformCmdWithGracefulShutdown(ctx, c.gracefulShutdownConfig.Period, mergeEnv, args...), nil
+	}
 	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
 }

tfexec/apply_test.go

@@ -22,7 +22,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, err := tf.applyCmd(context.Background(), newApplyConfig(
 			Backup("testbackup"),
 			LockTimeout("200s"),
 			State("teststate"),
@@ -40,7 +40,7 @@ func TestApplyCmd(t *testing.T) {
 			Var("var2=bar"),
 			Destroy(true),
 			DirOrPlan("testfile"),
-		)
+		))
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -72,7 +72,9 @@ func TestApplyCmd(t *testing.T) {
 
 	t.Run("refresh-only operation", func(t *testing.T) {
 		applyCmd, err := tf.applyCmd(context.Background(),
-			RefreshOnly(true),
+			newApplyConfig(
+				RefreshOnly(true),
+			),
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -104,22 +106,24 @@ func TestApplyJSONCmd(t *testing.T) {
 
 	t.Run("basic", func(t *testing.T) {
 		applyCmd, err := tf.applyJSONCmd(context.Background(),
-			Backup("testbackup"),
-			LockTimeout("200s"),
-			State("teststate"),
-			StateOut("teststateout"),
-			VarFile("foo.tfvars"),
-			VarFile("bar.tfvars"),
-			Lock(false),
-			Parallelism(99),
-			Refresh(false),
-			Replace("aws_instance.test"),
-			Replace("google_pubsub_topic.test"),
-			Target("target1"),
-			Target("target2"),
-			Var("var1=foo"),
-			Var("var2=bar"),
-			DirOrPlan("testfile"),
+			newApplyConfig(
+				Backup("testbackup"),
+				LockTimeout("200s"),
+				State("teststate"),
+				StateOut("teststateout"),
+				VarFile("foo.tfvars"),
+				VarFile("bar.tfvars"),
+				Lock(false),
+				Parallelism(99),
+				Refresh(false),
+				Replace("aws_instance.test"),
+				Replace("google_pubsub_topic.test"),
+				Target("target1"),
+				Target("target2"),
+				Var("var1=foo"),
+				Var("var2=bar"),
+				DirOrPlan("testfile"),
+			),
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -164,7 +168,9 @@ func TestApplyCmd_AllowDeferral(t *testing.T) {
 
 	t.Run("allow deferrals during apply", func(t *testing.T) {
 		applyCmd, err := tf.applyCmd(context.Background(),
-			AllowDeferral(true),
+			newApplyConfig(
+				AllowDeferral(true),
+			),
 		)
 		if err != nil {
 			t.Fatal(err)

tfexec/cmd.go

@@ -15,6 +15,8 @@ import (
 	"os"
 	"os/exec"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/hashicorp/terraform-exec/internal/version"
 )
@@ -193,6 +195,30 @@ func (tf *Terraform) buildTerraformCmd(ctx context.Context, mergeEnv map[string]
 	return cmd
 }
 
+func (tf *Terraform) buildTerraformCmdWithGracefulShutdown(ctx context.Context, gracefulShutdownPeriod time.Duration, mergeEnv map[string]string, args ...string) *exec.Cmd {
+	cmd := exec.CommandContext(ctx, tf.execPath, args...)
+
+	cmd.Env = tf.buildEnv(mergeEnv)
+	cmd.Dir = tf.workingDir
+
+	tf.logger.Printf("[INFO] configuring gracefully cancellation for command: %s", cmd.String())
+	cmd.WaitDelay = gracefulShutdownPeriod
+	cmd.Cancel = func() error {
+		tf.logger.Printf("[INFO] gracefully cancelling Terraform command: %s", cmd.String())
+		// os.Interrupt doesn't work on Windows so it will result into an immediate kill.
+		err := cmd.Process.Signal(os.Interrupt)
+		if err != nil && !errors.Is(err, os.ErrProcessDone) {
+			tf.logger.Printf("[ERROR] gracefully cancelling failed due to %s, forcefully cancelling Terraform command: %s", err.Error(), cmd.String())
+			err = cmd.Process.Kill()
+		}
+		return err
+	}
+
+	tf.logger.Printf("[INFO] running Terraform command: %s", cmd.String())
+
+	return cmd
+}
+
 func (tf *Terraform) runTerraformCmdJSON(ctx context.Context, cmd *exec.Cmd, v interface{}) error {
 	var outbuf = bytes.Buffer{}
 	cmd.Stdout = mergeWriters(cmd.Stdout, &outbuf)
@@ -275,3 +301,105 @@ func writeOutput(ctx context.Context, r io.ReadCloser, w io.Writer) error {
 		}
 	}
 }
+
+func (tf *Terraform) runTerraformCmdWithGracefulShutdown(ctx context.Context, cmd *exec.Cmd) error {
+	return tf.runTerraformCmdInner(ctx, cmd, true /* enableGracefulShutdown */)
+}
+
+func (tf *Terraform) runTerraformCmd(ctx context.Context, cmd *exec.Cmd) error {
+	return tf.runTerraformCmdInner(ctx, cmd, false /* enableGracefulShutdown */)
+}
+
+func (tf *Terraform) runTerraformCmdInner(ctx context.Context, cmd *exec.Cmd, enableGracefulShutdown bool) error {
+	var errBuf strings.Builder
+
+	cmd.SysProcAttr = defaultSysProcAttr
+
+	// check for early cancellation
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
+	// Read stdout / stderr logs from pipe instead of setting cmd.Stdout and
+	// cmd.Stderr because it can cause hanging when killing the command
+	// https://github.com/golang/go/issues/23019
+	stdoutWriter := mergeWriters(cmd.Stdout, tf.stdout)
+	stderrWriter := mergeWriters(tf.stderr, &errBuf)
+
+	cmd.Stderr = nil
+	cmd.Stdout = nil
+
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		return err
+	}
+
+	stderrPipe, err := cmd.StderrPipe()
+	if err != nil {
+		return err
+	}
+
+	err = cmd.Start()
+	if ctx.Err() != nil {
+		return cmdErr{
+			err:    err,
+			ctxErr: ctx.Err(),
+		}
+	}
+	if err != nil {
+		return err
+	}
+
+	var errStdout, errStderr error
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		errStdoutCtx := ctx
+		// When custom cancellation is enabled, avoid piping cancellation to prevent
+		// cancellation from being interrupted due to broken pipe.
+		if enableGracefulShutdown {
+			errStdoutCtx = context.WithoutCancel(ctx)
+		}
+		errStdout = writeOutput(errStdoutCtx, stdoutPipe, stdoutWriter)
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		errStderrCtx := ctx
+		// When custome cancellation is enabled, avoid piping cancellation to prevent
+		// cancellation from being interrupted due to broken pipe.
+		if enableGracefulShutdown {
+			errStderrCtx = context.WithoutCancel(ctx)
+		}
+		errStderr = writeOutput(errStderrCtx, stderrPipe, stderrWriter)
+	}()
+
+	// Reads from pipes must be completed before calling cmd.Wait(). Otherwise
+	// can cause a race condition
+	wg.Wait()
+
+	err = cmd.Wait()
+	if ctx.Err() != nil {
+		return cmdErr{
+			err:    err,
+			ctxErr: ctx.Err(),
+		}
+	}
+	if err != nil {
+		return fmt.Errorf("%w\n%s", err, errBuf.String())
+	}
+
+	// Return error if there was an issue reading the std out/err
+	if errStdout != nil && ctx.Err() != nil {
+		return fmt.Errorf("%w\n%s", errStdout, errBuf.String())
+	}
+	if errStderr != nil && ctx.Err() != nil {
+		return fmt.Errorf("%w\n%s", errStderr, errBuf.String())
+	}
+
+	return nil
+}