Merge pull request #11 from hashicorp/f-cancel

Ctrl-C

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

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{

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}"
+}

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
+}

terraform/hook.go

@@ -65,3 +65,19 @@ func (*NilHook) PreRefresh(string, *ResourceState) (HookAction, error) {
 func (*NilHook) PostRefresh(string, *ResourceState) (HookAction, error) {
 	return HookActionContinue, nil
 }
+
+// handleHook turns hook actions into panics. This lets you use the
+// panic/recover mechanism in Go as a flow control mechanism for hook
+// actions.
+func handleHook(a HookAction, err error) {
+	if err != nil {
+		// TODO: handle errors
+	}
+
+	switch a {
+	case HookActionContinue:
+		return
+	case HookActionHalt:
+		panic(HookActionHalt)
+	}
+}

terraform/hook_stop.go

@@ -0,0 +1,79 @@
+package terraform
+
+import (
+	"sync"
+)
+
+// stopHook is a private Hook implementation that Terraform uses to
+// signal when to stop or cancel actions.
+type stopHook struct {
+	sync.Mutex
+
+	// This should be incremented for every thing that can be stopped.
+	// When this is zero, a stopper can assume that everything is properly
+	// stopped.
+	count int
+
+	// This channel should be closed when it is time to stop
+	ch chan struct{}
+
+	serial    int
+	stoppedCh chan<- struct{}
+}
+
+func (h *stopHook) PreApply(string, *ResourceState, *ResourceDiff) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) PostApply(string, *ResourceState) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) PreDiff(string, *ResourceState) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) PostDiff(string, *ResourceDiff) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) PreRefresh(string, *ResourceState) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) PostRefresh(string, *ResourceState) (HookAction, error) {
+	return h.hook()
+}
+
+func (h *stopHook) hook() (HookAction, error) {
+	select {
+	case <-h.ch:
+		h.stoppedCh <- struct{}{}
+		return HookActionHalt, nil
+	default:
+		return HookActionContinue, nil
+	}
+}
+
+// reset should be called within the lock context
+func (h *stopHook) reset() {
+	h.ch = make(chan struct{})
+	h.count = 0
+	h.serial += 1
+	h.stoppedCh = nil
+}
+
+func (h *stopHook) ref() int {
+	h.Lock()
+	defer h.Unlock()
+	h.count++
+	return h.serial
+}
+
+func (h *stopHook) unref(s int) {
+	h.Lock()
+	defer h.Unlock()
+	if h.serial == s {
+		h.count--
+	}
+}

terraform/hook_stop_test.go

@@ -0,0 +1,9 @@
+package terraform
+
+import (
+	"testing"
+)
+
+func TestStopHook_impl(t *testing.T) {
+	var _ Hook = new(stopHook)
+}

terraform/terraform.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log"
 	"sync"
+	"sync/atomic"
 
 	"github.com/hashicorp/terraform/config"
 	"github.com/hashicorp/terraform/depgraph"
