From 6e098d1aaf9247a04eefa7457f781a0ae13419af Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 30 Aug 2013 23:39:29 -0700
Subject: [PATCH] packer: implement Cancel in ProvisionHook

---
 common/step_provision.go   |  1 +
 packer/build.go            |  4 ++-
 packer/provisioner.go      | 25 +++++++++++++++++-
 packer/provisioner_mock.go |  9 ++++++-
 packer/provisioner_test.go | 53 ++++++++++++++++++++++++++++++++++++--
 5 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/common/step_provision.go b/common/step_provision.go
index b51896208..d069a5c0c 100644
--- a/common/step_provision.go
+++ b/common/step_provision.go
@@ -42,6 +42,7 @@ func (*StepProvision) Run(state map[string]interface{}) multistep.StepAction {
 			return multistep.ActionContinue
 		case <-time.After(1 * time.Second):
 			if _, ok := state[multistep.StateCancelled]; ok {
+				log.Println("Cancelling provisioning due to interrupt...")
 				hook.Cancel()
 				return multistep.ActionHalt
 			}
diff --git a/packer/build.go b/packer/build.go
index cfd23d703..e324fdea7 100644
--- a/packer/build.go
+++ b/packer/build.go
@@ -207,7 +207,9 @@ func (b *coreBuild) Run(originalUi Ui, cache Cache) ([]Artifact, error) {
 			hooks[HookProvision] = make([]Hook, 0, 1)
 		}
 
-		hooks[HookProvision] = append(hooks[HookProvision], &ProvisionHook{provisioners})
+		hooks[HookProvision] = append(hooks[HookProvision], &ProvisionHook{
+			Provisioners: provisioners,
+		})
 	}
 
 	hook := &DispatchHook{Mapping: hooks}
diff --git a/packer/provisioner.go b/packer/provisioner.go
index 246795c35..3592fb919 100644
--- a/packer/provisioner.go
+++ b/packer/provisioner.go
@@ -1,5 +1,9 @@
 package packer
 
+import (
+	"sync"
+)
+
 // A provisioner is responsible for installing and configuring software
 // on a machine prior to building the actual image.
 type Provisioner interface {
@@ -25,11 +29,25 @@ type ProvisionHook struct {
 	// The provisioners to run as part of the hook. These should already
 	// be prepared (by calling Prepare) at some earlier stage.
 	Provisioners []Provisioner
+
+	lock               sync.Mutex
+	runningProvisioner Provisioner
 }
 
 // Runs the provisioners in order.
 func (h *ProvisionHook) Run(name string, ui Ui, comm Communicator, data interface{}) error {
+	defer func() {
+		h.lock.Lock()
+		defer h.lock.Unlock()
+
+		h.runningProvisioner = nil
+	}()
+
 	for _, p := range h.Provisioners {
+		h.lock.Lock()
+		h.runningProvisioner = p
+		h.lock.Unlock()
+
 		if err := p.Provision(ui, comm); err != nil {
 			return err
 		}
@@ -40,5 +58,10 @@ func (h *ProvisionHook) Run(name string, ui Ui, comm Communicator, data interfac
 
 // Cancels the privisioners that are still running.
 func (h *ProvisionHook) Cancel() {
-	// TODO(mitchellh): implement
+	h.lock.Lock()
+	defer h.lock.Unlock()
+
+	if h.runningProvisioner != nil {
+		h.runningProvisioner.Cancel()
+	}
 }
diff --git a/packer/provisioner_mock.go b/packer/provisioner_mock.go
index d05d8eda7..b61f642af 100644
--- a/packer/provisioner_mock.go
+++ b/packer/provisioner_mock.go
@@ -3,6 +3,8 @@ package packer
 // MockProvisioner is an implementation of Provisioner that can be
 // used for tests.
 type MockProvisioner struct {
+	ProvFunc func() error
+
 	PrepCalled   bool
 	PrepConfigs  []interface{}
 	ProvCalled   bool
@@ -19,7 +21,12 @@ func (t *MockProvisioner) Prepare(configs ...interface{}) error {
 func (t *MockProvisioner) Provision(ui Ui, comm Communicator) error {
 	t.ProvCalled = true
 	t.ProvUi = ui
-	return nil
+
+	if t.ProvFunc == nil {
+		return nil
+	}
+
+	return t.ProvFunc()
 }
 
 func (t *MockProvisioner) Cancel() {
diff --git a/packer/provisioner_test.go b/packer/provisioner_test.go
index 6107531e2..a3d97d511 100644
--- a/packer/provisioner_test.go
+++ b/packer/provisioner_test.go
@@ -1,6 +1,10 @@
 package packer
 
-import "testing"
+import (
+	"sync"
+	"testing"
+	"time"
+)
 
 func TestProvisionHook_Impl(t *testing.T) {
 	var raw interface{}
@@ -18,7 +22,10 @@ func TestProvisionHook(t *testing.T) {
 	var comm Communicator = nil
 	var data interface{} = nil
 
-	hook := &ProvisionHook{[]Provisioner{pA, pB}}
+	hook := &ProvisionHook{
+		Provisioners: []Provisioner{pA, pB},
+	}
+
 	hook.Run("foo", ui, comm, data)
 
 	if !pA.ProvCalled {
@@ -30,4 +37,46 @@ func TestProvisionHook(t *testing.T) {
 	}
 }
 
+func TestProvisionHook_cancel(t *testing.T) {
+	var lock sync.Mutex
+	order := make([]string, 0, 2)
+
+	p := &MockProvisioner{
+		ProvFunc: func() error {
+			time.Sleep(50 * time.Millisecond)
+
+			lock.Lock()
+			defer lock.Unlock()
+			order = append(order, "prov")
+
+			return nil
+		},
+	}
+
+	hook := &ProvisionHook{
+		Provisioners: []Provisioner{p},
+	}
+
+	finished := make(chan struct{})
+	go func() {
+		hook.Run("foo", nil, nil, nil)
+		close(finished)
+	}()
+
+	// Cancel it while it is running
+	time.Sleep(10 * time.Millisecond)
+	hook.Cancel()
+	lock.Lock()
+	order = append(order, "cancel")
+	lock.Unlock()
+
+	// Wait
+	<-finished
+
+	// Verify order
+	if order[0] != "cancel" || order[1] != "prov" {
+		t.Fatalf("bad: %#v", order)
+	}
+}
+
 // TODO(mitchellh): Test that they're run in the proper order