From 5aa6ada589be17b0eaf77e92bfd3694f9bf67b10 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Wed, 2 Jul 2014 17:01:02 -0700
Subject: [PATCH] command/apply: Ctrl-C works

---
 command/apply.go                             | 41 +++++++++-
 command/apply_test.go                        | 82 ++++++++++++++++++++
 command/test-fixtures/apply-shutdown/main.tf |  7 ++
 commands.go                                  | 23 +++++-
 4 files changed, 148 insertions(+), 5 deletions(-)
 create mode 100644 command/test-fixtures/apply-shutdown/main.tf

diff --git a/command/apply.go b/command/apply.go
index 4245187d87..028909b3d1 100644
--- a/command/apply.go
+++ b/command/apply.go
@@ -13,8 +13,9 @@ import (
 // ApplyCommand is a Command implementation that applies a Terraform
 // configuration and actually builds or changes infrastructure.
 type ApplyCommand struct {
-	TFConfig *terraform.Config
-	Ui       cli.Ui
+	ShutdownCh chan struct{}
+	TFConfig   *terraform.Config
+	Ui         cli.Ui
 }
 
 func (c *ApplyCommand) Run(args []string) int {
@@ -63,7 +64,41 @@ func (c *ApplyCommand) Run(args []string) int {
 		return 1
 	}
 
-	state, err := tf.Apply(plan)
+	errCh := make(chan error)
+	stateCh := make(chan *terraform.State)
+	go func() {
+		state, err := tf.Apply(plan)
+		if err != nil {
+			errCh <- err
+			return
+		}
+
+		stateCh <- state
+	}()
+
+	err = nil
+	var state *terraform.State
+	select {
+	case <-c.ShutdownCh:
+		c.Ui.Output("Interrupt received. Gracefully shutting down...")
+
+		// Stop execution
+		tf.Stop()
+
+		// Still get the result, since there is still one
+		select {
+		case <-c.ShutdownCh:
+			c.Ui.Error(
+				"Two interrupts received. Exiting immediately. Note that data\n" +
+					"loss may have occurred.")
+			return 1
+		case state = <-stateCh:
+		case err = <-errCh:
+		}
+	case state = <-stateCh:
+	case err = <-errCh:
+	}
+
 	if err != nil {
 		c.Ui.Error(fmt.Sprintf("Error applying plan: %s", err))
 		return 1
diff --git a/command/apply_test.go b/command/apply_test.go
index 9ea9c2dcc8..2e248b8244 100644
--- a/command/apply_test.go
+++ b/command/apply_test.go
@@ -85,6 +85,88 @@ func TestApply_plan(t *testing.T) {
 	}
 }
 
+func TestApply_shutdown(t *testing.T) {
+	stopped := false
+	stopCh := make(chan struct{})
+	stopReplyCh := make(chan struct{})
+
+	statePath := testTempFile(t)
+
+	p := testProvider()
+	shutdownCh := make(chan struct{})
+	ui := new(cli.MockUi)
+	c := &ApplyCommand{
+		ShutdownCh: shutdownCh,
+		TFConfig:   testTFConfig(p),
+		Ui:         ui,
+	}
+
+	p.DiffFn = func(
+		*terraform.ResourceState,
+		*terraform.ResourceConfig) (*terraform.ResourceDiff, error) {
+		return &terraform.ResourceDiff{
+			Attributes: map[string]*terraform.ResourceAttrDiff{
+				"ami": &terraform.ResourceAttrDiff{
+					New: "bar",
+				},
+			},
+		}, nil
+	}
+	p.ApplyFn = func(
+		*terraform.ResourceState,
+		*terraform.ResourceDiff) (*terraform.ResourceState, error) {
+		if !stopped {
+			stopped = true
+			close(stopCh)
+			<-stopReplyCh
+		}
+
+		return &terraform.ResourceState{
+			ID: "foo",
+			Attributes: map[string]string{
+				"ami": "2",
+			},
+		}, nil
+	}
+
+	go func() {
+		<-stopCh
+		shutdownCh <- struct{}{}
+		close(stopReplyCh)
+	}()
+
+	args := []string{
+		"-init",
+		statePath,
+		testFixturePath("apply-shutdown"),
+	}
+	if code := c.Run(args); code != 0 {
+		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
+	}
+
+	if _, err := os.Stat(statePath); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	f, err := os.Open(statePath)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer f.Close()
+
+	state, err := terraform.ReadState(f)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if state == nil {
+		t.Fatal("state should not be nil")
+	}
+
+	if len(state.Resources) != 1 {
+		t.Fatalf("bad: %d", len(state.Resources))
+	}
+}
+
 func TestApply_state(t *testing.T) {
 	originalState := &terraform.State{
 		Resources: map[string]*terraform.ResourceState{
diff --git a/command/test-fixtures/apply-shutdown/main.tf b/command/test-fixtures/apply-shutdown/main.tf
new file mode 100644
index 0000000000..1238f273ab
--- /dev/null
+++ b/command/test-fixtures/apply-shutdown/main.tf
@@ -0,0 +1,7 @@
+resource "test_instance" "foo" {
+    ami = "bar"
+}
+
+resource "test_instance" "bar" {
+    ami = "${test_instance.foo.ami}"
+}
diff --git a/commands.go b/commands.go
index c00210c529..2ff36f53df 100644
--- a/commands.go
+++ b/commands.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"os"
+	"os/signal"
 
 	"github.com/hashicorp/terraform/command"
 	"github.com/mitchellh/cli"
@@ -28,8 +29,9 @@ func init() {
 	Commands = map[string]cli.CommandFactory{
 		"apply": func() (cli.Command, error) {
 			return &command.ApplyCommand{
-				TFConfig: &TFConfig,
-				Ui:       Ui,
+				ShutdownCh: makeShutdownCh(),
+				TFConfig:   &TFConfig,
+				Ui:         Ui,
 			}, nil
 		},
 
@@ -64,3 +66,20 @@ func init() {
 		},
 	}
 }
+
+// makeShutdownCh creates an interrupt listener and returns a channel.
+// A message will be sent on the channel for every interrupt received.
+func makeShutdownCh() <-chan struct{} {
+	resultCh := make(chan struct{})
+
+	signalCh := make(chan os.Signal, 4)
+	signal.Notify(signalCh, os.Interrupt)
+	go func() {
+		for {
+			<-signalCh
+			resultCh <- struct{}{}
+		}
+	}()
+
+	return resultCh
+}