pass context into step.run

helper/multistep/basic_runner.go

@@ -28,7 +28,7 @@ type BasicRunner struct {
 }
 
 func (b *BasicRunner) Run(state StateBag) {
-	ctx, cancel := context.WithCancel(state.Context())
+	ctx, cancel := context.WithCancel(context.Background())
 
 	b.l.Lock()
 	if b.state != stateIdle {
@@ -70,7 +70,7 @@ func (b *BasicRunner) Run(state StateBag) {
 			break
 		}
 
-		action := step.Run(state)
+		action := step.Run(ctx, state)
 		defer step.Cleanup(state)
 
 		if _, ok := state.GetOk(StateCancelled); ok {

helper/multistep/debug_runner.go

@@ -1,6 +1,7 @@
 package multistep
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"sync"
@@ -113,7 +114,7 @@ type debugStepPause struct {
 	PauseFn  DebugPauseFn
 }
 
-func (s *debugStepPause) Run(state StateBag) StepAction {
+func (s *debugStepPause) Run(_ context.Context, state StateBag) StepAction {
 	s.PauseFn(DebugLocationAfterRun, s.StepName, state)
 	return ActionContinue
 }

helper/multistep/multistep.go

@@ -2,6 +2,8 @@
 // discrete steps.
 package multistep
 
+import "context"
+
 // A StepAction determines the next step to take regarding multi-step actions.
 type StepAction uint
 
@@ -26,7 +28,7 @@ type Step interface {
 	//
 	// The return value determines whether multi-step sequences continue
 	// or should halt.
-	Run(StateBag) StepAction
+	Run(context.Context, 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

helper/multistep/statebag.go

@@ -1,32 +1,31 @@
 package multistep
 
 import (
-	"context"
 	"sync"
 )
 
-// StateBag implements StateBag by using a normal map underneath
+// StateBag holds the state that is used by the Runner and Steps. The
+// StateBag implementation must be safe for concurrent access.
+type StateBag interface {
+	Get(string) interface{}
+	GetOk(string) (interface{}, bool)
+	Put(string, interface{})
+}
+
+// BasicStateBag implements StateBag by using a normal map underneath
 // protected by a RWMutex.
-type StateBag struct {
+type BasicStateBag struct {
 	data map[string]interface{}
 	l    sync.RWMutex
 	once sync.Once
-	ctx  context.Context
 }
 
-func (b *StateBag) Context() context.Context {
-	if b.ctx != nil {
-		return b.ctx
-	}
-	return context.Background()
-}
-
-func (b *StateBag) Get(k string) interface{} {
+func (b *BasicStateBag) Get(k string) interface{} {
 	result, _ := b.GetOk(k)
 	return result
 }
 
-func (b *StateBag) GetOk(k string) (interface{}, bool) {
+func (b *BasicStateBag) GetOk(k string) (interface{}, bool) {
 	b.l.RLock()
 	defer b.l.RUnlock()
 
@@ -34,7 +33,7 @@ func (b *StateBag) GetOk(k string) (interface{}, bool) {
 	return result, ok
 }
 
-func (b *StateBag) Put(k string, v interface{}) {
+func (b *BasicStateBag) Put(k string, v interface{}) {
 	b.l.Lock()
 	defer b.l.Unlock()