Make cache pruning more flexible

This commit implements 2 more CLI flags:
* --no-digest-only
  Only prune images without a digest specified ("fallback" images usually)
* --unreferenced-only
  Only prune downloads not referenced in any VM

Signed-off-by: Yury Bushmelev <jay4mail@gmail.com>

cmd/limactl/prune.go

@@ -4,6 +4,10 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/lima-vm/lima/pkg/downloader"
+	"github.com/lima-vm/lima/pkg/limayaml"
+	"github.com/lima-vm/lima/pkg/store"
+	"github.com/opencontainers/go-digest"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -16,15 +20,101 @@ func newPruneCommand() *cobra.Command {
 		RunE:              pruneAction,
 		ValidArgsFunction: cobra.NoFileCompletions,
 	}
+
+	pruneCommand.Flags().Bool("no-digest-only", false, "Only prune images without a digest specified (\"fallback\" images usually)")
+	pruneCommand.Flags().Bool("unreferenced-only", false, "Only prune downloads not referenced in any VM")
 	return pruneCommand
 }
 
 func pruneAction(cmd *cobra.Command, args []string) error {
+	pruneWithoutDigest, err := cmd.Flags().GetBool("no-digest-only")
+	if err != nil {
+		return err
+	}
+	pruneUnreferenced, err := cmd.Flags().GetBool("unreferenced-only")
+	if err != nil {
+		return err
+	}
+
+	if pruneWithoutDigest || pruneUnreferenced {
+		files, err := getReferencedDownloads()
+		if err != nil {
+			return err
+		}
+
+		cacheEntries, err := downloader.CachedDownloads(downloader.WithCache())
+		if err != nil {
+			return err
+		}
+
+		for _, entry := range cacheEntries {
+			entryFields := logrus.Fields{
+				"id":       entry.ID,
+				"location": entry.Location,
+			}
+
+			logrus.WithFields(entryFields).Debug("cache entry")
+
+			// check if the cache entry is referenced
+			if refFile, refFound := files[entry.ID]; refFound {
+				if refFile.Location != entry.Location { // sanity check
+					logrus.WithFields(logrus.Fields{
+						"id":                  entry.ID,
+						"location":            entry.Location,
+						"referenced_location": refFile.Location,
+					}).Warnf("Sanity check failed! URL mismatch")
+				}
+
+				if pruneWithoutDigest && refFile.Digest == "" {
+					// delete the fallback image entry (entry w/o digest) even if referenced
+					logrus.WithFields(entryFields).Infof("Deleting fallback entry")
+					if err := os.RemoveAll(entry.Path); err != nil {
+						logrus.WithError(err).WithFields(logrus.Fields{
+							"path": entry.Path,
+						}).Errorf("Cannot delete directory. Skipping...")
+					}
+				}
+			} else {
+				if pruneUnreferenced {
+					// delete the unreferenced cached entry
+					logrus.WithFields(entryFields).Infof("Deleting unreferenced entry")
+					if err := os.RemoveAll(entry.Path); err != nil {
+						logrus.WithError(err).WithFields(logrus.Fields{
+							"path": entry.Path,
+						}).Errorf("Cannot delete directory. Skipping...")
+					}
+				}
+			}
+		}
+		return nil
+	}
+
+	// prune everything if no options specified
 	ucd, err := os.UserCacheDir()
 	if err != nil {
 		return err
 	}
 	cacheDir := filepath.Join(ucd, "lima")
-	logrus.Infof("Pruning %q", cacheDir)
+	logrus.Infof("Pruning everything in %q", cacheDir)
 	return os.RemoveAll(cacheDir)
 }
+
+// Collect all downloads referenced in VM definitions and templates
+func getReferencedDownloads() (map[string]limayaml.File, error) {
+	digests := make(map[string]limayaml.File)
+
+	vmRefs, err := store.Downloads()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, f := range vmRefs {
+		d := digest.SHA256.FromString(f.Location).Encoded()
+		logrus.WithFields(logrus.Fields{
+			"id":       d,
+			"location": f.Location,
+		}).Debugf("referenced file")
+		digests[d] = f
+	}
+	return digests, nil
+}

