From 62e3d1362fd26e48c8b525fdc1b78e2691fd13a8 Mon Sep 17 00:00:00 2001
From: Matthew Hooker <mwhooker@gmail.com>
Date: Fri, 19 Jan 2018 19:55:27 -0800
Subject: [PATCH] pass context through step.run again

---
 helper/multistep/basic_runner.go |  2 +-
 helper/multistep/debug_runner.go |  3 ++-
 helper/multistep/multistep.go    |  4 ++-
 helper/multistep/statebag.go     | 46 +++++---------------------------
 4 files changed, 13 insertions(+), 42 deletions(-)

diff --git a/helper/multistep/basic_runner.go b/helper/multistep/basic_runner.go
index 0c928321e..eb5018d36 100644
--- a/helper/multistep/basic_runner.go
+++ b/helper/multistep/basic_runner.go
@@ -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 c2f1e8a84..9efb6d998 100644
--- a/helper/multistep/statebag.go
+++ b/helper/multistep/statebag.go
@@ -1,9 +1,6 @@
 package multistep
 
-import (
-	"context"
-	"sync"
-)
+import "sync"
 
 // Add context to state bag to prevent changing step signature
 
@@ -13,8 +10,6 @@ type StateBag interface {
 	Get(string) interface{}
 	GetOk(string) (interface{}, bool)
 	Put(string, interface{})
-	Context() context.Context
-	WithContext(context.Context) StateBag
 }
 
 // BasicStateBag implements StateBag by using a normal map underneath
@@ -22,13 +17,7 @@ type StateBag interface {
 type BasicStateBag struct {
 	data map[string]interface{}
 	l    sync.RWMutex
-	ctx  context.Context
-}
-
-func NewBasicStateBag() *BasicStateBag {
-	b := new(BasicStateBag)
-	b.data = make(map[string]interface{})
-	return b
+	once sync.Once
 }
 
 func (b *BasicStateBag) Get(k string) interface{} {
@@ -48,32 +37,11 @@ func (b *BasicStateBag) Put(k string, v interface{}) {
 	b.l.Lock()
 	defer b.l.Unlock()
 
+	// Make sure the map is initialized one time, on write
+	b.once.Do(func() {
+		b.data = make(map[string]interface{})
+	})
+
 	// Write the data
 	b.data[k] = v
 }
-
-func (b *BasicStateBag) Context() context.Context {
-	if b.ctx != nil {
-		return b.ctx
-	}
-	return context.Background()
-}
-
-// WithContext returns a copy of BasicStateBag with the provided context
-// We copy the state bag
-func (b *BasicStateBag) WithContext(ctx context.Context) *BasicStateBag {
-	if ctx == nil {
-		panic("nil context")
-	}
-	// read lock because copying is a read operation
-	b.l.RLock()
-	defer b.l.RUnlock()
-
-	b2 := NewBasicStateBag()
-
-	for k, v := range b.data {
-		b2.data[k] = v
-	}
-	b2.ctx = ctx
-	return b2
-}