From 1d0adbf2c2ada88c6a17176c7e50dfef51eba97b Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 25 Jul 2013 20:49:15 -0500
Subject: [PATCH] builder/amazon/*: sigint while waiting for state change works

---
 CHANGELOG.md                                  |  9 +++---
 builder/amazon/common/instance.go             | 30 ++++++++++++++-----
 .../amazon/common/step_run_source_instance.go | 19 ++++++++++--
 builder/amazon/ebs/step_stop_instance.go      |  9 +++++-
 4 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a27cfdd71..83041b9a1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,12 +8,13 @@ FEATURES:
 
 IMPROVEMENTS:
 
-* amazon-ebs: Can now launch instances into a VPC for added protection [GH-210]
-* virtualbox, vmware: Add backspace, delete, and F1-F12 keys to the boot
+* builder/amazon/all: Ctrl-C while waiting for state change works
+* builder/amazon/ebs: Can now launch instances into a VPC for added protection [GH-210]
+* builder/virtualbox,vmware: Add backspace, delete, and F1-F12 keys to the boot
   command.
-* virtualbox: massive performance improvements with big ISO files because
+* builder/virtualbox: massive performance improvements with big ISO files because
   an expensive copy is avoided. [GH-202]
-* vmware: CD is removed prior to exporting final machine. [GH-198]
+* builder/vmware: CD is removed prior to exporting final machine. [GH-198]
 
 BUG FIXES:
 
diff --git a/builder/amazon/common/instance.go b/builder/amazon/common/instance.go
index 53dab50ce..442ed597c 100644
--- a/builder/amazon/common/instance.go
+++ b/builder/amazon/common/instance.go
@@ -1,19 +1,35 @@
 package common
 
 import (
+	"errors"
 	"fmt"
 	"github.com/mitchellh/goamz/ec2"
+	"github.com/mitchellh/multistep"
 	"log"
 	"time"
 )
 
-func WaitForState(ec2conn *ec2.EC2, originalInstance *ec2.Instance, pending []string, target string) (i *ec2.Instance, err error) {
-	log.Printf("Waiting for instance state to become: %s", target)
+type StateChangeConf struct {
+	Conn      *ec2.EC2
+	Instance  *ec2.Instance
+	Pending   []string
+	StepState map[string]interface{}
+	Target    string
+}
+
+func WaitForState(conf *StateChangeConf) (i *ec2.Instance, err error) {
+	log.Printf("Waiting for instance state to become: %s", conf.Target)
+
+	i = conf.Instance
+	for i.State.Name != conf.Target {
+		if conf.StepState != nil {
+			if _, ok := conf.StepState[multistep.StateCancelled]; ok {
+				return nil, errors.New("interrupted")
+			}
+		}
 
-	i = originalInstance
-	for i.State.Name != target {
 		found := false
-		for _, allowed := range pending {
+		for _, allowed := range conf.Pending {
 			if i.State.Name == allowed {
 				found = true
 				break
@@ -21,12 +37,12 @@ func WaitForState(ec2conn *ec2.EC2, originalInstance *ec2.Instance, pending []st
 		}
 
 		if !found {
-			fmt.Errorf("unexpected state '%s', wanted target '%s'", i.State.Name, target)
+			fmt.Errorf("unexpected state '%s', wanted target '%s'", i.State.Name, conf.Target)
 			return
 		}
 
 		var resp *ec2.InstancesResp
-		resp, err = ec2conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
+		resp, err = conf.Conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
 		if err != nil {
 			return
 		}
diff --git a/builder/amazon/common/step_run_source_instance.go b/builder/amazon/common/step_run_source_instance.go
index 2c7254b49..22018246b 100644
--- a/builder/amazon/common/step_run_source_instance.go
+++ b/builder/amazon/common/step_run_source_instance.go
@@ -60,7 +60,14 @@ func (s *StepRunSourceInstance) Run(state map[string]interface{}) multistep.Step
 	log.Printf("instance id: %s", s.instance.InstanceId)
 
 	ui.Say(fmt.Sprintf("Waiting for instance (%s) to become ready...", s.instance.InstanceId))
-	s.instance, err = WaitForState(ec2conn, s.instance, []string{"pending"}, "running")
+	stateChange := StateChangeConf{
+		Conn:      ec2conn,
+		Instance:  s.instance,
+		Pending:   []string{"pending"},
+		Target:    "running",
+		StepState: state,
+	}
+	s.instance, err = WaitForState(&stateChange)
 	if err != nil {
 		err := fmt.Errorf("Error waiting for instance (%s) to become ready: %s", s.instance.InstanceId, err)
 		state["error"] = err
@@ -87,6 +94,12 @@ func (s *StepRunSourceInstance) Cleanup(state map[string]interface{}) {
 		return
 	}
 
-	pending := []string{"pending", "running", "shutting-down", "stopped", "stopping"}
-	WaitForState(ec2conn, s.instance, pending, "terminated")
+	stateChange := StateChangeConf{
+		Conn:     ec2conn,
+		Instance: s.instance,
+		Pending:  []string{"pending", "running", "shutting-down", "stopped", "stopping"},
+		Target:   "running",
+	}
+
+	WaitForState(&stateChange)
 }
diff --git a/builder/amazon/ebs/step_stop_instance.go b/builder/amazon/ebs/step_stop_instance.go
index 5c608632f..f0c7b47a9 100644
--- a/builder/amazon/ebs/step_stop_instance.go
+++ b/builder/amazon/ebs/step_stop_instance.go
@@ -27,7 +27,14 @@ func (s *stepStopInstance) Run(state map[string]interface{}) multistep.StepActio
 
 	// Wait for the instance to actual stop
 	ui.Say("Waiting for the instance to stop...")
-	instance, err = awscommon.WaitForState(ec2conn, instance, []string{"running", "stopping"}, "stopped")
+	stateChange := awscommon.StateChangeConf{
+		Conn:      ec2conn,
+		Instance:  instance,
+		Pending:   []string{"running", "stopping"},
+		Target:    "stopped",
+		StepState: state,
+	}
+	instance, err = awscommon.WaitForState(&stateChange)
 	if err != nil {
 		err := fmt.Errorf("Error waiting for instance to stop: %s", err)
 		state["error"] = err