Merge pull request #16816 from hashicorp/jbardin/plan-shutdown

Make plan command cancellable

command/apply.go

@@ -24,9 +24,6 @@ type ApplyCommand struct {
 	// If true, then this apply command will become the "destroy"
 	// command. It is just like apply but only processes a destroy.
 	Destroy bool
-
-	// When this channel is closed, the apply will be cancelled.
-	ShutdownCh <-chan struct{}
 }
 
 func (c *ApplyCommand) Run(args []string) int {
@@ -186,6 +183,11 @@ func (c *ApplyCommand) Run(args []string) int {
 		// Cancel our context so we can start gracefully exiting
 		ctxCancel()
 
+		// notify tests that the command context was canceled
+		if testShutdownHook != nil {
+			testShutdownHook()
+		}
+
 		// Notify the user
 		c.Ui.Output(outputInterrupt)
 

command/apply_test.go

@@ -824,22 +824,27 @@ func TestApply_refresh(t *testing.T) {
 }
 
 func TestApply_shutdown(t *testing.T) {
-	stopped := false
-	stopCh := make(chan struct{})
-	stopReplyCh := make(chan struct{})
+	cancelled := false
+	cancelDone := make(chan struct{})
+	testShutdownHook = func() {
+		cancelled = true
+		close(cancelDone)
+	}
+	defer func() {
+		testShutdownHook = nil
+	}()
 
 	statePath := testTempFile(t)
-
 	p := testProvider()
 	shutdownCh := make(chan struct{})
+
 	ui := new(cli.MockUi)
 	c := &ApplyCommand{
 		Meta: Meta{
 			testingOverrides: metaOverridesForProvider(p),
 			Ui:               ui,
-		},
-
 			ShutdownCh:       shutdownCh,
+		},
 	}
 
 	p.DiffFn = func(
@@ -858,10 +863,10 @@ func TestApply_shutdown(t *testing.T) {
 		*terraform.InstanceInfo,
 		*terraform.InstanceState,
 		*terraform.InstanceDiff) (*terraform.InstanceState, error) {
-		if !stopped {
-			stopped = true
-			close(stopCh)
-			<-stopReplyCh
+
+		if !cancelled {
+			shutdownCh <- struct{}{}
+			<-cancelDone
 		}
 
 		return &terraform.InstanceState{
@@ -872,18 +877,6 @@ func TestApply_shutdown(t *testing.T) {
 		}, nil
 	}
 
-	go func() {
-		<-stopCh
-		shutdownCh <- struct{}{}
-
-		// This is really dirty, but we have no other way to assure that
-		// tf.Stop() has been called. This doesn't assure it either, but
-		// it makes it much more certain.
-		time.Sleep(50 * time.Millisecond)
-
-		close(stopReplyCh)
-	}()
-
 	args := []string{
 		"-state", statePath,
 		"-auto-approve",
@@ -897,6 +890,10 @@ func TestApply_shutdown(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
+	if !cancelled {
+		t.Fatal("command not cancelled")
+	}
+
 	state := testStateRead(t, statePath)
 	if state == nil {
 		t.Fatal("state should not be nil")

command/console.go

@@ -17,9 +17,6 @@ import (
 // configuration and actually builds or changes infrastructure.
 type ConsoleCommand struct {
 	Meta
-
-	// When this channel is closed, the apply will be cancelled.
-	ShutdownCh <-chan struct{}
 }
 
 func (c *ConsoleCommand) Run(args []string) int {

command/meta.go

@@ -76,6 +76,9 @@ type Meta struct {
 	// is not suitable, e.g. because of a read-only filesystem.
 	OverrideDataDir string
 
+	// When this channel is closed, the command will be cancelled.
+	ShutdownCh <-chan struct{}
+
 	//----------------------------------------------------------
 	// Protected: commands can set these
 	//----------------------------------------------------------
@@ -638,3 +641,7 @@ func isAutoVarFile(path string) bool {
 	return strings.HasSuffix(path, ".auto.tfvars") ||
 		strings.HasSuffix(path, ".auto.tfvars.json")
 }
+
+// testShutdownHook is used by tests to verify that a command context has been
+// canceled
+var testShutdownHook func()

command/plan.go

@@ -104,18 +104,42 @@ func (c *PlanCommand) Run(args []string) int {
 	opReq.Type = backend.OperationTypePlan
 
 	// Perform the operation
-	op, err := b.Operation(context.Background(), opReq)
+	ctx, ctxCancel := context.WithCancel(context.Background())
+	defer ctxCancel()
+
+	op, err := b.Operation(ctx, opReq)
 	if err != nil {
 		c.Ui.Error(fmt.Sprintf("Error starting operation: %s", err))
 		return 1
 	}
 
-	// Wait for the operation to complete
-	<-op.Done()
+	select {
+	case <-c.ShutdownCh:
+		// Cancel our context so we can start gracefully exiting
+		ctxCancel()
+
+		// notify tests that the command context was canceled
+		if testShutdownHook != nil {
+			testShutdownHook()
+		}
+
+		// Notify the user
+		c.Ui.Output(outputInterrupt)
+
+		// Still get the result, since there is still one
+		select {
+		case <-c.ShutdownCh:
+			c.Ui.Error(
+				"Two interrupts received. Exiting immediately")
+			return 1
+		case <-op.Done():
+		}
+	case <-op.Done():
 		if err := op.Err; err != nil {
 			c.showDiagnostics(err)
 			return 1
 		}
+	}
 
 	/*
 		err = terraform.SetDebugInfo(DefaultDataDir)

command/plan_test.go

@@ -831,6 +831,56 @@ func TestPlan_detailedExitcode_emptyDiff(t *testing.T) {
 	}
 }
 
+func TestPlan_shutdown(t *testing.T) {
+	cancelled := false
+	cancelDone := make(chan struct{})
+	testShutdownHook = func() {
+		cancelled = true
+		close(cancelDone)
+	}
+	defer func() {
+		testShutdownHook = nil
+	}()
+
+	shutdownCh := make(chan struct{})
+	p := testProvider()
+	ui := new(cli.MockUi)
+	c := &PlanCommand{
+		Meta: Meta{
+			testingOverrides: metaOverridesForProvider(p),
+			Ui:               ui,
+			ShutdownCh:       shutdownCh,
+		},
+	}
+
+	p.DiffFn = func(
+		*terraform.InstanceInfo,
+		*terraform.InstanceState,
+		*terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
+
+		if !cancelled {
+			shutdownCh <- struct{}{}
+			<-cancelDone
+		}
+
+		return &terraform.InstanceDiff{
+			Attributes: map[string]*terraform.ResourceAttrDiff{
+				"ami": &terraform.ResourceAttrDiff{
+					New: "bar",
+				},
+			},
+		}, nil
+	}
+
+	if code := c.Run([]string{testFixturePath("apply-shutdown")}); code != 0 {
+		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
+	}
+
+	if !cancelled {
+		t.Fatal("command not cancelled")
+	}
+}
+
 const planVarFile = `
 foo = "bar"
 `

commands.go

@@ -63,6 +63,8 @@ func initCommands(config *Config) {
 		RunningInAutomation: inAutomation,
 		PluginCacheDir:      config.PluginCacheDir,
 		OverrideDataDir:     dataDir,
+
+		ShutdownCh: makeShutdownCh(),
 	}
 
 	// The command list is included in the terraform -help
@@ -81,14 +83,12 @@ func initCommands(config *Config) {
 		"apply": func() (cli.Command, error) {
 			return &command.ApplyCommand{
 				Meta: meta,
-				ShutdownCh: makeShutdownCh(),
 			}, nil
 		},
 
 		"console": func() (cli.Command, error) {
 			return &command.ConsoleCommand{
 				Meta: meta,
-				ShutdownCh: makeShutdownCh(),
 			}, nil
 		},
 
@@ -96,7 +96,6 @@ func initCommands(config *Config) {
 			return &command.ApplyCommand{
 				Meta:    meta,
 				Destroy: true,
-				ShutdownCh: makeShutdownCh(),
 			}, nil
 		},