Skip to content

Commit ec0c968

Browse files
estespestesp
authored
Merge pull request #3520 from dmcgowan/image-usage
Add image usage function to client
2 parents f650414 + c017e0e commit ec0c968

3 files changed

Lines changed: 239 additions & 5 deletions

File tree

‎image.go‎

Lines changed: 137 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ package containerd
1919
import (
2020
"context"
2121
"fmt"
22+
"strings"
23+
"sync/atomic"
2224

2325
"github.com/containerd/containerd/content"
2426
"github.com/containerd/containerd/diff"
@@ -31,6 +33,7 @@ import (
3133
"github.com/opencontainers/image-spec/identity"
3234
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
3335
"github.com/pkg/errors"
36+
"golang.org/x/sync/semaphore"
3437
)
3538

3639
// Image describes an image used by containers
@@ -47,6 +50,8 @@ type Image interface {
4750
RootFS(ctx context.Context) ([]digest.Digest, error)
4851
// Size returns the total size of the image's packed resources.
4952
Size(ctx context.Context) (int64, error)
53+
// Usage returns a usage calculation for the image.
54+
Usage(context.Context, ...UsageOpt) (int64, error)
5055
// Config descriptor for the image.
5156
Config(ctx context.Context) (ocispec.Descriptor, error)
5257
// IsUnpacked returns whether or not an image is unpacked.
@@ -55,6 +60,49 @@ type Image interface {
5560
ContentStore() content.Store
5661
}
5762

63+
type usageOptions struct {
64+
manifestLimit *int
65+
manifestOnly bool
66+
snapshots bool
67+
}
68+
69+
// UsageOpt is used to configure the usage calculation
70+
type UsageOpt func(*usageOptions) error
71+
72+
// WithUsageManifestLimit sets the limit to the number of manifests which will
73+
// be walked for usage. Setting this value to 0 will require all manifests to
74+
// be walked, returning ErrNotFound if manifests are missing.
75+
// NOTE: By default all manifests which exist will be walked
76+
// and any non-existent manifests and their subobjects will be ignored.
77+
func WithUsageManifestLimit(i int) UsageOpt {
78+
// If 0 then don't filter any manifests
79+
// By default limits to current platform
80+
return func(o *usageOptions) error {
81+
o.manifestLimit = &i
82+
return nil
83+
}
84+
}
85+
86+
// WithSnapshotUsage will check for referenced snapshots from the image objects
87+
// and include the snapshot size in the total usage.
88+
func WithSnapshotUsage() UsageOpt {
89+
return func(o *usageOptions) error {
90+
o.snapshots = true
91+
return nil
92+
}
93+
}
94+
95+
// WithManifestUsage is used to get the usage for an image based on what is
96+
// reported by the manifests rather than what exists in the content store.
97+
// NOTE: This function is best used with the manifest limit set to get a
98+
// consistent value, otherwise non-existent manifests will be excluded.
99+
func WithManifestUsage() UsageOpt {
100+
return func(o *usageOptions) error {
101+
o.manifestOnly = true
102+
return nil
103+
}
104+
}
105+
58106
var _ = (Image)(&image{})
59107

60108
// NewImage returns a client image object from the metadata image
@@ -100,8 +148,95 @@ func (i *image) RootFS(ctx context.Context) ([]digest.Digest, error) {
100148
}
101149

102150
func (i *image) Size(ctx context.Context) (int64, error) {
103-
provider := i.client.ContentStore()
104-
return i.i.Size(ctx, provider, i.platform)
151+
return i.Usage(ctx, WithUsageManifestLimit(1), WithManifestUsage())
152+
}
153+
154+
func (i *image) Usage(ctx context.Context, opts ...UsageOpt) (int64, error) {
155+
var config usageOptions
156+
for _, opt := range opts {
157+
if err := opt(&config); err != nil {
158+
return 0, err
159+
}
160+
}
161+
162+
var (
163+
provider = i.client.ContentStore()
164+
handler = images.ChildrenHandler(provider)
165+
size int64
166+
mustExist bool
167+
)
168+
169+
if config.manifestLimit != nil {
170+
handler = images.LimitManifests(handler, i.platform, *config.manifestLimit)
171+
mustExist = true
172+
}
173+
174+
var wh images.HandlerFunc = func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
175+
var usage int64
176+
children, err := handler(ctx, desc)
177+
if err != nil {
178+
if !errdefs.IsNotFound(err) || mustExist {
179+
return nil, err
180+
}
181+
if !config.manifestOnly {
182+
// Do not count size of non-existent objects
183+
desc.Size = 0
184+
}
185+
} else if config.snapshots || !config.manifestOnly {
186+
info, err := provider.Info(ctx, desc.Digest)
187+
if err != nil {
188+
if !errdefs.IsNotFound(err) {
189+
return nil, err
190+
}
191+
if !config.manifestOnly {
192+
// Do not count size of non-existent objects
193+
desc.Size = 0
194+
}
195+
} else if info.Size > desc.Size {
196+
// Count actual usage, Size may be unset or -1
197+
desc.Size = info.Size
198+
}
199+
200+
for k, v := range info.Labels {
201+
const prefix = "containerd.io/gc.ref.snapshot."
202+
if !strings.HasPrefix(k, prefix) {
203+
continue
204+
}
205+
206+
sn := i.client.SnapshotService(k[len(prefix):])
207+
if sn == nil {
208+
continue
209+
}
210+
211+
u, err := sn.Usage(ctx, v)
212+
if err != nil {
213+
if !errdefs.IsNotFound(err) && !errdefs.IsInvalidArgument(err) {
214+
return nil, err
215+
}
216+
} else {
217+
usage += u.Size
218+
}
219+
}
220+
}
221+
222+
// Ignore unknown sizes. Generally unknown sizes should
223+
// never be set in manifests, however, the usage
224+
// calculation does not need to enforce this.
225+
if desc.Size >= 0 {
226+
usage += desc.Size
227+
}
228+
229+
atomic.AddInt64(&size, usage)
230+
231+
return children, nil
232+
}
233+
234+
l := semaphore.NewWeighted(3)
235+
if err := images.Dispatch(ctx, wh, l, i.i.Target); err != nil {
236+
return 0, err
237+
}
238+
239+
return size, nil
105240
}
106241

107242
func (i *image) Config(ctx context.Context) (ocispec.Descriptor, error) {

‎image_test.go‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,101 @@ func TestImagePullWithDistSourceLabel(t *testing.T) {
135135
t.Fatal(err)
136136
}
137137
}
138+
139+
func TestImageUsage(t *testing.T) {
140+
if testing.Short() || runtime.GOOS == "windows" {
141+
t.Skip()
142+
}
143+
144+
imageName := "docker.io/library/busybox:latest"
145+
ctx, cancel := testContext(t)
146+
defer cancel()
147+
148+
client, err := newClient(t, address)
149+
if err != nil {
150+
t.Fatal(err)
151+
}
152+
defer client.Close()
153+
154+
// Cleanup
155+
err = client.ImageService().Delete(ctx, imageName)
156+
if err != nil && !errdefs.IsNotFound(err) {
157+
t.Fatal(err)
158+
}
159+
160+
testPlatform := platforms.Only(ocispec.Platform{
161+
OS: "linux",
162+
Architecture: "amd64",
163+
})
164+
165+
// Pull single platform, do not unpack
166+
image, err := client.Pull(ctx, imageName, WithPlatformMatcher(testPlatform))
167+
if err != nil {
168+
t.Fatal(err)
169+
}
170+
171+
s1, err := image.Usage(ctx, WithUsageManifestLimit(1))
172+
if err != nil {
173+
t.Fatal(err)
174+
}
175+
176+
if _, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage()); err == nil {
177+
t.Fatal("expected NotFound with missing manifests")
178+
} else if !errdefs.IsNotFound(err) {
179+
t.Fatalf("unexpected error: %+v", err)
180+
}
181+
182+
// Pin image name to specific version for future fetches
183+
imageName = imageName + "@" + image.Target().Digest.String()
184+
185+
// Fetch single platforms, but all manifests pulled
186+
if _, err := client.Fetch(ctx, imageName, WithPlatformMatcher(testPlatform)); err != nil {
187+
t.Fatal(err)
188+
}
189+
190+
if s, err := image.Usage(ctx, WithUsageManifestLimit(1)); err != nil {
191+
t.Fatal(err)
192+
} else if s != s1 {
193+
t.Fatalf("unexpected usage %d, expected %d", s, s1)
194+
}
195+
196+
s2, err := image.Usage(ctx, WithUsageManifestLimit(0))
197+
if err != nil {
198+
t.Fatal(err)
199+
}
200+
201+
if s2 <= s1 {
202+
t.Fatalf("Expected larger usage counting all manifests: %d <= %d", s2, s1)
203+
}
204+
205+
s3, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage())
206+
if err != nil {
207+
t.Fatal(err)
208+
}
209+
210+
if s3 <= s2 {
211+
t.Fatalf("Expected larger usage counting all manifest reported sizes: %d <= %d", s3, s2)
212+
}
213+
214+
// Fetch everything
215+
if _, err = client.Fetch(ctx, imageName); err != nil {
216+
t.Fatal(err)
217+
}
218+
219+
if s, err := image.Usage(ctx); err != nil {
220+
t.Fatal(err)
221+
} else if s != s3 {
222+
t.Fatalf("Expected actual usage to equal manifest reported usage of %d: got %d", s3, s)
223+
}
224+
225+
err = image.Unpack(ctx, DefaultSnapshotter)
226+
if err != nil {
227+
t.Fatal(err)
228+
}
229+
230+
if s, err := image.Usage(ctx, WithSnapshotUsage()); err != nil {
231+
t.Fatal(err)
232+
} else if s <= s3 {
233+
t.Fatalf("Expected actual usage with snapshots to be greater: %d <= %d", s, s3)
234+
}
235+
}

‎images/handlers.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func Walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) err
117117
//
118118
// If any handler returns an error, the dispatch session will be canceled.
119119
func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, descs ...ocispec.Descriptor) error {
120-
eg, ctx := errgroup.WithContext(ctx)
120+
eg, ctx2 := errgroup.WithContext(ctx)
121121
for _, desc := range descs {
122122
desc := desc
123123

@@ -126,10 +126,11 @@ func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted,
126126
return err
127127
}
128128
}
129+
129130
eg.Go(func() error {
130131
desc := desc
131132

132-
children, err := handler.Handle(ctx, desc)
133+
children, err := handler.Handle(ctx2, desc)
133134
if limiter != nil {
134135
limiter.Release(1)
135136
}
@@ -141,7 +142,7 @@ func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted,
141142
}
142143

143144
if len(children) > 0 {
144-
return Dispatch(ctx, handler, limiter, children...)
145+
return Dispatch(ctx2, handler, limiter, children...)
145146
}
146147

147148
return nil

0 commit comments

Comments
 (0)