From e21015140838446a07b73d597c2f1a9276cd441c Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 30 Aug 2013 17:26:51 -0700
Subject: [PATCH] packer: implement Cancel on DispatchHook

---
 packer/build.go     |  2 +-
 packer/hook.go      | 34 ++++++++++++++++++++++++-
 packer/hook_test.go | 60 ++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/packer/build.go b/packer/build.go
index b8e0fa165..cfd23d703 100644
--- a/packer/build.go
+++ b/packer/build.go
@@ -210,7 +210,7 @@ func (b *coreBuild) Run(originalUi Ui, cache Cache) ([]Artifact, error) {
 		hooks[HookProvision] = append(hooks[HookProvision], &ProvisionHook{provisioners})
 	}
 
-	hook := &DispatchHook{hooks}
+	hook := &DispatchHook{Mapping: hooks}
 	artifacts := make([]Artifact, 0, 1)
 
 	// The builder just has a normal Ui, but targetted
diff --git a/packer/hook.go b/packer/hook.go
index f31b62a5d..d151bb8c4 100644
--- a/packer/hook.go
+++ b/packer/hook.go
@@ -1,5 +1,9 @@
 package packer
 
+import (
+	"sync"
+)
+
 // This is the hook that should be fired for provisioners to run.
 const HookProvision = "packer_provision"
 
@@ -24,12 +28,20 @@ type Hook interface {
 // A Hook implementation that dispatches based on an internal mapping.
 type DispatchHook struct {
 	Mapping map[string][]Hook
+
+	l           sync.Mutex
+	cancelled   bool
+	runningHook Hook
 }
 
 // Runs the hook with the given name by dispatching it to the proper
 // hooks if a mapping exists. If a mapping doesn't exist, then nothing
 // happens.
 func (h *DispatchHook) Run(name string, ui Ui, comm Communicator, data interface{}) error {
+	h.l.Lock()
+	h.cancelled = false
+	h.l.Unlock()
+
 	hooks, ok := h.Mapping[name]
 	if !ok {
 		// No hooks for that name. No problem.
@@ -37,6 +49,15 @@ func (h *DispatchHook) Run(name string, ui Ui, comm Communicator, data interface
 	}
 
 	for _, hook := range hooks {
+		h.l.Lock()
+		if h.cancelled {
+			h.l.Unlock()
+			return nil
+		}
+
+		h.runningHook = hook
+		h.l.Unlock()
+
 		if err := hook.Run(name, ui, comm, data); err != nil {
 			return err
 		}
@@ -45,4 +66,15 @@ func (h *DispatchHook) Run(name string, ui Ui, comm Communicator, data interface
 	return nil
 }
 
-func (h *DispatchHook) Cancel() {}
+// Cancels all the hooks that are currently in-flight, if any. This will
+// block until the hooks are all cancelled.
+func (h *DispatchHook) Cancel() {
+	h.l.Lock()
+	defer h.l.Unlock()
+
+	if h.runningHook != nil {
+		h.runningHook.Cancel()
+	}
+
+	h.cancelled = true
+}
diff --git a/packer/hook_test.go b/packer/hook_test.go
index 06db0073d..bf88e65df 100644
--- a/packer/hook_test.go
+++ b/packer/hook_test.go
@@ -2,21 +2,57 @@ package packer
 
 import (
 	"cgl.tideland.biz/asserts"
+	"sync"
 	"testing"
+	"time"
 )
 
+// A helper Hook implementation for testing cancels.
+type CancelHook struct {
+	sync.Mutex
+	cancelCh chan struct{}
+	doneCh   chan struct{}
+
+	Cancelled bool
+}
+
+func (h *CancelHook) Run(string, Ui, Communicator, interface{}) error {
+	h.Lock()
+	h.cancelCh = make(chan struct{})
+	h.doneCh = make(chan struct{})
+	h.Unlock()
+
+	defer close(h.doneCh)
+
+	select {
+	case <-h.cancelCh:
+		h.Cancelled = true
+	case <-time.After(1 * time.Second):
+	}
+
+	return nil
+}
+
+func (h *CancelHook) Cancel() {
+	h.Lock()
+	close(h.cancelCh)
+	h.Unlock()
+
+	<-h.doneCh
+}
+
 func TestDispatchHook_Implements(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
 	var r Hook
-	c := &DispatchHook{nil}
+	c := &DispatchHook{}
 
 	assert.Implementor(c, &r, "should be a Hook")
 }
 
 func TestDispatchHook_Run_NoHooks(t *testing.T) {
 	// Just make sure nothing blows up
-	dh := &DispatchHook{make(map[string][]Hook)}
+	dh := &DispatchHook{}
 	dh.Run("foo", nil, nil, nil)
 }
 
@@ -27,10 +63,28 @@ func TestDispatchHook_Run(t *testing.T) {
 
 	mapping := make(map[string][]Hook)
 	mapping["foo"] = []Hook{hook}
-	dh := &DispatchHook{mapping}
+	dh := &DispatchHook{Mapping: mapping}
 	dh.Run("foo", nil, nil, 42)
 
 	assert.True(hook.RunCalled, "run should be called")
 	assert.Equal(hook.RunName, "foo", "should be proper event")
 	assert.Equal(hook.RunData, 42, "should be correct data")
 }
+
+func TestDispatchHook_cancel(t *testing.T) {
+	hook := new(CancelHook)
+
+	dh := &DispatchHook{
+		Mapping: map[string][]Hook{
+			"foo": []Hook{hook},
+		},
+	}
+
+	go dh.Run("foo", nil, nil, 42)
+	time.Sleep(100 * time.Millisecond)
+	dh.Cancel()
+
+	if !hook.Cancelled {
+		t.Fatal("hook should've cancelled")
+	}
+}