From e98f2016027578887926762377da8ad9a9fc8288 Mon Sep 17 00:00:00 2001
From: Matthew Hooker <mwhooker@gmail.com>
Date: Fri, 19 Jan 2018 20:09:05 -0800
Subject: [PATCH] working with opt-in

---
 builder/amazon/common/step_run_source_instance.go | 13 ++++---------
 helper/multistep/basic_runner.go                  |  8 +++++++-
 helper/multistep/debug_runner.go                  |  3 +--
 helper/multistep/multistep.go                     |  6 +++++-
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/builder/amazon/common/step_run_source_instance.go b/builder/amazon/common/step_run_source_instance.go
index 9186c868c..dcf6fea70 100644
--- a/builder/amazon/common/step_run_source_instance.go
+++ b/builder/amazon/common/step_run_source_instance.go
@@ -37,6 +37,10 @@ type StepRunSourceInstance struct {
 }
 
 func (s *StepRunSourceInstance) Run(state multistep.StateBag) multistep.StepAction {
+	return s.RunWithContext(context.Background(), state)
+}
+
+func (s *StepRunSourceInstance) RunWithContext(ctx context.Context, state multistep.StateBag) multistep.StepAction {
 	ec2conn := state.Get("ec2").(*ec2.EC2)
 	var keyName string
 	if name, ok := state.GetOk("keyPair"); ok {
@@ -179,15 +183,6 @@ func (s *StepRunSourceInstance) Run(state multistep.StateBag) multistep.StepActi
 	describeInstance := &ec2.DescribeInstancesInput{
 		InstanceIds: []*string{aws.String(instanceId)},
 	}
-	ctx, cancel := context.WithCancel(context.Background())
-
-	go func() {
-		for {
-			if _, ok := state.GetOk(multistep.StateCancelled); ok {
-				cancel()
-			}
-		}
-	}()
 
 	if err := ec2conn.WaitUntilInstanceRunningWithContext(ctx, describeInstance); err != nil {
 		err := fmt.Errorf("Error waiting for instance (%s) to become ready: %s", instanceId, err)
diff --git a/helper/multistep/basic_runner.go b/helper/multistep/basic_runner.go
index eb5018d36..f152178bc 100644
--- a/helper/multistep/basic_runner.go
+++ b/helper/multistep/basic_runner.go
@@ -70,7 +70,13 @@ func (b *BasicRunner) Run(state StateBag) {
 			break
 		}
 
-		action := step.Run(ctx, state)
+		var action StepAction
+
+		if stepCtx, ok := step.(StepRunnableWithContext); ok {
+			action = stepCtx.RunWithContext(ctx, state)
+		} else {
+			action = step.Run(state)
+		}
 		defer step.Cleanup(state)
 
 		if _, ok := state.GetOk(StateCancelled); ok {
diff --git a/helper/multistep/debug_runner.go b/helper/multistep/debug_runner.go
index 88af01c54..882009494 100644
--- a/helper/multistep/debug_runner.go
+++ b/helper/multistep/debug_runner.go
@@ -1,7 +1,6 @@
 package multistep
 
 import (
-	"context"
 	"fmt"
 	"reflect"
 	"sync"
@@ -114,7 +113,7 @@ type debugStepPause struct {
 	PauseFn  DebugPauseFn
 }
 
-func (s *debugStepPause) Run(_ context.Context, state StateBag) StepAction {
+func (s *debugStepPause) Run(state StateBag) StepAction {
 	s.PauseFn(DebugLocationAfterRun, s.StepName, state)
 	return ActionContinue
 }
diff --git a/helper/multistep/multistep.go b/helper/multistep/multistep.go
index 7b4e0801f..9b8a40b6d 100644
--- a/helper/multistep/multistep.go
+++ b/helper/multistep/multistep.go
@@ -19,6 +19,10 @@ const StateCancelled = "cancelled"
 // This is the key set in the state bag when a step halted the sequence.
 const StateHalted = "halted"
 
+type StepRunnableWithContext interface {
+	RunWithContext(context.Context, StateBag) StepAction
+}
+
 // Step is a single step that is part of a potentially large sequence
 // of other steps, responsible for performing some specific action.
 type Step interface {
@@ -28,7 +32,7 @@ type Step interface {
 	//
 	// The return value determines whether multi-step sequences continue
 	// or should halt.
-	Run(context.Context, StateBag) StepAction
+	Run(StateBag) StepAction
 
 	// Cleanup is called in reverse order of the steps that have run
 	// and allow steps to clean up after themselves. Do not assume if this