fix(compose): improve image pruning during uninstall

Either remove all unused images (when the "dangling" filter is false),
or remove only images referenced by apps being uninstalled, unless they
are still used by one or more containers.

Update the integration test to verify that no images remain after
uninstall with pruning enabled.

Signed-off-by: Mike Sul <mike.sul@foundries.io>

Author: mike-sul

pkg/compose/uninstall.go

@@ -3,13 +3,19 @@ package compose
 import (
 	"context"
 	"errors"
-	"github.com/docker/docker/api/types/filters"
+	"fmt"
 	"os"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/filters"
+	dockerClient "github.com/docker/docker/client"
 )
 
 type (
 	UninstallOpts struct {
-		Prune bool
+		Prune                 bool
+		RemoveAllUnusedImages bool
 	}
 	UninstallOpt func(*UninstallOpts)
 )
@@ -18,9 +24,12 @@ var (
 	ErrUninstallRunningApps = errors.New("failed to uninstall apps: some apps are still running, please stop them first")
 )
 
-func WithImagePruning() UninstallOpt {
+func WithImagePruning(allUnused ...bool) UninstallOpt {
 	return func(opts *UninstallOpts) {
 		opts.Prune = true
+		if len(allUnused) > 0 {
+			opts.RemoveAllUnusedImages = allUnused[0]
+		}
 	}
 }
 
@@ -64,16 +73,61 @@ func UninstallApps(ctx context.Context, cfg *Config, appRefs []string, options .
 	}
 
 	if opts.Prune {
-		cli, errClient := GetDockerClient(cfg.DockerHost)
+		dockerClient, errClient := GetDockerClient(cfg.DockerHost)
 		if errClient != nil {
 			return errClient
 		}
-		// Prune only dangling images.
-		// The dangling images are the ones that are not tagged and not referenced by any container.
-		// TODO: consider pruning volumes and networks if needed.
-		// TODO: consider pruning only those images that are related to the uninstalled apps,
-		//       otherwise it prunes all dangling images including those that are not managed by composectl
-		_, err = cli.ImagesPrune(ctx, filters.NewArgs(filters.Arg("dangling", "true")))
+
+		if opts.RemoveAllUnusedImages {
+			// Prune all unused images, including dangling and non-dangling ones.
+			// The non-dangling images are the ones that are not referenced by any container, but they can still be tagged.
+			_, err = dockerClient.ImagesPrune(ctx, filters.NewArgs(filters.Arg("dangling", "false")))
+			return err
+		}
+
+		// Get all containers to find out which images are still referenced by the containers, so that we don't prune those images.
+		allContainers, errCtrList := dockerClient.ContainerList(ctx, container.ListOptions{})
+		if errCtrList != nil {
+			return fmt.Errorf("failed to list containers: %w", errCtrList)
+		}
+		imagesWithContainer := make(map[string]types.Container)
+		for _, container := range allContainers {
+			imagesWithContainer[container.Image] = container
+		}
+
+		// Remove or untag images that thar are referenced by the compose apps to be uninstalled, but not referenced by any container.
+		// We cannot even untag images that are still referenced by the containers, if we do then
+		// composectl will think the app that uses that image is uninstalled.
+		for _, app := range status.Apps {
+			for _, imageRoot := range app.GetComposeRoot().Children {
+				if _, hasContainer := imagesWithContainer[imageRoot.Ref()]; !hasContainer {
+					removeImage(ctx, dockerClient, imageRoot)
+				}
+			}
+		}
+		// Prune dangling images, which are the ones that are not tagged and not referenced by any container.
+		_, err = dockerClient.ImagesPrune(ctx, filters.NewArgs(filters.Arg("dangling", "true")))
 	}
 	return err
 }
+
+func removeImage(ctx context.Context, client *dockerClient.Client, imageRoot *TreeNode) {
+	imageNode := imageRoot
+	for {
+		// remove image (untag if more than 2 references) referenced by the URI with a digest
+		// Ignore error because the image may have already been removed as a child of another image,
+		// or the image may be referenced by other compose apps that are running or not uninstalled.
+		_, _ = client.ImageRemove(ctx, imageNode.Ref(), types.ImageRemoveOptions{Force: false, PruneChildren: true})
+		if imageRef, refParseErr := ParseImageRef(imageNode.Ref()); refParseErr == nil {
+			// remove image (untag if more than 2 references) referenced by the URI with a tag
+			// Ignore error because the image may have already been removed as a child of another image,
+			// or the image may be referenced by other compose apps that are running or not uninstalled.
+			_, _ = client.ImageRemove(ctx, imageRef.GetTagRef(), types.ImageRemoveOptions{Force: false, PruneChildren: true})
+		}
+		if imageNode.Type == BlobTypeImageManifest || len(imageNode.Children) == 0 {
+			break
+		}
+		// The image root points to the image index, which was removed, now remove the image manifest
+		imageNode = imageNode.Children[0]
+	}
+}

test/integration/update_test.go

@@ -2,6 +2,7 @@ package e2e_tests
 
 import (
 	"context"
+	"github.com/docker/docker/api/types"
 	"testing"
 
 	"github.com/foundriesio/composeapp/pkg/compose"
@@ -621,21 +622,35 @@ services:
 	f.Check(t, updateRunner.Start(ctx))
 	// Complete with pruning to remove the second app
 	f.Check(t, updateRunner.Complete(ctx, update.CompleteWithPruning()))
-
 	appsStatus, err = compose.CheckAppsStatus(ctx, cfg, nil)
 	f.Check(t, err)
-	if !appsStatus.AreRunning() {
-		t.Fatal("app is expected to be running")
-	}
 	if len(appsStatus.Apps) > 1 || len(appsStatus.Apps) == 0 {
 		t.Fatalf("only one app is expected to be running, found %d", len(appsStatus.Apps))
 	}
 	if appsStatus.Apps[0].Ref().String() != appURIs[0] {
 		t.Fatalf("expected app URI %s, found %s", appURIs[0], appsStatus.Apps[0].Ref().String())
 	}
+	appsStatus, err = compose.CheckAppsStatus(ctx, cfg, oneAppURI)
+	f.Check(t, err)
+	if !appsStatus.AreRunning() {
+		t.Fatalf("app is expected to be running; uri=%s", oneAppURI[0])
+	}
 
 	// stop, uninstall, and remove all apps
 	f.Check(t, compose.StopApps(ctx, cfg, oneAppURI))
 	f.Check(t, compose.UninstallApps(ctx, cfg, oneAppURI, compose.WithImagePruning()))
 	f.Check(t, compose.RemoveApps(ctx, cfg, oneAppURI))
+
+	// check if no images are left after pruning in the docker storage
+	cli, err := compose.GetDockerClient("")
+	f.Check(t, err)
+	defer cli.Close()
+	images, err := cli.ImageList(ctx, types.ImageListOptions{All: true})
+	f.Check(t, err)
+	if len(images) != 0 {
+		for _, img := range images {
+			t.Logf("unexpected image left after pruning: ID=%s, RepoTags=%v, RepoDigests=%v\n", img.ID, img.RepoTags, img.RepoDigests)
+		}
+		t.Fatalf("no images are expected to be left after pruning, found %d", len(images))
+	}
 }