From 91608843a4f3ba0fec01d73ec3e93b4f03151a4d Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Fri, 3 Feb 2017 14:23:24 -0500
Subject: [PATCH] Add state locking in taint/untaint

---
 command/taint.go   | 21 +++++++++++++++++----
 command/untaint.go | 21 +++++++++++++++++----
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/command/taint.go b/command/taint.go
index 1d60729a6..6d652fdf4 100644
--- a/command/taint.go
+++ b/command/taint.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"strings"
 
+	"github.com/hashicorp/terraform/state"
 	"github.com/hashicorp/terraform/terraform"
 )
 
@@ -25,6 +26,7 @@ func (c *TaintCommand) Run(args []string) int {
 	cmdFlags.StringVar(&c.Meta.statePath, "state", DefaultStateFilename, "path")
 	cmdFlags.StringVar(&c.Meta.stateOutPath, "state-out", "", "path")
 	cmdFlags.StringVar(&c.Meta.backupPath, "backup", "", "path")
+	cmdFlags.BoolVar(&c.Meta.lockState, "state-lock", true, "lock state")
 	cmdFlags.Usage = func() { c.Ui.Error(c.Help()) }
 	if err := cmdFlags.Parse(args); err != nil {
 		return 1
@@ -64,14 +66,23 @@ func (c *TaintCommand) Run(args []string) int {
 	}
 
 	// Get the state
-	state, err := b.State()
+	st, err := b.State()
 	if err != nil {
 		c.Ui.Error(fmt.Sprintf("Failed to load state: %s", err))
 		return 1
 	}
 
+	if s, ok := st.(state.Locker); c.Meta.lockState && ok {
+		if err := s.Lock("taint"); err != nil {
+			c.Ui.Error(fmt.Sprintf("Failed to lock state: %s", err))
+			return 1
+		}
+
+		defer s.Unlock()
+	}
+
 	// Get the actual state structure
-	s := state.State()
+	s := st.State()
 	if s.Empty() {
 		if allowMissing {
 			return c.allowMissingExit(name, module)
@@ -129,11 +140,11 @@ func (c *TaintCommand) Run(args []string) int {
 	rs.Taint()
 
 	log.Printf("[INFO] Writing state output to: %s", c.Meta.StateOutPath())
-	if err := state.WriteState(s); err != nil {
+	if err := st.WriteState(s); err != nil {
 		c.Ui.Error(fmt.Sprintf("Error writing state file: %s", err))
 		return 1
 	}
-	if err := state.PersistState(); err != nil {
+	if err := st.PersistState(); err != nil {
 		c.Ui.Error(fmt.Sprintf("Error writing state file: %s", err))
 		return 1
 	}
@@ -166,6 +177,8 @@ Options:
                       modifying. Defaults to the "-state-out" path with
                       ".backup" extension. Set to "-" to disable backup.
 
+  -lock-state=true    Lock the state file when locking is supported.
+
   -module=path        The module path where the resource lives. By
                       default this will be root. Child modules can be specified
                       by names. Ex. "consul" or "consul.vpc" (nested modules).
diff --git a/command/untaint.go b/command/untaint.go
index 4be0590fa..925a73fc3 100644
--- a/command/untaint.go
+++ b/command/untaint.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"log"
 	"strings"
+
+	"github.com/hashicorp/terraform/state"
 )
 
 // UntaintCommand is a cli.Command implementation that manually untaints
@@ -51,14 +53,23 @@ func (c *UntaintCommand) Run(args []string) int {
 	}
 
 	// Get the state
-	state, err := b.State()
+	st, err := b.State()
 	if err != nil {
 		c.Ui.Error(fmt.Sprintf("Failed to load state: %s", err))
 		return 1
 	}
 
+	if s, ok := st.(state.Locker); c.Meta.lockState && ok {
+		if err := s.Lock("untaint"); err != nil {
+			c.Ui.Error(fmt.Sprintf("Failed to lock state: %s", err))
+			return 1
+		}
+
+		defer s.Unlock()
+	}
+
 	// Get the actual state structure
-	s := state.State()
+	s := st.State()
 	if s.Empty() {
 		if allowMissing {
 			return c.allowMissingExit(name, module)
@@ -116,11 +127,11 @@ func (c *UntaintCommand) Run(args []string) int {
 	rs.Untaint()
 
 	log.Printf("[INFO] Writing state output to: %s", c.Meta.StateOutPath())
-	if err := state.WriteState(s); err != nil {
+	if err := st.WriteState(s); err != nil {
 		c.Ui.Error(fmt.Sprintf("Error writing state file: %s", err))
 		return 1
 	}
-	if err := state.PersistState(); err != nil {
+	if err := st.PersistState(); err != nil {
 		c.Ui.Error(fmt.Sprintf("Error writing state file: %s", err))
 		return 1
 	}
@@ -153,6 +164,8 @@ Options:
                       modifying. Defaults to the "-state-out" path with
                       ".backup" extension. Set to "-" to disable backup.
 
+  -lock-state=true    Lock the state file when locking is supported.
+
   -module=path        The module path where the resource lives. By
                       default this will be root. Child modules can be specified
                       by names. Ex. "consul" or "consul.vpc" (nested modules).