build: label built images for reliable cleanup on down

pkg/api/labels.go

@@ -47,14 +47,14 @@ const (
 	OneoffLabel = "com.docker.compose.oneoff"
 	// SlugLabel stores unique slug used for one-off container identity
 	SlugLabel = "com.docker.compose.slug"
-	// ImageNameLabel stores the content of the image section in the compose file
-	ImageNameLabel = "com.docker.compose.image_name"
 	// ImageDigestLabel stores digest of the container image used to run service
 	ImageDigestLabel = "com.docker.compose.image"
 	// DependenciesLabel stores service dependencies
 	DependenciesLabel = "com.docker.compose.depends_on"
-	// VersionLabel stores the compose tool version used to run application
+	// VersionLabel stores the compose tool version used to build/run application
 	VersionLabel = "com.docker.compose.version"
+	// ImageBuilderLabel stores the builder (classic or BuildKit) used to produce the image.
+	ImageBuilderLabel = "com.docker.compose.image.builder"
 )
 
 // ComposeVersion is the compose tool version as declared by label VersionLabel

pkg/compose/build.go

@@ -139,7 +139,6 @@ func (s *composeService) ensureImagesExists(ctx context.Context, project *types.
 				project.Services[i].Labels = types.Labels{}
 			}
 			project.Services[i].CustomLabels.Add(api.ImageDigestLabel, digest)
-			project.Services[i].CustomLabels.Add(api.ImageNameLabel, service.Image)
 		}
 	}
 	return nil
@@ -192,7 +191,6 @@ func (s *composeService) getLocalImagesDigests(ctx context.Context, project *typ
 		digest, ok := images[imgName]
 		if ok {
 			project.Services[i].CustomLabels.Add(api.ImageDigestLabel, digest)
-			project.Services[i].CustomLabels.Add(api.ImageNameLabel, project.Services[i].Image)
 		}
 	}
 
@@ -263,6 +261,8 @@ func (s *composeService) toBuildOptions(project *types.Project, service types.Se
 		tags = append(tags, service.Build.Tags...)
 	}
 
