working with opt-in

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)

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 {

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
 }

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