From 75771c4634892c1dc79b9d2856e11c51a017dd47 Mon Sep 17 00:00:00 2001 From: Derek McGowan Date: Thu, 8 Aug 2019 10:06:18 -0700 Subject: [PATCH] Add usage function to client The usage function allows more configurable and accurate calculations of the usage for an image. Signed-off-by: Derek McGowan --- image.go | 139 +++++++++++++++++++++++++++++++++++++++++++++++++- image_test.go | 98 +++++++++++++++++++++++++++++++++++ 2 files changed, 235 insertions(+), 2 deletions(-) diff --git a/image.go b/image.go index 77c95eaa4..df0e3196f 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/errdefs" @@ -29,6 +31,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 @@ -45,6 +48,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. @@ -53,6 +58,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 @@ -98,8 +146,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) + } +}