+	imageLabels := getImageBuildLabels(project, service)
+
 	return build.Options{
 		Inputs: build.Inputs{
 			ContextPath:    service.Build.Context,
@@ -277,7 +277,7 @@ func (s *composeService) toBuildOptions(project *types.Project, service types.Se
 		Target:      service.Build.Target,
 		Exports:     []bclient.ExportEntry{{Type: "image", Attrs: map[string]string{}}},
 		Platforms:   plats,
-		Labels:      service.Build.Labels,
+		Labels:      imageLabels,
 		NetworkMode: service.Build.Network,
 		ExtraHosts:  service.Build.ExtraHosts.AsList(),
 		Session:     sessionConfig,
@@ -327,7 +327,6 @@ func sshAgentProvider(sshKeys types.SSHConfig) (session.Attachable, error) {
 }
 
 func addSecretsConfig(project *types.Project, service types.ServiceConfig) (session.Attachable, error) {
-
 	var sources []secretsprovider.Source
 	for _, secret := range service.Build.Secrets {
 		config := project.Secrets[secret.Source]
@@ -352,3 +351,17 @@ func addSecretsConfig(project *types.Project, service types.ServiceConfig) (sess
 	}
 	return secretsprovider.NewSecretProvider(store), nil
 }
+
+func getImageBuildLabels(project *types.Project, service types.ServiceConfig) types.Labels {
+	ret := make(types.Labels)
+	if service.Build != nil {
+		for k, v := range service.Build.Labels {
+			ret.Add(k, v)
+		}
+	}
+
+	ret.Add(api.VersionLabel, api.ComposeVersion)
+	ret.Add(api.ProjectLabel, project.Name)
+	ret.Add(api.ServiceLabel, service.Name)
+	return ret
+}

pkg/compose/build_buildkit.go

@@ -25,6 +25,8 @@ import (
 	"github.com/docker/buildx/build"
 	"github.com/docker/buildx/driver"
 	xprogress "github.com/docker/buildx/util/progress"
+
+	"github.com/docker/compose/v2/pkg/api"
 )
 
 func (s *composeService) doBuildBuildkit(ctx context.Context, project *types.Project, opts map[string]build.Options, mode string) (map[string]string, error) {
@@ -47,6 +49,15 @@ func (s *composeService) doBuildBuildkit(ctx context.Context, project *types.Pro
 	defer cancel()
 	w := xprogress.NewPrinter(progressCtx, s.stdout(), os.Stdout, mode)
 
+	for k := range opts {
+		if opts[k].Labels == nil {
+			opt := opts[k]
+			opt.Labels = make(map[string]string)
+			opts[k] = opt
+		}
+		opts[k].Labels[api.ImageBuilderLabel] = "buildkit"
+	}
+
 	// We rely on buildx "docker" builder integrated in docker engine, so don't need a DockerAPI here
 	response, err := build.Build(ctx, driverInfo, opts, nil, filepath.Dir(s.configFile().Filename), w)
 	errW := w.Wait()

pkg/compose/build_classic.go

@@ -89,6 +89,11 @@ func (s *composeService) doBuildClassicSimpleImage(ctx context.Context, options
 		}
 	}
 
+	if options.Labels == nil {
+		options.Labels = make(map[string]string)
+	}
+	options.Labels[api.ImageBuilderLabel] = "classic"
+
 	switch {
 	case isLocalDir(specifiedContext):
 		contextDir, relDockerfile, err = build.GetContextFromLocalDir(specifiedContext, dockerfileName)

pkg/compose/compose.go

@@ -30,11 +30,12 @@ import (
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/config/configfile"
 	"github.com/docker/cli/cli/streams"
-	"github.com/docker/compose/v2/pkg/api"
 	moby "github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/client"
 	"github.com/pkg/errors"
+
+	"github.com/docker/compose/v2/pkg/api"
 )
 
 // NewComposeService create a local implementation of the compose.Service API
@@ -130,13 +131,9 @@ func (s *composeService) projectFromName(containers Containers, projectName stri
 		serviceLabel := c.Labels[api.ServiceLabel]
 		_, ok := set[serviceLabel]
 		if !ok {
-			serviceImage := c.Image
-			if serviceNameFromLabel, ok := c.Labels[api.ImageNameLabel]; ok {
-				serviceImage = serviceNameFromLabel
-			}
 			set[serviceLabel] = &types.ServiceConfig{
 				Name:   serviceLabel,
-				Image:  serviceImage,
+				Image:  c.Image,
 				Labels: c.Labels,
 			}
 		}

pkg/compose/down.go

@@ -23,6 +23,7 @@ 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"
@@ -31,6 +32,7 @@ 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
@@ -86,7 +88,11 @@ func (s *composeService) down(ctx context.Context, projectName string, options a
 	ops := s.ensureNetworksDown(ctx, project, w)
 
 	if options.Images != "" {
-		ops = append(ops, s.ensureImagesDown(ctx, project, options, w)...)
+		imgOps, err := s.ensureImagesDown(ctx, project, options, w)
+		if err != nil {
+			return err
+		}
+		ops = append(ops, imgOps...)
 	}
 
 	if options.Volumes {
@@ -118,15 +124,20 @@ func (s *composeService) ensureVolumesDown(ctx context.Context, project *types.P
 	return ops
 }
 
-func (s *composeService) ensureImagesDown(ctx context.Context, project *types.Project, options api.DownOptions, w progress.Writer) []downOp {
+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)
+	if err != nil {
+		return nil, err
+	}
+
 	var ops []downOp
-	for image := range s.getServiceImagesToRemove(options, project) {
-		image := image
+	for i := range images {
+		img := images[i]
 		ops = append(ops, func() error {
-			return s.removeImage(ctx, image, w)
+			return s.removeImage(ctx, img, w)
 		})
 	}
-	return ops
+	return ops, nil
 }
 
 func (s *composeService) ensureNetworksDown(ctx context.Context, project *types.Project, w progress.Writer) []downOp {
@@ -190,17 +201,108 @@ func (s *composeService) removeNetwork(ctx context.Context, name string, w progr
 	return nil
 }
 
-func (s *composeService) getServiceImagesToRemove(options api.DownOptions, project *types.Project) map[string]struct{} {
-	images := map[string]struct{}{}
+//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 {
-		image, ok := service.Labels[api.ImageNameLabel] // Information on the compose file at the creation of the container
-		if !ok || (options.Images == "local" && image != "") {
+		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
 		}
-		image = api.GetImageNameOrDefault(service, project.Name)
-		images[image] = struct{}{}
+		imgName := api.GetImageNameOrDefault(service, project.Name)
+		addImageToRemove(imgName, true)
 	}
-	return images
+	return imagesToRemove, nil
 }
 
 func (s *composeService) removeImage(ctx context.Context, image string, w progress.Writer) error {