limactl prune: add --keep-referred option to keep objects that are referred by some instances or templates

cmd/limactl/prune.go

@@ -2,8 +2,11 @@ package main
 
 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/lima-vm/lima/pkg/templatestore"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -17,15 +20,98 @@ func newPruneCommand() *cobra.Command {
 		ValidArgsFunction: cobra.NoFileCompletions,
 		GroupID:           advancedCommand,
 	}
+	pruneCommand.Flags().Bool("keep-referred", false, "Keep objects that are referred by some instances or templates")
 	return pruneCommand
 }
 
-func pruneAction(_ *cobra.Command, _ []string) error {
-	ucd, err := os.UserCacheDir()
+func pruneAction(cmd *cobra.Command, _ []string) error {
+	keepReferred, err := cmd.Flags().GetBool("keep-referred")
 	if err != nil {
 		return err
 	}
-	cacheDir := filepath.Join(ucd, "lima")
-	logrus.Infof("Pruning %q", cacheDir)
-	return os.RemoveAll(cacheDir)
+	opt := downloader.WithCache()
+	if !keepReferred {
+		return downloader.RemoveAllCacheDir(opt)
+	}
+
+	// Prune downloads that are not used by any instances or templates
+	cacheEntries, err := downloader.CacheEntries(opt)
+	if err != nil {
+		return err
+	}
+	knownLocations, err := knownLocations()
+	if err != nil {
+		return err
+	}
+	for cacheKey, cachePath := range cacheEntries {
+		if file, exists := knownLocations[cacheKey]; exists {
+			logrus.Debugf("Keep %q caching %q", cacheKey, file.Location)
+		} else {
+			logrus.Debug("Deleting ", cacheKey)
+			if err := os.RemoveAll(cachePath); err != nil {
+				logrus.Warnf("Failed to delete %q: %v", cacheKey, err)
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func knownLocations() (map[string]limayaml.File, error) {
+	locations := make(map[string]limayaml.File)
+
+	// Collect locations from instances
+	instances, err := store.Instances()
+	if err != nil {
+		return nil, err
+	}
+	for _, instanceName := range instances {
+		instance, err := store.Inspect(instanceName)
+		if err != nil {
+			return nil, err
+		}
+		for k, v := range locationsFromLimaYAML(instance.Config) {
+			locations[k] = v
+		}
+	}
+
+	// Collect locations from templates
+	templates, err := templatestore.Templates()
+	if err != nil {
+		return nil, err
+	}
+	for _, t := range templates {
+		b, err := templatestore.Read(t.Name)
+		if err != nil {
+			return nil, err
+		}
+		y, err := limayaml.Load(b, t.Name)
+		if err != nil {
+			return nil, err
+		}
+		for k, v := range locationsFromLimaYAML(y) {
+			locations[k] = v
+		}
+	}
+	return locations, nil
+}
+
+func locationsFromLimaYAML(y *limayaml.LimaYAML) map[string]limayaml.File {
+	locations := make(map[string]limayaml.File)
+	for _, f := range y.Images {
+		locations[downloader.CacheKey(f.Location)] = f.File
+		if f.Kernel != nil {
+			locations[downloader.CacheKey(f.Kernel.Location)] = f.Kernel.File
+		}
+		if f.Initrd != nil {
+			locations[downloader.CacheKey(f.Initrd.Location)] = *f.Initrd
+		}
+	}
+	for _, f := range y.Containerd.Archives {
+		locations[downloader.CacheKey(f.Location)] = f
+	}
+	for _, f := range y.Firmware.Images {
+		locations[downloader.CacheKey(f.Location)] = f.File
+	}
+	return locations
 }

pkg/downloader/downloader.go

@@ -341,7 +341,7 @@ func Cached(remote string, opts ...Opt) (*Result, error) {
 //   - "time" file contains the time (Last-Modified header)
 //   - "type" file contains the type (Content-Type header)
 func cacheDirectoryPath(cacheDir, remote string) string {
-	return filepath.Join(cacheDir, "download", "by-url-sha256", fmt.Sprintf("%x", sha256.Sum256([]byte(remote))))
+	return filepath.Join(cacheDir, "download", "by-url-sha256", CacheKey(remote))
 }
 
 // cacheDigestPath returns the cache digest file path.
@@ -601,3 +601,55 @@ func downloadHTTP(ctx context.Context, localPath, lastModified, contentType, url
 	}
 	return os.Rename(localPathTmp, localPath)
 }
+
+// CacheEntries returns a map of cache entries.
+// The key is the SHA256 of the URL.
+// The value is the path to the cache entry.
+func CacheEntries(opt ...Opt) (map[string]string, error) {
+	entries := make(map[string]string)
+	var o options
+	for _, f := range opt {
+		if err := f(&o); err != nil {
+			return nil, err
+		}
+	}
+	if o.cacheDir == "" {
+		return entries, nil
+	}
+	downloadDir := filepath.Join(o.cacheDir, "download", "by-url-sha256")
+	_, err := os.Stat(downloadDir)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return entries, nil
+		}
+		return nil, err
+	}
+	cacheEntries, err := os.ReadDir(downloadDir)
+	if err != nil {
+		return nil, err
+	}
+	for _, entry := range cacheEntries {
+		entries[entry.Name()] = filepath.Join(downloadDir, entry.Name())
+	}
+	return entries, nil
+}
+
+// CacheKey returns the key for a cache entry of the remote URL.
+func CacheKey(remote string) string {
+	return fmt.Sprintf("%x", sha256.Sum256([]byte(remote)))
+}
+
+// RemoveAllCacheDir removes the cache directory.
+func RemoveAllCacheDir(opt ...Opt) error {
+	var o options
+	for _, f := range opt {
+		if err := f(&o); err != nil {
+			return err
+		}
+	}
+	if o.cacheDir == "" {
+		return nil
+	}
+	logrus.Infof("Pruning %q", o.cacheDir)
+	return os.RemoveAll(o.cacheDir)
+}