Snapshotter with caching

Signed-off-by: apostasie <spam_blackhole@farcloser.world>

pkg/cmd/container/inspect.go

@@ -28,6 +28,7 @@ import (
 	"github.com/containerd/nerdctl/v2/pkg/containerinspector"
 	"github.com/containerd/nerdctl/v2/pkg/formatter"
 	"github.com/containerd/nerdctl/v2/pkg/idutil/containerwalker"
+	"github.com/containerd/nerdctl/v2/pkg/imgutil"
 	"github.com/containerd/nerdctl/v2/pkg/inspecttypes/dockercompat"
 )
 
@@ -36,7 +37,7 @@ func Inspect(ctx context.Context, client *containerd.Client, containers []string
 	f := &containerInspector{
 		mode:        options.Mode,
 		size:        options.Size,
-		snapshotter: client.SnapshotService(options.GOptions.Snapshotter),
+		snapshotter: imgutil.SnapshotServiceWithCache(client.SnapshotService(options.GOptions.Snapshotter)),
 	}
 
 	walker := &containerwalker.ContainerWalker{
@@ -61,39 +62,6 @@ type containerInspector struct {
 	entries     []interface{}
 }
 
-// resourceTotal will return:
-// - the Usage value of the resource referenced by ID
-// - the cumulative Usage value of the resource, and all parents, recursively
-// Typically, for a running container, this will equal the size of the read-write layer, plus the sum of the size of all layers in the base image
-func resourceTotal(ctx context.Context, snapshotter snapshots.Snapshotter, resourceID string) (snapshots.Usage, snapshots.Usage, error) {
-	var first snapshots.Usage
-	var total snapshots.Usage
-	var info snapshots.Info
-
-	for next := resourceID; next != ""; next = info.Parent {
-		// Get the resource usage info
-		usage, err := snapshotter.Usage(ctx, next)
-		if err != nil {
-			return first, total, err
-		}
-		// In case that's the first one, store that
-		if next == resourceID {
-			first = usage
-		}
-		// And increment totals
-		total.Size += usage.Size
-		total.Inodes += usage.Inodes
-
-		// Now, get the parent, if any and iterate
-		info, err = snapshotter.Stat(ctx, next)
-		if err != nil {
-			return first, total, err
-		}
-	}
-
-	return first, total, nil
-}
-
 func (x *containerInspector) Handler(ctx context.Context, found containerwalker.Found) error {
 	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
@@ -111,7 +79,7 @@ func (x *containerInspector) Handler(ctx context.Context, found containerwalker.
 			return err
 		}
 		if x.size {
-			resourceUsage, allResourceUsage, err := resourceTotal(ctx, x.snapshotter, d.ID)
+			resourceUsage, allResourceUsage, err := imgutil.ResourceUsage(ctx, x.snapshotter, d.ID)
 			if err == nil {
 				d.SizeRw = &resourceUsage.Size
 				d.SizeRootFs = &allResourceUsage.Size

pkg/cmd/container/list.go

@@ -25,9 +25,9 @@ import (
 	"time"
 
 	"github.com/containerd/containerd"
-	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/pkg/progress"
+	"github.com/containerd/containerd/snapshots"
 	"github.com/containerd/log"
 	"github.com/containerd/nerdctl/v2/pkg/api/types"
 	"github.com/containerd/nerdctl/v2/pkg/formatter"
@@ -107,8 +107,14 @@ func (x *ListItem) Label(s string) string {
 
 func prepareContainers(ctx context.Context, client *containerd.Client, containers []containerd.Container, options types.ContainerListOptions) ([]ListItem, error) {
 	listItems := make([]ListItem, len(containers))
+	snapshottersCache := map[string]snapshots.Snapshotter{}
 	for i, c := range containers {
 		info, err := c.Info(ctx, containerd.WithoutRefreshedMetadata)
+		snapshotter, ok := snapshottersCache[info.Snapshotter]
+		if !ok {
+			snapshottersCache[info.Snapshotter] = imgutil.SnapshotServiceWithCache(client.SnapshotService(info.Snapshotter))
+			snapshotter = snapshottersCache[info.Snapshotter]
+		}
 		if err != nil {
 			if errdefs.IsNotFound(err) {
 				log.G(ctx).Warn(err)
@@ -141,7 +147,7 @@ func prepareContainers(ctx context.Context, client *containerd.Client, container
 			Labels:    info.Labels,
 		}
 		if options.Size {
-			containerSize, err := getContainerSize(ctx, client, c, info)
+			containerSize, err := getContainerSize(ctx, snapshotter, info.SnapshotKey)
 			if err != nil {
 				return nil, err
 			}
@@ -180,30 +186,18 @@ func getContainerNetworks(containerLables map[string]string) []string {
 	return networks
 }
 
-func getContainerSize(ctx context.Context, client *containerd.Client, c containerd.Container, info containers.Container) (string, error) {
+func getContainerSize(ctx context.Context, snapshotter snapshots.Snapshotter, snapshotKey string) (string, error) {
 	// get container snapshot size
-	snapshotKey := info.SnapshotKey
 	var containerSize int64
+	var imageSize int64
 
 	if snapshotKey != "" {
-		usage, err := client.SnapshotService(info.Snapshotter).Usage(ctx, snapshotKey)
+		rw, all, err := imgutil.ResourceUsage(ctx, snapshotter, snapshotKey)
 		if err != nil {
 			return "", err
 		}
-		containerSize = usage.Size
-	}
-
-	// get the image interface
-	image, err := c.Image(ctx)
-	if err != nil {
-		return "", err
-	}
-
-	sn := client.SnapshotService(info.Snapshotter)
-
-	imageSize, err := imgutil.UnpackedImageSize(ctx, sn, image)
-	if err != nil {
-		return "", err
+		containerSize = rw.Size
+		imageSize = all.Size
 	}
 
 	return fmt.Sprintf("%s (virtual %s)", progress.Bytes(containerSize).String(), progress.Bytes(imageSize).String()), nil

pkg/cmd/image/inspect.go

@@ -27,6 +27,7 @@ import (
 	"github.com/containerd/nerdctl/v2/pkg/formatter"
 	"github.com/containerd/nerdctl/v2/pkg/idutil/imagewalker"
 	"github.com/containerd/nerdctl/v2/pkg/imageinspector"
+	"github.com/containerd/nerdctl/v2/pkg/imgutil"
 	"github.com/containerd/nerdctl/v2/pkg/inspecttypes/dockercompat"
 )
 
@@ -35,13 +36,14 @@ func Inspect(ctx context.Context, client *containerd.Client, images []string, op
 	f := &imageInspector{
 		mode: options.Mode,
 	}
+	snapshotter := imgutil.SnapshotServiceWithCache(client.SnapshotService(options.GOptions.Snapshotter))
 	walker := &imagewalker.ImageWalker{
 		Client: client,
 		OnFound: func(ctx context.Context, found imagewalker.Found) error {
 			ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 			defer cancel()
 
-			n, err := imageinspector.Inspect(ctx, client, found.Image, options.GOptions.Snapshotter)
+			n, err := imageinspector.Inspect(ctx, client, found.Image, snapshotter)
 			if err != nil {
 				return err
 			}

pkg/cmd/image/list.go

@@ -169,7 +169,7 @@ func printImages(ctx context.Context, client *containerd.Client, imageList []ima
 		tmpl:         tmpl,
 		client:       client,
 		contentStore: client.ContentStore(),
-		snapshotter:  client.SnapshotService(options.GOptions.Snapshotter),
+		snapshotter:  imgutil.SnapshotServiceWithCache(client.SnapshotService(options.GOptions.Snapshotter)),
 	}
 
 	for _, img := range imageList {

pkg/imageinspector/imageinspector.go

@@ -21,13 +21,14 @@ import (
 
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/snapshots"
 	"github.com/containerd/log"
-	imgutil "github.com/containerd/nerdctl/v2/pkg/imgutil"
+	"github.com/containerd/nerdctl/v2/pkg/imgutil"
 	"github.com/containerd/nerdctl/v2/pkg/inspecttypes/native"
 )
 
 // Inspect inspects the image, for the platform specified in image.platform.
-func Inspect(ctx context.Context, client *containerd.Client, image images.Image, snapshotter string) (*native.Image, error) {
+func Inspect(ctx context.Context, client *containerd.Client, image images.Image, snapshotter snapshots.Snapshotter) (*native.Image, error) {
 
 	n := &native.Image{}
 
@@ -55,8 +56,7 @@ func Inspect(ctx context.Context, client *containerd.Client, image images.Image,
 		n.ImageConfigDesc = imageConfigDesc
 		n.ImageConfig = imageConfig
 	}
-	snapSvc := client.SnapshotService(snapshotter)
-	n.Size, err = imgutil.UnpackedImageSize(ctx, snapSvc, img)
+	n.Size, err = imgutil.UnpackedImageSize(ctx, snapshotter, img)
 	if err != nil {
 		log.G(ctx).WithError(err).WithField("id", image.Name).Warnf("failed to inspect calculate size")
 	}

pkg/imgutil/imgutil.go

@@ -377,27 +377,36 @@ func ParseRepoTag(imgName string) (string, string) {
 	return repository, tag
 }
 
-type snapshotKey string
-
-// recursive function to calculate total usage of key's parent
-func (key snapshotKey) add(ctx context.Context, s snapshots.Snapshotter, usage *snapshots.Usage) error {
-	if key == "" {
-		return nil
-	}
-	u, err := s.Usage(ctx, string(key))
-	if err != nil {
-		return err
-	}
-
-	usage.Add(u)
+// ResourceUsage will return:
+// - the Usage value of the resource referenced by ID
+// - the cumulative Usage value of the resource, and all parents, recursively
+// Typically, for a running container, this will equal the size of the read-write layer, plus the sum of the size of all layers in the base image
+func ResourceUsage(ctx context.Context, snapshotter snapshots.Snapshotter, resourceID string) (snapshots.Usage, snapshots.Usage, error) {
+	var first snapshots.Usage
+	var total snapshots.Usage
+	var info snapshots.Info
+	for next := resourceID; next != ""; next = info.Parent {
+		// Get the resource usage info
+		usage, err := snapshotter.Usage(ctx, next)
+		if err != nil {
+			return first, total, err
+		}
+		// In case that's the first one, store that
+		if next == resourceID {
+			first = usage
+		}
+		// And increment totals
+		total.Size += usage.Size
+		total.Inodes += usage.Inodes
 
-	info, err := s.Stat(ctx, string(key))
-	if err != nil {
-		return err
+		// Now, get the parent, if any and iterate
+		info, err = snapshotter.Stat(ctx, next)
+		if err != nil {
+			return first, total, err
+		}
 	}
 
-	key = snapshotKey(info.Parent)
-	return key.add(ctx, s, usage)
+	return first, total, nil
 }
 
 // UnpackedImageSize is the size of the unpacked snapshots.
@@ -409,23 +418,7 @@ func UnpackedImageSize(ctx context.Context, s snapshots.Snapshotter, img contain
 	}
 
 	chainID := identity.ChainID(diffIDs).String()
-	usage, err := s.Usage(ctx, chainID)
-	if err != nil {
-		if errdefs.IsNotFound(err) {
-			log.G(ctx).WithError(err).Debugf("image %q seems not unpacked", img.Name())
-			return 0, nil
-		}
-		return 0, err
-	}
-
-	info, err := s.Stat(ctx, chainID)
-	if err != nil {
-		return 0, err
-	}
+	_, total, err := ResourceUsage(ctx, s, chainID)
 
-	//add ChainID's parent usage to the total usage
-	if err := snapshotKey(info.Parent).add(ctx, s, &usage); err != nil {
-		return 0, err
-	}
-	return usage.Size, nil
+	return total.Size, err
 }

pkg/imgutil/snapshotter.go

@@ -17,12 +17,15 @@
 package imgutil
 
 import (
+	"context"
 	"strings"
 
 	socisource "github.com/awslabs/soci-snapshotter/fs/source"
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/mount"
 	ctdsnapshotters "github.com/containerd/containerd/pkg/snapshotters"
+	"github.com/containerd/containerd/snapshots"
 	"github.com/containerd/log"
 	"github.com/containerd/nerdctl/v2/pkg/api/types"
 	"github.com/containerd/nerdctl/v2/pkg/imgutil/pull"
@@ -117,3 +120,71 @@ func stargzExtraLabels(f func(images.Handler) images.Handler, rFlags types.Remot
 func sociExtraLabels(f func(images.Handler) images.Handler, rFlags types.RemoteSnapshotterFlags) func(images.Handler) images.Handler {
 	return socisource.AppendDefaultLabelsHandlerWrapper(rFlags.SociIndexDigest, f)
 }
+
+func SnapshotServiceWithCache(nativeSnapshotter snapshots.Snapshotter) snapshots.Snapshotter {
+	return &CachingSnapshotter{
+		containerdSnapshotter: nativeSnapshotter,
+		statCache:             map[string]snapshots.Info{},
+		usageCache:            map[string]snapshots.Usage{},
+	}
+}
+
+type CachingSnapshotter struct {
+	containerdSnapshotter snapshots.Snapshotter
+	statCache             map[string]snapshots.Info
+	usageCache            map[string]snapshots.Usage
+}
+
+func (snap *CachingSnapshotter) Stat(ctx context.Context, key string) (snapshots.Info, error) {
+	if stat, ok := snap.statCache[key]; ok {
+		return stat, nil
+	}
+	stat, err := snap.containerdSnapshotter.Stat(ctx, key)
+	if err == nil {
+		snap.statCache[key] = stat
+	}
+	return stat, err
+}
+
+func (snap *CachingSnapshotter) Update(ctx context.Context, info snapshots.Info, fieldpaths ...string) (snapshots.Info, error) {
+	return snap.containerdSnapshotter.Update(ctx, info, fieldpaths...)
+}
+
+func (snap *CachingSnapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, error) {
+	if usage, ok := snap.usageCache[key]; ok {
+		return usage, nil
+	}
+	usage, err := snap.containerdSnapshotter.Usage(ctx, key)
+	if err == nil {
+		snap.usageCache[key] = usage
+	}
+	return usage, err
+}
+
+func (snap *CachingSnapshotter) Mounts(ctx context.Context, key string) ([]mount.Mount, error) {
+	return snap.containerdSnapshotter.Mounts(ctx, key)
+}
+
+func (snap *CachingSnapshotter) Prepare(ctx context.Context, key, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
+	return snap.containerdSnapshotter.Prepare(ctx, key, parent, opts...)
+}
+
+func (snap *CachingSnapshotter) View(ctx context.Context, key, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
+	return snap.containerdSnapshotter.View(ctx, key, parent, opts...)
+}
+
+func (snap *CachingSnapshotter) Commit(ctx context.Context, name, key string, opts ...snapshots.Opt) error {
+	return snap.containerdSnapshotter.Commit(ctx, name, key, opts...)
+}
+
+func (snap *CachingSnapshotter) Remove(ctx context.Context, key string) error {
+	return snap.containerdSnapshotter.Remove(ctx, key)
+}
+
+func (snap *CachingSnapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, filters ...string) error {
+	return snap.containerdSnapshotter.Walk(ctx, fn, filters...)
+}
+
+func (snap *CachingSnapshotter) Close() error {
+	return snap.containerdSnapshotter.Close()
+}