From 89d43256bb28fc91c2154e95330a8ac3207ee9f9 Mon Sep 17 00:00:00 2001
From: Matthew Hooker <mwhooker@gmail.com>
Date: Fri, 19 Jan 2018 15:59:33 -0800
Subject: [PATCH] pass context into step.run

---
 helper/multistep/basic_runner.go |  4 ++--
 helper/multistep/debug_runner.go |  3 ++-
 helper/multistep/multistep.go    |  4 +++-
 helper/multistep/statebag.go     | 27 +++++++++++++--------------
 4 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/helper/multistep/basic_runner.go b/helper/multistep/basic_runner.go
index 305b801dd..eb5018d36 100644
--- a/helper/multistep/basic_runner.go
+++ b/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 {
diff --git a/helper/multistep/debug_runner.go b/helper/multistep/debug_runner.go
index 882009494..88af01c54 100644
--- a/helper/multistep/debug_runner.go
+++ b/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
 }
diff --git a/helper/multistep/multistep.go b/helper/multistep/multistep.go
index 4e3478dac..7b4e0801f 100644
--- a/helper/multistep/multistep.go
+++ b/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
diff --git a/helper/multistep/statebag.go b/helper/multistep/statebag.go
index d9f121561..dab712316 100644
--- a/helper/multistep/statebag.go
+++ b/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()