From 3f0c4b0e1939d8c6af2318b67507f143f761444e Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 29 Jul 2013 18:47:43 -0700
Subject: [PATCH] builder/amazon/common: generic wait for state to wait for any
 state

---
 builder/amazon/common/instance.go             | 40 +++++++++++++------
 .../amazon/common/step_run_source_instance.go |  7 ++--
 builder/amazon/ebs/step_stop_instance.go      |  5 ++-
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/builder/amazon/common/instance.go b/builder/amazon/common/instance.go
index 442ed597c..5411916b0 100644
--- a/builder/amazon/common/instance.go
+++ b/builder/amazon/common/instance.go
@@ -11,17 +11,38 @@ import (
 
 type StateChangeConf struct {
 	Conn      *ec2.EC2
-	Instance  *ec2.Instance
 	Pending   []string
+	Refresh func() (interface{}, string, error)
 	StepState map[string]interface{}
 	Target    string
 }
 
-func WaitForState(conf *StateChangeConf) (i *ec2.Instance, err error) {
+func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) func() (interface{}, string, error) {
+	return func() (interface{}, string, error) {
+		resp, err := conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
+		if err != nil {
+			return nil, "", err
+		}
+
+		i = &resp.Reservations[0].Instances[0]
+		return i, i.State.Name, nil
+	}
+}
+
+func WaitForState(conf *StateChangeConf) (i interface{}, err error) {
 	log.Printf("Waiting for instance state to become: %s", conf.Target)
 
-	i = conf.Instance
-	for i.State.Name != conf.Target {
+	for {
+		var currentState string
+		i, currentState, err = conf.Refresh()
+		if err != nil {
+			return
+		}
+
+		if currentState == conf.Target {
+			return
+		}
+
 		if conf.StepState != nil {
 			if _, ok := conf.StepState[multistep.StateCancelled]; ok {
 				return nil, errors.New("interrupted")
@@ -30,24 +51,17 @@ func WaitForState(conf *StateChangeConf) (i *ec2.Instance, err error) {
 
 		found := false
 		for _, allowed := range conf.Pending {
-			if i.State.Name == allowed {
+			if currentState == allowed {
 				found = true
 				break
 			}
 		}
 
 		if !found {
-			fmt.Errorf("unexpected state '%s', wanted target '%s'", i.State.Name, conf.Target)
+			fmt.Errorf("unexpected state '%s', wanted target '%s'", currentState, conf.Target)
 			return
 		}
 
-		var resp *ec2.InstancesResp
-		resp, err = conf.Conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
-		if err != nil {
-			return
-		}
-
-		i = &resp.Reservations[0].Instances[0]
 		time.Sleep(2 * time.Second)
 	}
 
diff --git a/builder/amazon/common/step_run_source_instance.go b/builder/amazon/common/step_run_source_instance.go
index 22018246b..ec823ffc6 100644
--- a/builder/amazon/common/step_run_source_instance.go
+++ b/builder/amazon/common/step_run_source_instance.go
@@ -62,12 +62,13 @@ func (s *StepRunSourceInstance) Run(state map[string]interface{}) multistep.Step
 	ui.Say(fmt.Sprintf("Waiting for instance (%s) to become ready...", s.instance.InstanceId))
 	stateChange := StateChangeConf{
 		Conn:      ec2conn,
-		Instance:  s.instance,
 		Pending:   []string{"pending"},
 		Target:    "running",
+		Refresh:  InstanceStateRefreshFunc(ec2conn, s.instance),
 		StepState: state,
 	}
-	s.instance, err = WaitForState(&stateChange)
+	latestInstance, err := WaitForState(&stateChange)
+	s.instance = latestInstance.(*ec2.Instance)
 	if err != nil {
 		err := fmt.Errorf("Error waiting for instance (%s) to become ready: %s", s.instance.InstanceId, err)
 		state["error"] = err
@@ -96,8 +97,8 @@ func (s *StepRunSourceInstance) Cleanup(state map[string]interface{}) {
 
 	stateChange := StateChangeConf{
 		Conn:     ec2conn,
-		Instance: s.instance,
 		Pending:  []string{"pending", "running", "shutting-down", "stopped", "stopping"},
+		Refresh:  InstanceStateRefreshFunc(ec2conn, s.instance),
 		Target:   "running",
 	}
 
diff --git a/builder/amazon/ebs/step_stop_instance.go b/builder/amazon/ebs/step_stop_instance.go
index f0c7b47a9..616e91b29 100644
--- a/builder/amazon/ebs/step_stop_instance.go
+++ b/builder/amazon/ebs/step_stop_instance.go
@@ -29,12 +29,13 @@ func (s *stepStopInstance) Run(state map[string]interface{}) multistep.StepActio
 	ui.Say("Waiting for the instance to stop...")
 	stateChange := awscommon.StateChangeConf{
 		Conn:      ec2conn,
-		Instance:  instance,
 		Pending:   []string{"running", "stopping"},
 		Target:    "stopped",
+		Refresh:  awscommon.InstanceStateRefreshFunc(ec2conn, instance),
 		StepState: state,
 	}
-	instance, err = awscommon.WaitForState(&stateChange)
+	instanceRaw, err := awscommon.WaitForState(&stateChange)
+	instance = instanceRaw.(*ec2.Instance)
 	if err != nil {
 		err := fmt.Errorf("Error waiting for instance to stop: %s", err)
 		state["error"] = err