terraform: stopHook and tests

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

@@ -1,9 +1,11 @@
 package terraform
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"sync"
+	"sync/atomic"
 
 	"github.com/hashicorp/terraform/config"
 	"github.com/hashicorp/terraform/depgraph"
@@ -15,12 +17,17 @@ 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
 // tree internally on the Terraform structure.
 type genericWalkFunc func(*Resource) (map[string]string, error)
 
+// genericWalkStop is a special return value that can be returned from a
+// genericWalkFunc that causes the walk to cease immediately.
+var genericWalkStop error
+
 // Config is the configuration that must be given to instantiate
 // a Terraform structure.
 type Config struct {
@@ -28,6 +35,10 @@ type Config struct {
 	Providers map[string]ResourceProviderFactory
 }
 
+func init() {
+	genericWalkStop = errors.New("genericWalkStop")
+}
+
 // New creates a new Terraform structure, initializes resource providers
 // for the given configuration, etc.
 //
@@ -35,13 +46,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 +86,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 +135,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,
@@ -175,11 +239,19 @@ 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)
+			a, err := h.PreApply(r.Id, r.State, diff)
+			if err != nil {
+				return nil, err
+			}
+
+			switch a {
+			case HookActionHalt:
+				return nil, genericWalkStop
+			}
 		}
 
 		// 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 +291,15 @@ func (t *Terraform) applyWalkFn(
 		r.State = rs
 
 		for _, h := range t.hooks {
-			// TODO: return value
-			h.PostApply(r.Id, r.State)
+			a, err := h.PostApply(r.Id, r.State)
+			if err != nil {
+				return nil, err
+			}
+
+			switch a {
+			case HookActionHalt:
+				return nil, genericWalkStop
+			}
 		}
 
 		// Determine the new state and update variables
@@ -305,12 +384,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:
@@ -363,6 +450,11 @@ func (t *Terraform) genericWalkFn(
 		log.Printf("[INFO] Walking: %s", rn.Resource.Id)
 		newVars, err := cb(rn.Resource)
 		if err != nil {
+			if err == genericWalkStop {
+				atomic.StoreUint32(&stop, 1)
+				return nil
+			}
+
 			return err
 		}
 

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