pkg/downloader/downloader.go

@@ -39,6 +39,13 @@ type Result struct {
 	ValidatedDigest bool
 }
 
+type CacheEntry struct {
+	ID       string        // sha256 of the Location
+	Location string        // source URL of the entry
+	Path     string        // Absolute path of the entry
+	Digest   digest.Digest // Checksum of the entry's data (a file downloaded)
+}
+
 type options struct {
 	cacheDir       string // default: empty (disables caching)
 	decompress     bool   // default: false (keep compression)
@@ -483,3 +490,78 @@ func downloadHTTP(localPath, url string, description string, expectedDigest dige
 
 	return nil
 }
+
+// Read first *.digest file found and return the Digest object from its contents
+func guessDigestFromDir(dir string) (digest.Digest, error) {
+	dirList, err := os.ReadDir(dir)
+	if err != nil {
+		return "", err
+	}
+
+	for _, f := range dirList {
+		if strings.HasSuffix(f.Name(), ".digest") && !(f.IsDir()) {
+			data, err := os.ReadFile(filepath.Join(dir, f.Name()))
+			if err != nil {
+				return "", err
+			}
+			d, err := digest.Parse(string(data))
+			if err != nil {
+				return "", err
+			}
+			return d, nil
+		}
+	}
+	return "", nil
+}
+
+func CachedDownloads(opts ...Opt) ([]CacheEntry, error) {
+	// Get the default cache directory
+	var o options
+	for _, f := range opts {
+		if err := f(&o); err != nil {
+			return nil, err
+		}
+	}
+	if o.cacheDir == "" {
+		return []CacheEntry{}, nil // No cache => empty results
+	}
+
+	shaDir := filepath.Join(o.cacheDir, "download", "by-url-sha256")
+
+	shaDirList, err := os.ReadDir(shaDir)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	entries := []CacheEntry{}
+	for _, f := range shaDirList {
+		if strings.HasPrefix(f.Name(), ".") || strings.HasPrefix(f.Name(), "_") {
+			continue
+		}
+		if !f.IsDir() {
+			continue
+		}
+
+		path := filepath.Join(shaDir, f.Name())
+
+		x := CacheEntry{
+			ID:   f.Name(),
+			Path: path,
+		}
+
+		data, err := os.ReadFile(filepath.Join(path, "url"))
+		if err != nil {
+			logrus.WithError(err).Infof("Cannot read url of %q", f.Name())
+		}
+		x.Location = string(data)
+
+		if x.Digest, err = guessDigestFromDir(path); err != nil {
+			logrus.WithError(err).Infof("Cannot read digest of %q", f.Name())
+		}
+		entries = append(entries, x)
+	}
+	return entries, nil
+}

pkg/downloader/downloader_test.go

@@ -18,11 +18,27 @@ const (
 	dummyRemoteFileDigest = "sha256:58d2de96f9d91f0acd93cb1e28bf7c42fc86079037768d6aa63b4e7e7b3c9be0"
 )
 
+var dummyRemoteFileURLSHA256 = digest.SHA256.FromString(dummyRemoteFileURL).Encoded()
+
 func TestMain(m *testing.M) {
 	HideProgress = true
 	os.Exit(m.Run())
 }
 
+// helper to check if the file referred by the URL is cached
+func findInCache(url string, opts ...Opt) (*CacheEntry, error) {
+	c, err := CachedDownloads(opts...)
+	if err != nil {
+		return nil, err
+	}
+	for _, entry := range c {
+		if entry.Location == url {
+			return &entry, nil
+		}
+	}
+	return nil, nil
+}
+
 func TestDownloadRemote(t *testing.T) {
 	if testing.Short() {
 		t.Skip()
@@ -34,6 +50,11 @@ func TestDownloadRemote(t *testing.T) {
 			assert.NilError(t, err)
 			assert.Equal(t, StatusDownloaded, r.Status)
 
+			// check the cache is empty
+			c, err := CachedDownloads()
+			assert.NilError(t, err)
+			assert.Equal(t, len(c), 0)
+
 			// download again, make sure StatusSkippedIsReturned
 			r, err = Download(localPath, dummyRemoteFileURL)
 			assert.NilError(t, err)
@@ -49,6 +70,11 @@ func TestDownloadRemote(t *testing.T) {
 			assert.NilError(t, err)
 			assert.Equal(t, StatusDownloaded, r.Status)
 
+			// check the cache is empty
+			c, err := CachedDownloads()
+			assert.NilError(t, err)
+			assert.Equal(t, len(c), 0)
+
 			r, err = Download(localPath, dummyRemoteFileURL, WithExpectedDigest(dummyRemoteFileDigest))
 			assert.NilError(t, err)
 			assert.Equal(t, StatusSkipped, r.Status)
@@ -61,6 +87,12 @@ func TestDownloadRemote(t *testing.T) {
 		assert.NilError(t, err)
 		assert.Equal(t, StatusDownloaded, r.Status)
 
+		// check the download is in cache
+		entry, err := findInCache(dummyRemoteFileURL, WithCacheDir(cacheDir))
+		assert.NilError(t, err)
+		assert.Equal(t, entry.Digest.String(), dummyRemoteFileDigest)
+		assert.Equal(t, entry.ID, dummyRemoteFileURLSHA256)
+
 		r, err = Download(localPath, dummyRemoteFileURL, WithExpectedDigest(dummyRemoteFileDigest), WithCacheDir(cacheDir))
 		assert.NilError(t, err)
 		assert.Equal(t, StatusSkipped, r.Status)
@@ -79,6 +111,12 @@ func TestDownloadRemote(t *testing.T) {
 		assert.NilError(t, err)
 		assert.Equal(t, StatusDownloaded, r.Status)
 
+		// check the download is in cache
+		entry, err := findInCache(dummyRemoteFileURL, WithCacheDir(cacheDir))
+		assert.NilError(t, err)
+		assert.Equal(t, entry.Digest.String(), dummyRemoteFileDigest)
+		assert.Equal(t, entry.ID, dummyRemoteFileURLSHA256)
+
 		r, err = Download("", dummyRemoteFileURL, WithExpectedDigest(dummyRemoteFileDigest), WithCacheDir(cacheDir))
 		assert.NilError(t, err)
 		assert.Equal(t, StatusUsedCache, r.Status)

pkg/limayaml/limayaml.go

@@ -38,6 +38,24 @@ type LimaYAML struct {
 	Rosetta           Rosetta        `yaml:"rosetta,omitempty" json:"rosetta,omitempty"`
 }
 
+// Returns list of all the files referred by the LimaYAML instance
+func (y *LimaYAML) ReferredFiles() []File {
+	files := []File{}
+	for _, img := range y.Images {
+		files = append(files, img.File)
+		if img.Kernel != nil {
+			files = append(files, img.Kernel.File)
+		}
+		if img.Initrd != nil {
+			files = append(files, *img.Initrd)
+		}
+	}
+	for _, archive := range y.Containerd.Archives {
+		files = append(files, archive)
+	}
+	return files
+}
+
 type Arch = string
 type MountType = string
 type VMType = string

pkg/store/store.go

@@ -102,3 +102,20 @@ func LoadYAMLByFilePath(filePath string) (*limayaml.LimaYAML, error) {
 	}
 	return y, nil
 }
+
+func Downloads() ([]limayaml.File, error) {
+	instanceNames, err := Instances()
+	if err != nil {
+		return nil, err
+	}
+	files := []limayaml.File{}
+
+	for _, instanceName := range instanceNames {
+		instance, err := Inspect(instanceName)
+		if err != nil {
+			return nil, err
+		}
+		files = append(files, instance.Config.ReferredFiles()...)
+	}
+	return files, nil
+}