diff --git a/image.go b/image.go index 9cfc03a30..f4c66da9f 100644 --- a/image.go +++ b/image.go @@ -19,6 +19,8 @@ package containerd import ( "context" "fmt" + "strings" + "sync/atomic" "github.com/containerd/containerd/content" "github.com/containerd/containerd/diff" @@ -31,6 +33,7 @@ import ( "github.com/opencontainers/image-spec/identity" ocispec "github.com/opencontainers/image-spec/specs-go/v1" "github.com/pkg/errors" + "golang.org/x/sync/semaphore" ) // Image describes an image used by containers @@ -47,6 +50,8 @@ type Image interface { RootFS(ctx context.Context) ([]digest.Digest, error) // Size returns the total size of the image's packed resources. Size(ctx context.Context) (int64, error) + // Usage returns a usage calculation for the image. + Usage(context.Context, ...UsageOpt) (int64, error) // Config descriptor for the image. Config(ctx context.Context) (ocispec.Descriptor, error) // IsUnpacked returns whether or not an image is unpacked. @@ -55,6 +60,49 @@ type Image interface { ContentStore() content.Store } +type usageOptions struct { + manifestLimit *int + manifestOnly bool + snapshots bool +} + +// UsageOpt is used to configure the usage calculation +type UsageOpt func(*usageOptions) error + +// WithUsageManifestLimit sets the limit to the number of manifests which will +// be walked for usage. Setting this value to 0 will require all manifests to +// be walked, returning ErrNotFound if manifests are missing. +// NOTE: By default all manifests which exist will be walked +// and any non-existent manifests and their subobjects will be ignored. +func WithUsageManifestLimit(i int) UsageOpt { + // If 0 then don't filter any manifests + // By default limits to current platform + return func(o *usageOptions) error { + o.manifestLimit = &i + return nil + } +} + +// WithSnapshotUsage will check for referenced snapshots from the image objects +// and include the snapshot size in the total usage. +func WithSnapshotUsage() UsageOpt { + return func(o *usageOptions) error { + o.snapshots = true + return nil + } +} + +// WithManifestUsage is used to get the usage for an image based on what is +// reported by the manifests rather than what exists in the content store. +// NOTE: This function is best used with the manifest limit set to get a +// consistent value, otherwise non-existent manifests will be excluded. +func WithManifestUsage() UsageOpt { + return func(o *usageOptions) error { + o.manifestOnly = true + return nil + } +} + var _ = (Image)(&image{}) // NewImage returns a client image object from the metadata image @@ -100,8 +148,95 @@ func (i *image) RootFS(ctx context.Context) ([]digest.Digest, error) { } func (i *image) Size(ctx context.Context) (int64, error) { - provider := i.client.ContentStore() - return i.i.Size(ctx, provider, i.platform) + return i.Usage(ctx, WithUsageManifestLimit(1), WithManifestUsage()) +} + +func (i *image) Usage(ctx context.Context, opts ...UsageOpt) (int64, error) { + var config usageOptions + for _, opt := range opts { + if err := opt(&config); err != nil { + return 0, err + } + } + + var ( + provider = i.client.ContentStore() + handler = images.ChildrenHandler(provider) + size int64 + mustExist bool + ) + + if config.manifestLimit != nil { + handler = images.LimitManifests(handler, i.platform, *config.manifestLimit) + mustExist = true + } + + var wh images.HandlerFunc = func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) { + var usage int64 + children, err := handler(ctx, desc) + if err != nil { + if !errdefs.IsNotFound(err) || mustExist { + return nil, err + } + if !config.manifestOnly { + // Do not count size of non-existent objects + desc.Size = 0 + } + } else if config.snapshots || !config.manifestOnly { + info, err := provider.Info(ctx, desc.Digest) + if err != nil { + if !errdefs.IsNotFound(err) { + return nil, err + } + if !config.manifestOnly { + // Do not count size of non-existent objects + desc.Size = 0 + } + } else if info.Size > desc.Size { + // Count actual usage, Size may be unset or -1 + desc.Size = info.Size + } + + for k, v := range info.Labels { + const prefix = "containerd.io/gc.ref.snapshot." + if !strings.HasPrefix(k, prefix) { + continue + } + + sn := i.client.SnapshotService(k[len(prefix):]) + if sn == nil { + continue + } + + u, err := sn.Usage(ctx, v) + if err != nil { + if !errdefs.IsNotFound(err) && !errdefs.IsInvalidArgument(err) { + return nil, err + } + } else { + usage += u.Size + } + } + } + + // Ignore unknown sizes. Generally unknown sizes should + // never be set in manifests, however, the usage + // calculation does not need to enforce this. + if desc.Size >= 0 { + usage += desc.Size + } + + atomic.AddInt64(&size, usage) + + return children, nil + } + + l := semaphore.NewWeighted(3) + if err := images.Dispatch(ctx, wh, l, i.i.Target); err != nil { + return 0, err + } + + return size, nil } func (i *image) Config(ctx context.Context) (ocispec.Descriptor, error) { diff --git a/image_test.go b/image_test.go index cae4694eb..33203764f 100644 --- a/image_test.go +++ b/image_test.go @@ -135,3 +135,101 @@ func TestImagePullWithDistSourceLabel(t *testing.T) { t.Fatal(err) } } + +func TestImageUsage(t *testing.T) { + if testing.Short() || runtime.GOOS == "windows" { + t.Skip() + } + + imageName := "docker.io/library/busybox:latest" + ctx, cancel := testContext(t) + defer cancel() + + client, err := newClient(t, address) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + // Cleanup + err = client.ImageService().Delete(ctx, imageName) + if err != nil && !errdefs.IsNotFound(err) { + t.Fatal(err) + } + + testPlatform := platforms.Only(ocispec.Platform{ + OS: "linux", + Architecture: "amd64", + }) + + // Pull single platform, do not unpack + image, err := client.Pull(ctx, imageName, WithPlatformMatcher(testPlatform)) + if err != nil { + t.Fatal(err) + } + + s1, err := image.Usage(ctx, WithUsageManifestLimit(1)) + if err != nil { + t.Fatal(err) + } + + if _, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage()); err == nil { + t.Fatal("expected NotFound with missing manifests") + } else if !errdefs.IsNotFound(err) { + t.Fatalf("unexpected error: %+v", err) + } + + // Pin image name to specific version for future fetches + imageName = imageName + "@" + image.Target().Digest.String() + + // Fetch single platforms, but all manifests pulled + if _, err := client.Fetch(ctx, imageName, WithPlatformMatcher(testPlatform)); err != nil { + t.Fatal(err) + } + + if s, err := image.Usage(ctx, WithUsageManifestLimit(1)); err != nil { + t.Fatal(err) + } else if s != s1 { + t.Fatalf("unexpected usage %d, expected %d", s, s1) + } + + s2, err := image.Usage(ctx, WithUsageManifestLimit(0)) + if err != nil { + t.Fatal(err) + } + + if s2 <= s1 { + t.Fatalf("Expected larger usage counting all manifests: %d <= %d", s2, s1) + } + + s3, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage()) + if err != nil { + t.Fatal(err) + } + + if s3 <= s2 { + t.Fatalf("Expected larger usage counting all manifest reported sizes: %d <= %d", s3, s2) + } + + // Fetch everything + if _, err = client.Fetch(ctx, imageName); err != nil { + t.Fatal(err) + } + + if s, err := image.Usage(ctx); err != nil { + t.Fatal(err) + } else if s != s3 { + t.Fatalf("Expected actual usage to equal manifest reported usage of %d: got %d", s3, s) + } + + err = image.Unpack(ctx, DefaultSnapshotter) + if err != nil { + t.Fatal(err) + } + + if s, err := image.Usage(ctx, WithSnapshotUsage()); err != nil { + t.Fatal(err) + } else if s <= s3 { + t.Fatalf("Expected actual usage with snapshots to be greater: %d <= %d", s, s3) + } +} diff --git a/images/handlers.go b/images/handlers.go index dac701bb8..04c2d5a60 100644 --- a/images/handlers.go +++ b/images/handlers.go @@ -117,7 +117,7 @@ func Walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) err // // If any handler returns an error, the dispatch session will be canceled. func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, descs ...ocispec.Descriptor) error { - eg, ctx := errgroup.WithContext(ctx) + eg, ctx2 := errgroup.WithContext(ctx) for _, desc := range descs { desc := desc @@ -126,10 +126,11 @@ func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, return err } } + eg.Go(func() error { desc := desc - children, err := handler.Handle(ctx, desc) + children, err := handler.Handle(ctx2, desc) if limiter != nil { limiter.Release(1) } @@ -141,7 +142,7 @@ func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, } if len(children) > 0 { - return Dispatch(ctx, handler, limiter, children...) + return Dispatch(ctx2, handler, limiter, children...) } return nil