2 files changed, 21 insertions, 36 deletions

diff --git a/libpod/runtime_ctr.go b/libpod/runtime_ctr.go
index 32902d9c6..8613950f6 100644
--- a/libpod/runtime_ctr.go
+++ b/libpod/runtime_ctr.go
@@ -1,7 +1,6 @@
 package libpod
 
 import (
-	"github.com/containers/storage"
 	spec "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -86,6 +85,12 @@ func (r *Runtime) RemoveContainer(c *Container, force bool) error {
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
+	return r.removeContainer(c, force)
+}
+
+// Internal function to remove a container
+// Locks the container, but does not lock the runtime
+func (r *Runtime) removeContainer(c *Container, force bool) error {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
@@ -210,30 +215,3 @@ func (r *Runtime) GetContainers(filters ...ContainerFilter) ([]*Container, error
 
 	return ctrsFiltered, nil
 }
-
-// getContainersWithImage returns a list of containers referencing imageID
-func (r *Runtime) getContainersWithImage(imageID string) ([]storage.Container, error) {
-	var matchingContainers []storage.Container
-	containers, err := r.store.Containers()
-	if err != nil {
-		return nil, err
-	}
-
-	for _, ctr := range containers {
-		if ctr.ImageID == imageID {
-			matchingContainers = append(matchingContainers, ctr)
-		}
-	}
-	return matchingContainers, nil
-}
-
-// removeMultipleContainers deletes a list of containers from the store
-// TODO refactor this to remove libpod Containers
-func (r *Runtime) removeMultipleContainers(containers []storage.Container) error {
-	for _, ctr := range containers {
-		if err := r.store.DeleteContainer(ctr.ID); err != nil {
-			return errors.Wrapf(err, "could not remove container %q", ctr)
-		}
-	}
-	return nil
-}
diff --git a/libpod/runtime_img.go b/libpod/runtime_img.go
index 2b18a5322..ba1cec96f 100644
--- a/libpod/runtime_img.go
+++ b/libpod/runtime_img.go
@@ -798,19 +798,26 @@ func (r *Runtime) RemoveImage(image *storage.Image, force bool) (string, error)
 		return "", ErrRuntimeStopped
 	}
 
-	containersWithImage, err := r.getContainersWithImage(image.ID)
+	// Get all containers, filter to only those using the image, and remove those containers
+	ctrs, err := r.state.AllContainers()
 	if err != nil {
-		return "", errors.Wrapf(err, "error getting containers for image %q", image.ID)
+		return "", err
+	}
+	imageCtrs := []*Container{}
+	for _, ctr := range ctrs {
+		if ctr.config.RootfsImageID == image.ID {
+			imageCtrs = append(imageCtrs, ctr)
+		}
 	}
-	if len(containersWithImage) > 0 && len(image.Names) <= 1 {
+	if len(imageCtrs) > 0 && len(image.Names) <= 1 {
 		if force {
-			if err := r.removeMultipleContainers(containersWithImage); err != nil {
-				return "", err
+			for _, ctr := range imageCtrs {
+				if err := r.removeContainer(ctr, true); err != nil {
+					return "", errors.Wrapf(err, "error removing image %s: container %s using image could not be removed", image.ID, ctr.ID())
+				}
 			}
 		} else {
-			for _, ctr := range containersWithImage {
-				return "", fmt.Errorf("Could not remove image %q (must force) - container %q is using its reference image", image.ID, ctr.ImageID)
-			}
+			return "", fmt.Errorf("could not remove image %s as it is being used by %d containers", image.ID, len(imageCtrs))
 		}
 	}