down: refactor image pruning

Signed-off-by: Milas Bowman <milas.bowman@docker.com>

pkg/compose/down.go

@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/compose-spec/compose-go/types"
-	"github.com/distribution/distribution/v3/reference"
 	moby "github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/errdefs"
@@ -32,7 +31,6 @@ import (
 
 	"github.com/docker/compose/v2/pkg/api"
 	"github.com/docker/compose/v2/pkg/progress"
-	"github.com/docker/compose/v2/pkg/utils"
 )
 
 type downOp func() error
@@ -125,7 +123,12 @@ func (s *composeService) ensureVolumesDown(ctx context.Context, project *types.P
 }
 
 func (s *composeService) ensureImagesDown(ctx context.Context, project *types.Project, options api.DownOptions, w progress.Writer) ([]downOp, error) {
-	images, err := s.getServiceImagesToRemove(ctx, options, project)
+	imagePruner := NewImagePruner(s.apiClient(), project)
+	pruneOpts := ImagePruneOptions{
+		Mode:          ImagePruneMode(options.Images),
+		RemoveOrphans: options.RemoveOrphans,
+	}
+	images, err := imagePruner.ImagesToPrune(ctx, pruneOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -201,110 +204,6 @@ func (s *composeService) removeNetwork(ctx context.Context, name string, w progr
 	return nil
 }
 
-//nolint:gocyclo
-func (s *composeService) getServiceImagesToRemove(ctx context.Context, options api.DownOptions, project *types.Project) ([]string, error) {
-	if options.Images == "" {
-		return nil, nil
-	}
-
-	var localServiceImages []string
-	var imagesToRemove []string
-	addImageToRemove := func(img string, checkExistence bool) {
-		// since some references come from user input (service.image) and some
-		// come from the engine API, we standardize them, opting for the
-		// familiar name format since they'll also be displayed in the CLI
-		ref, err := reference.ParseNormalizedNamed(img)
-		if err != nil {
-			return
-		}
-		ref = reference.TagNameOnly(ref)
-		img = reference.FamiliarString(ref)
-		if utils.StringContains(imagesToRemove, img) {
-			return
-		}
-
-		if checkExistence {
-			_, _, err := s.apiClient().ImageInspectWithRaw(ctx, img)
-			if errdefs.IsNotFound(err) {
-				// err on the side of caution: only skip if we successfully
-				// queried the API and got back a definitive "not exists"
-				return
-			}
-		}
-
-		imagesToRemove = append(imagesToRemove, img)
-	}
-
-	imageListOpts := moby.ImageListOptions{
-		Filters: filters.NewArgs(
-			projectFilter(project.Name),
-			// TODO(milas): we should really clean up the dangling images as
-			// well (historically we have NOT); need to refactor this to handle
-			// it gracefully without producing confusing CLI output, i.e. we
-			// do not want to print out a bunch of untagged/dangling image IDs,
-			// they should be grouped into a logical operation for the relevant
-			// service
-			filters.Arg("dangling", "false"),
-		),
-	}
-	projectImages, err := s.apiClient().ImageList(ctx, imageListOpts)
-	if err != nil {
-		return nil, err
-	}
-
-	// 1. Remote / custom-named images - only deleted on `--rmi="all"`
-	for _, service := range project.Services {
-		if service.Image == "" {
-			localServiceImages = append(localServiceImages, service.Name)
-			continue
-		}
-
-		if options.Images == "all" {
-			addImageToRemove(service.Image, true)
-		}
-	}
-
-	// 2. *LABELED* Locally-built images with implicit image names
-	//
-	// If `--remove-orphans` is being used, then ALL images for the project
-	// will be selected for removal. Otherwise, only those that match a known
-	// service based on the loaded project will be included.
-	for _, img := range projectImages {
-		if len(img.RepoTags) == 0 {
-			// currently, we're only removing the tagged references, but
-			// if we start removing the dangling images and grouping by
-			// service, we can remove this (and should rely on `Image::ID`)
-			continue
-		}
-
-		shouldRemove := options.RemoveOrphans
-		for _, service := range localServiceImages {
-			if img.Labels[api.ServiceLabel] == service {
-				shouldRemove = true
-				break
-			}
-		}
-
-		if shouldRemove {
-			addImageToRemove(img.RepoTags[0], false)
-		}
-	}
-
-	// 3. *UNLABELED* Locally-built images with implicit image names
-	//
-	// This is a fallback for (2) to handle images built by previous
-	// versions of Compose, which did not label their built images.
-	for _, serviceName := range localServiceImages {
-		service, err := project.GetService(serviceName)
-		if err != nil || service.Image != "" {
-			continue
-		}
-		imgName := api.GetImageNameOrDefault(service, project.Name)
-		addImageToRemove(imgName, true)
-	}
-	return imagesToRemove, nil
-}
-
 func (s *composeService) removeImage(ctx context.Context, image string, w progress.Writer) error {
 	id := fmt.Sprintf("Image %s", image)
 	w.Event(progress.NewEvent(id, progress.Working, "Removing"))

pkg/compose/down_test.go

@@ -192,10 +192,11 @@ func TestDownRemoveImages(t *testing.T) {
 	}, nil).AnyTimes()
 
 	imagesToBeInspected := map[string]bool{
-		"local-named-image:latest":               true,
+		"testproject-local-anonymous":     true,
-		"remote-image:latest":                    true,
+		"local-named-image":               true,
-		"testproject-no-images-anonymous:latest": false,
+		"remote-image":                    true,
-		"missing-named-image:latest":             false,
+		"testproject-no-images-anonymous": false,
+		"missing-named-image":             false,
 	}
 	for img, exists := range imagesToBeInspected {
 		var resp moby.ImageInspect
@@ -278,7 +279,7 @@ func TestDownRemoveImages_NoLabel(t *testing.T) {
 		),
 	}).Return(nil, nil)
 
-	api.EXPECT().ImageInspectWithRaw(gomock.Any(), "testproject-service1:latest").
+	api.EXPECT().ImageInspectWithRaw(gomock.Any(), "testproject-service1").
 		Return(moby.ImageInspect{}, nil, nil)
 
 	api.EXPECT().ContainerStop(gomock.Any(), "123", nil).Return(nil)

pkg/compose/image_pruner.go

@@ -0,0 +1,202 @@
+package compose
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"sync"
+
+	"github.com/compose-spec/compose-go/types"
+	"github.com/distribution/distribution/v3/reference"
+	moby "github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/filters"
+	"github.com/docker/docker/client"
+	"github.com/docker/docker/errdefs"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/docker/compose/v2/pkg/api"
+)
+
+// ImagePruneMode controls how aggressively images associated with the project
+// are removed from the engine.
+type ImagePruneMode string
+
+const (
+	// ImagePruneNone indicates that no project images should be removed.
+	ImagePruneNone ImagePruneMode = ""
+	// ImagePruneLocal indicates that only images built locally by Compose
+	// should be removed.
+	ImagePruneLocal ImagePruneMode = "local"
+	// ImagePruneAll indicates that all project-associated images, including
+	// remote images should be removed.
+	ImagePruneAll ImagePruneMode = "all"
+)
+
+// ImagePruneOptions controls the behavior of image pruning.
+type ImagePruneOptions struct {
+	Mode ImagePruneMode
+
+	// RemoveOrphans will result in the removal of images that were built for
+	// the project regardless of whether they are for a known service if true.
+	RemoveOrphans bool
+}
+
+// ImagePruner handles image removal during Compose `down` operations.
+type ImagePruner struct {
+	client  client.ImageAPIClient
+	project *types.Project
+}
+
+// NewImagePruner creates an ImagePruner object for a project.
+func NewImagePruner(imageClient client.ImageAPIClient, project *types.Project) *ImagePruner {
+	return &ImagePruner{
+		client:  imageClient,
+		project: project,
+	}
+}
+
+// ImagesToPrune returns the set of images that should be removed.
+func (p *ImagePruner) ImagesToPrune(ctx context.Context, opts ImagePruneOptions) ([]string, error) {
+	if opts.Mode == ImagePruneNone {
+		return nil, nil
+	} else if opts.Mode != ImagePruneLocal && opts.Mode != ImagePruneAll {
+		return nil, fmt.Errorf("unsupported image prune mode: %s", opts.Mode)
+	}
+
+	var images []string
+
+	if opts.Mode == ImagePruneAll {
+		namedImages, err := p.namedImages(ctx)
+		if err != nil {
+			return nil, err
+		}
+		images = append(images, namedImages...)
+	}
+
+	projectImages, err := p.builtImagesForProject(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, img := range projectImages {
+		if len(img.RepoTags) == 0 {
+			// currently, we're only removing the tagged references, but
+			// if we start removing the dangling images and grouping by
+			// service, we can remove this (and should rely on `Image::ID`)
+			continue
+		}
+
+		removeImage := opts.RemoveOrphans
+		if !removeImage {
+			service, err := p.project.GetService(img.Labels[api.ServiceLabel])
+			if err == nil && (opts.Mode == ImagePruneAll || service.Image == "") {
+				removeImage = true
+			}
+		}
+
+		if removeImage {
+			images = append(images, img.RepoTags[0])
+		}
+	}
+
+	fallbackImages, err := p.unlabeledLocalImages(ctx)
+	if err != nil {
+		return nil, err
+	}
+	images = append(images, fallbackImages...)
+
+	images = normalizeAndDedupeImages(images)
+	return images, nil
+}
+
+func (p *ImagePruner) builtImagesForProject(ctx context.Context) ([]moby.ImageSummary, error) {
+	imageListOpts := moby.ImageListOptions{
+		Filters: filters.NewArgs(
+			projectFilter(p.project.Name),
+			// TODO(milas): we should really clean up the dangling images as
+			// well (historically we have NOT); need to refactor this to handle
+			// it gracefully without producing confusing CLI output, i.e. we
+			// do not want to print out a bunch of untagged/dangling image IDs,
+			// they should be grouped into a logical operation for the relevant
+			// service
+			filters.Arg("dangling", "false"),
+		),
+	}
+	projectImages, err := p.client.ImageList(ctx, imageListOpts)
+	if err != nil {
+		return nil, err
+	}
+	return projectImages, nil
+}
+
+func (p *ImagePruner) namedImages(ctx context.Context) ([]string, error) {
+	var images []string
+	for _, service := range p.project.Services {
+		if service.Image == "" {
+			continue
+		}
+		images = append(images, service.Image)
+	}
+	return p.filterImagesByExistence(ctx, images)
+}
+
+func (p *ImagePruner) filterImagesByExistence(ctx context.Context, imageNames []string) ([]string, error) {
+	var mu sync.Mutex
+	var ret []string
+
+	eg, ctx := errgroup.WithContext(ctx)
+	for _, img := range imageNames {
+		img := img
+		eg.Go(func() error {
+			_, _, err := p.client.ImageInspectWithRaw(ctx, img)
+			if errdefs.IsNotFound(err) {
+				// err on the side of caution: only skip if we successfully
+				// queried the API and got back a definitive "not exists"
+				return nil
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			ret = append(ret, img)
+			return nil
+		})
+	}
+
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
+	return ret, nil
+}
+
+func (p *ImagePruner) unlabeledLocalImages(ctx context.Context) ([]string, error) {
+	var images []string
+	for _, service := range p.project.Services {
+		if service.Image != "" {
+			continue
+		}
+		img := api.GetImageNameOrDefault(service, p.project.Name)
+		images = append(images, img)
+	}
+	return p.filterImagesByExistence(ctx, images)
+}
+
+func normalizeAndDedupeImages(images []string) []string {
+	seen := make(map[string]struct{}, len(images))
+	for _, img := range images {
+		// since some references come from user input (service.image) and some
+		// come from the engine API, we standardize them, opting for the
+		// familiar name format since they'll also be displayed in the CLI
+		ref, err := reference.ParseNormalizedNamed(img)
+		if err == nil {
+			ref = reference.TagNameOnly(ref)
+			img = reference.FamiliarString(ref)
+		}
+		seen[img] = struct{}{}
+	}
+	ret := make([]string, 0, len(seen))
+	for v := range seen {
+		ret = append(ret, v)
+	}
+	sort.Strings(ret)
+	return ret
+}

pkg/e2e/build_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	"gotest.tools/v3/assert"
 	"gotest.tools/v3/icmd"
 )
@@ -211,6 +212,10 @@ func TestBuildImageDependencies(t *testing.T) {
 	doTest := func(t *testing.T, cli *CLI) {
 		resetState := func() {
 			cli.RunDockerComposeCmd(t, "down", "--rmi=all", "-t=0")
+			res := cli.RunDockerOrExitError(t, "image", "rm", "build-dependencies-service")
+			if res.Error != nil {
+				require.Contains(t, res.Stderr(), `Error: No such image: build-dependencies-service`)
+			}
 		}
 		resetState()
 		t.Cleanup(resetState)
@@ -229,6 +234,15 @@ func TestBuildImageDependencies(t *testing.T) {
 			"image", "inspect", "--format={{ index .RepoTags 0 }}",
 			"build-dependencies-service")
 		res.Assert(t, icmd.Expected{Out: "build-dependencies-service:latest"})
+
+		res = cli.RunDockerComposeCmd(t, "down", "-t0", "--rmi=all", "--remove-orphans")
+		t.Log(res.Combined())
+
+		res = cli.RunDockerOrExitError(t, "image", "inspect", "build-dependencies-service")
+		res.Assert(t, icmd.Expected{
+			ExitCode: 1,
+			Err:      "Error: No such image: build-dependencies-service",
+		})
 	}
 
 	t.Run("ClassicBuilder", func(t *testing.T) {