fix(server): prune files

server/images.go

@@ -602,132 +602,6 @@ func CopyModel(src, dst model.Name) error {
 	return err
 }
 
-func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
-	fp, err := GetManifestPath()
-	if err != nil {
-		return err
-	}
-
-	walkFunc := func(path string, info os.FileInfo, _ error) error {
-		if info.IsDir() {
-			return nil
-		}
-
-		dir, file := filepath.Split(path)
-		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
-		tag := strings.Join([]string{dir, file}, ":")
-		fmp := ParseModelPath(tag)
-
-		// skip the manifest we're trying to delete
-		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
-			return nil
-		}
-
-		// save (i.e. delete from the deleteMap) any files used in other manifests
-		manifest, _, err := GetManifest(fmp)
-		if err != nil {
-			// nolint: nilerr
-			return nil
-		}
-
-		for _, layer := range manifest.Layers {
-			delete(deleteMap, layer.Digest)
-		}
-
-		delete(deleteMap, manifest.Config.Digest)
-		return nil
-	}
-
-	if err := filepath.Walk(fp, walkFunc); err != nil {
-		return err
-	}
-
-	// only delete the files which are still in the deleteMap
-	for k := range deleteMap {
-		fp, err := GetBlobsPath(k)
-		if err != nil {
-			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
-			continue
-		}
-		if !dryRun {
-			if err := os.Remove(fp); err != nil {
-				slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
-				continue
-			}
-		} else {
-			slog.Info(fmt.Sprintf("wanted to remove: %s", fp))
-		}
-	}
-
-	return nil
-}
-
-func PruneLayers() error {
-	deleteMap := make(map[string]struct{})
-	p, err := GetBlobsPath("")
-	if err != nil {
-		return err
-	}
-
-	blobs, err := os.ReadDir(p)
-	if err != nil {
-		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
-		return err
-	}
-
-	for _, blob := range blobs {
-		name := blob.Name()
-		name = strings.ReplaceAll(name, "-", ":")
-		if strings.HasPrefix(name, "sha256:") {
-			deleteMap[name] = struct{}{}
-		}
-	}
-
-	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
-
-	err = deleteUnusedLayers(nil, deleteMap, false)
-	if err != nil {
-		return err
-	}
-
-	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
-
-	return nil
-}
-
-func PruneDirectory(path string) error {
-	info, err := os.Lstat(path)
-	if err != nil {
-		return err
-	}
-
-	if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
-		entries, err := os.ReadDir(path)
-		if err != nil {
-			return err
-		}
-
-		for _, entry := range entries {
-			if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
-				return err
-			}
-		}
-
-		entries, err = os.ReadDir(path)
-		if err != nil {
-			return err
-		}
-
-		if len(entries) > 0 {
-			return nil
-		}
-
-		return os.Remove(path)
-	}
-
-	return nil
-}
-
 func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
 	mp := ParseModelPath(name)
 	fn(api.ProgressResponse{Status: "retrieving manifest"})

server/layer.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
+	"strings"
 )
 
 type Layer struct {
@@ -111,3 +113,36 @@ func (l *Layer) Remove() error {
 
 	return os.Remove(blob)
 }
+
+func Layers() (map[string]*Layer, error) {
+	blobs, err := GetBlobsPath("")
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO(mxyng): use something less brittle
+	matches, err := filepath.Glob(filepath.Join(blobs, "*"))
+	if err != nil {
+		return nil, err
+	}
+
+	ds := make(map[string]*Layer)
+	for _, match := range matches {
+		rel, err := filepath.Rel(blobs, match)
+		if err != nil {
+			return nil, err
+		}
+
+		// TODO(mxyng): this should ideally use model.Digest but
+		// that's currently incompatible with the manifest digest
+		d := strings.Replace(rel, "sha256-", "sha256:", 1)
+		layer, err := NewLayerFromLayer(d, "", "")
+		if err != nil {
+			return nil, err
+		}
+
+		ds[d] = layer
+	}
+
+	return ds, nil
+}

server/manifest.go

@@ -37,7 +37,42 @@ func (m *Manifest) Remove() error {
 		return err
 	}
 
-	return PruneDirectory(manifests)
+	return pruneEmptyDirectory(manifests)
+}
+
+func pruneEmptyDirectory(p string) error {
+	fi, err := os.Lstat(p)
+	if err != nil {
+		return err
+	}
+
+	if fi.Mode()&os.ModeSymlink == 0 {
+		entries, err := os.ReadDir(p)
+		if err != nil {
+			return err
+		}
+
+		for _, entry := range entries {
+			if entry.IsDir() {
+				if err := pruneEmptyDirectory(filepath.Join(p, entry.Name())); err != nil {
+					return err
+				}
+			}
+		}
+
+		entries, err = os.ReadDir(p)
+		if err != nil {
+			return err
+		}
+
+		if len(entries) == 0 {
+			if err := os.Remove(p); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
 }
 
 func (m *Manifest) RemoveLayers() error {

server/routes.go

@@ -970,17 +970,15 @@ func Serve(ln net.Listener) error {
 
 	if !envconfig.NoPrune {
 		// clean up unused layers and manifests
-		if err := PruneLayers(); err != nil {
-			return err
-		}
-
-		manifestsPath, err := GetManifestPath()
+		layers, err := Layers()
 		if err != nil {
 			return err
 		}
 
-		if err := PruneDirectory(manifestsPath); err != nil {
-			return err
+		for _, layer := range layers {
+			if err := layer.Remove(); err != nil {
+				return err
+			}
 		}
 	}