From cf06204c823eacd34b00bced36359e1765d2c439 Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Tue, 17 Jan 2017 16:41:59 -0800
Subject: [PATCH] Add CreatedResources.Remove and use it

---
 client/driver/docker.go      | 13 ++++++------
 client/driver/driver.go      | 17 +++++++++++++++
 client/driver/driver_test.go | 40 +++++++++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/client/driver/docker.go b/client/driver/docker.go
index 2932e2647..08217d81a 100644
--- a/client/driver/docker.go
+++ b/client/driver/docker.go
@@ -504,19 +504,18 @@ func (d *DockerDriver) Cleanup(_ *ExecContext, res *CreatedResources) error {
 	for key, resources := range res.Resources {
 		switch key {
 		case dockerImageResKey:
-			// Remove and only add back images that failed to be
-			// removed in a retryable way
-			delete(res.Resources, key)
 			for _, value := range resources {
-				if err := d.cleanupImage(value); err != nil {
+				err := d.cleanupImage(value)
+				if err != nil {
 					if structs.IsRecoverable(err) {
 						retry = true
 					}
 					merr.Errors = append(merr.Errors, err)
-
-					// Re-add all failures to the map
-					res.Add(key, value)
+					continue
 				}
+
+				// Remove cleaned image from resources
+				res.Remove(dockerImageResKey, value)
 			}
 		default:
 			d.logger.Printf("[WARN] driver.docker: unknown resource to cleanup: %q", key)
diff --git a/client/driver/driver.go b/client/driver/driver.go
index 3c08455e3..52db1f485 100644
--- a/client/driver/driver.go
+++ b/client/driver/driver.go
@@ -85,6 +85,23 @@ func (r *CreatedResources) Add(k, v string) {
 	return
 }
 
+// Remove a resource. Return true if removed, otherwise false.
+//
+// Removes the entire key if the needle is the last value in the list.
+func (r *CreatedResources) Remove(k, needle string) bool {
+	haystack := r.Resources[k]
+	for i, item := range haystack {
+		if item == needle {
+			r.Resources[k] = append(haystack[:i], haystack[i+1:]...)
+			if len(r.Resources[k]) == 0 {
+				delete(r.Resources, k)
+			}
+			return true
+		}
+	}
+	return false
+}
+
 // Copy returns a new deep copy of CreatedResrouces.
 func (r *CreatedResources) Copy() *CreatedResources {
 	newr := CreatedResources{
diff --git a/client/driver/driver_test.go b/client/driver/driver_test.go
index 51e0dbe58..6cca2b4c7 100644
--- a/client/driver/driver_test.go
+++ b/client/driver/driver_test.go
@@ -257,7 +257,7 @@ func TestMapMergeStrStr(t *testing.T) {
 	}
 }
 
-func TestCreatedResources(t *testing.T) {
+func TestCreatedResources_AddMerge(t *testing.T) {
 	res1 := NewCreatedResources()
 	res1.Add("k1", "v1")
 	res1.Add("k1", "v2")
@@ -295,3 +295,41 @@ func TestCreatedResources(t *testing.T) {
 		t.Fatalf("3.  %#v != expected %#v", res1.Resources, expected)
 	}
 }
+
+func TestCreatedResources_CopyRemove(t *testing.T) {
+	res1 := NewCreatedResources()
+	res1.Add("k1", "v1")
+	res1.Add("k1", "v2")
+	res1.Add("k1", "v3")
+	res1.Add("k2", "v1")
+
+	// Assert Copy creates a deep copy
+	res2 := res1.Copy()
+
+	if !reflect.DeepEqual(res1, res2) {
+		t.Fatalf("%#v != %#v", res1, res2)
+	}
+
+	// Assert removing v1 from k1 returns true and updates Resources slice
+	if removed := res2.Remove("k1", "v1"); !removed {
+		t.Fatalf("expected v1 to be removed: %#v", res2)
+	}
+
+	if expected := []string{"v2", "v3"}; !reflect.DeepEqual(expected, res2.Resources["k1"]) {
+		t.Fatalf("unpexpected list for k1: %#v", res2.Resources["k1"])
+	}
+
+	// Assert removing the only value from a key removes the key
+	if removed := res2.Remove("k2", "v1"); !removed {
+		t.Fatalf("expected v1 to be removed from k2: %#v", res2.Resources)
+	}
+
+	if _, found := res2.Resources["k2"]; found {
+		t.Fatalf("k2 should have been removed from Resources: %#v", res2.Resources)
+	}
+
+	// Make sure res1 wasn't updated
+	if reflect.DeepEqual(res1, res2) {
+		t.Fatalf("res1 should not equal res2: #%v", res1)
+	}
+}