@@ -15,6 +16,7 @@ import (
 type Terraform struct {
 	hooks     []Hook
 	providers map[string]ResourceProviderFactory
+	stopHook  *stopHook
 }
 
 // This is a function type used to implement a walker for the resource
@@ -35,13 +37,29 @@ type Config struct {
 // time, as well as richer checks such as verifying that the resource providers
 // can be properly initialized, can be configured, etc.
 func New(c *Config) (*Terraform, error) {
+	sh := new(stopHook)
+	sh.Lock()
+	sh.reset()
+	sh.Unlock()
+
+	// Copy all the hooks and add our stop hook. We don't append directly
+	// to the Config so that we're not modifying that in-place.
+	hooks := make([]Hook, len(c.Hooks)+1)
+	copy(hooks, c.Hooks)
+	hooks[len(c.Hooks)] = sh
+
 	return &Terraform{
-		hooks:     c.Hooks,
+		hooks:     hooks,
+		stopHook:  sh,
 		providers: c.Providers,
 	}, nil
 }
 
 func (t *Terraform) Apply(p *Plan) (*State, error) {
+	// Increase the count on the stop hook so we know when to stop
+	serial := t.stopHook.ref()
+	defer t.stopHook.unref(serial)
+
 	// Make sure we're working with a plan that doesn't have null pointers
 	// everywhere, and is instead just empty otherwise.
 	p.init()
@@ -59,7 +77,40 @@ func (t *Terraform) Apply(p *Plan) (*State, error) {
 	return t.apply(g, p)
 }
 
+// Stop stops all running tasks (applies, plans, refreshes).
+//
+// This will block until all running tasks are stopped. While Stop is
+// blocked, any new calls to Apply, Plan, Refresh, etc. will also block. New
+// calls, however, will start once this Stop has returned.
+func (t *Terraform) Stop() {
+	log.Printf("[INFO] Terraform stopping tasks")
+
+	t.stopHook.Lock()
+	defer t.stopHook.Unlock()
+
+	// Setup the stoppedCh
+	stoppedCh := make(chan struct{}, t.stopHook.count)
+	t.stopHook.stoppedCh = stoppedCh
+
+	// Close the channel to signal that we're done
+	close(t.stopHook.ch)
+
+	// Expect the number of count stops...
+	log.Printf("[DEBUG] Waiting for %d tasks to stop", t.stopHook.count)
+	for i := 0; i < t.stopHook.count; i++ {
+		<-stoppedCh
+	}
+	log.Printf("[DEBUG] Stopped!")
+
+	// Success, everything stopped, reset everything
+	t.stopHook.reset()
+}
+
 func (t *Terraform) Plan(opts *PlanOpts) (*Plan, error) {
+	// Increase the count on the stop hook so we know when to stop
+	serial := t.stopHook.ref()
+	defer t.stopHook.unref(serial)
+
 	g, err := Graph(&GraphOpts{
 		Config:    opts.Config,
 		Providers: t.providers,
@@ -75,6 +126,10 @@ func (t *Terraform) Plan(opts *PlanOpts) (*Plan, error) {
 // Refresh goes through all the resources in the state and refreshes them
 // to their latest status.
 func (t *Terraform) Refresh(c *config.Config, s *State) (*State, error) {
+	// Increase the count on the stop hook so we know when to stop
+	serial := t.stopHook.ref()
+	defer t.stopHook.unref(serial)
+
 	g, err := Graph(&GraphOpts{
 		Config:    c,
 		Providers: t.providers,
@@ -119,8 +174,7 @@ func (t *Terraform) refreshWalkFn(result *State) depgraph.WalkFunc {
 
 	cb := func(r *Resource) (map[string]string, error) {
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PreRefresh(r.Id, r.State)
+			handleHook(h.PreRefresh(r.Id, r.State))
 		}
 
 		rs, err := r.Provider.Refresh(r.State)
@@ -139,8 +193,7 @@ func (t *Terraform) refreshWalkFn(result *State) depgraph.WalkFunc {
 		l.Unlock()
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PostRefresh(r.Id, rs)
+			handleHook(h.PostRefresh(r.Id, rs))
 		}
 
 		return nil, nil
@@ -175,11 +228,11 @@ func (t *Terraform) applyWalkFn(
 		// anything and that the diff has no computed values (pre-computed)
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PreApply(r.Id, r.State, diff)
+			handleHook(h.PreApply(r.Id, r.State, diff))
 		}
 
 		// With the completed diff, apply!
+		log.Printf("[DEBUG] %s: Executing Apply", r.Id)
 		rs, err := r.Provider.Apply(r.State, diff)
 		if err != nil {
 			return nil, err
@@ -219,8 +272,7 @@ func (t *Terraform) applyWalkFn(
 		r.State = rs
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PostApply(r.Id, r.State)
+			handleHook(h.PostApply(r.Id, r.State))
 		}
 
 		// Determine the new state and update variables
@@ -245,8 +297,7 @@ func (t *Terraform) planWalkFn(result *Plan, opts *PlanOpts) depgraph.WalkFunc {
 		var diff *ResourceDiff
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PreDiff(r.Id, r.State)
+			handleHook(h.PreDiff(r.Id, r.State))
 		}
 
 		if opts.Destroy {
@@ -279,8 +330,7 @@ func (t *Terraform) planWalkFn(result *Plan, opts *PlanOpts) depgraph.WalkFunc {
 		l.Unlock()
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PostDiff(r.Id, diff)
+			handleHook(h.PostDiff(r.Id, diff))
 		}
 
 		// Determine the new state and update variables
@@ -305,12 +355,20 @@ func (t *Terraform) genericWalkFn(
 		vars[fmt.Sprintf("var.%s", k)] = v
 	}
 
+	// This will keep track of whether we're stopped or not
+	var stop uint32 = 0
+
 	return func(n *depgraph.Noun) error {
 		// If it is the root node, ignore
 		if n.Name == GraphRootNode {
 			return nil
 		}
 
+		// If we're stopped, return right away
+		if atomic.LoadUint32(&stop) != 0 {
+			return nil
+		}
+
 		switch m := n.Meta.(type) {
 		case *GraphNodeResource:
 		case *GraphNodeResourceProvider:
@@ -359,6 +417,17 @@ func (t *Terraform) genericWalkFn(
 			rn.Resource.Config = nil
 		}
 
+		// Handle recovery of special panic scenarios
+		defer func() {
+			if v := recover(); v != nil {
+				if v == HookActionHalt {
+					atomic.StoreUint32(&stop, 1)
+				} else {
+					panic(v)
+				}
+			}
+		}()
+
 		// Call the callack
 		log.Printf("[INFO] Walking: %s", rn.Resource.Id)
 		newVars, err := cb(rn.Resource)

terraform/terraform_test.go

@@ -39,6 +39,87 @@ func TestTerraformApply(t *testing.T) {
 	}
 }
 
+func TestTerraformApply_cancel(t *testing.T) {
+	stopped := false
+	stopCh := make(chan struct{})
+	stopReplyCh := make(chan struct{})
+
+	rpAWS := new(MockResourceProvider)
+	rpAWS.ResourcesReturn = []ResourceType{
+		ResourceType{Name: "aws_instance"},
+	}
+	rpAWS.DiffFn = func(*ResourceState, *ResourceConfig) (*ResourceDiff, error) {
+		return &ResourceDiff{
+			Attributes: map[string]*ResourceAttrDiff{
+				"num": &ResourceAttrDiff{
+					New: "bar",
+				},
+			},
+		}, nil
+	}
+	rpAWS.ApplyFn = func(*ResourceState, *ResourceDiff) (*ResourceState, error) {
+		if !stopped {
+			stopped = true
+			close(stopCh)
+			<-stopReplyCh
+		}
+
+		return &ResourceState{
+			ID: "foo",
+			Attributes: map[string]string{
+				"num": "2",
+			},
+		}, nil
+	}
+
+	c := testConfig(t, "apply-cancel")
+	tf := testTerraform2(t, &Config{
+		Providers: map[string]ResourceProviderFactory{
+			"aws": testProviderFuncFixed(rpAWS),
+		},
+	})
+
+	p, err := tf.Plan(&PlanOpts{Config: c})
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Start the Apply in a goroutine
+	stateCh := make(chan *State)
+	go func() {
+		state, err := tf.Apply(p)
+		if err != nil {
+			panic(err)
+		}
+
+		stateCh <- state
+	}()
+
+	// Start a goroutine so we can inject exactly when we stop
+	s := tf.stopHook.ref()
+	go func() {
+		defer tf.stopHook.unref(s)
+		<-tf.stopHook.ch
+		close(stopReplyCh)
+		tf.stopHook.stoppedCh <- struct{}{}
+	}()
+
+	<-stopCh
+	tf.Stop()
+
+	state := <-stateCh
+
+	if len(state.Resources) != 1 {
+		t.Fatalf("bad: %#v", state.Resources)
+	}
+
+	actual := strings.TrimSpace(state.String())
+	expected := strings.TrimSpace(testTerraformApplyCancelStr)
+	if actual != expected {
+		t.Fatalf("bad: \n%s", actual)
+	}
+}
+
 func TestTerraformApply_compute(t *testing.T) {
 	// This tests that computed variables are properly re-diffed
 	// to get the value prior to application (Apply).
@@ -683,6 +764,12 @@ aws_instance.foo:
   type = aws_instance
 `
 
+const testTerraformApplyCancelStr = `
+aws_instance.foo:
+  ID = foo
+  num = 2
+`
+
 const testTerraformApplyComputeStr = `
 aws_instance.bar:
   ID = foo

terraform/test-fixtures/apply-cancel/main.tf

@@ -0,0 +1,7 @@
+resource "aws_instance" "foo" {
+    num = "2"
+}
+
+resource "aws_instance" "bar" {
+    foo = "${aws_instance.foo.num}"
+}