From 2e014fa2ac27d7c7d89d639d1cecb47ad70b98e8 Mon Sep 17 00:00:00 2001 From: Iceber Gu Date: Sun, 5 Nov 2023 15:02:12 +0800 Subject: [PATCH] cri: fix update of pinned label for images Signed-off-by: Iceber Gu --- pkg/cri/store/image/image.go | 106 ++++++++++++++++++++++++++++-- pkg/cri/store/image/image_test.go | 73 +++++++++++++++++++- 2 files changed, 171 insertions(+), 8 deletions(-) diff --git a/pkg/cri/store/image/image.go b/pkg/cri/store/image/image.go index df1679439..364562431 100644 --- a/pkg/cri/store/image/image.go +++ b/pkg/cri/store/image/image.go @@ -30,6 +30,7 @@ import ( "github.com/containerd/containerd/v2/pkg/cri/util" "github.com/containerd/containerd/v2/platforms" docker "github.com/distribution/reference" + "k8s.io/apimachinery/pkg/util/sets" "github.com/opencontainers/go-digest" imagedigest "github.com/opencontainers/go-digest" @@ -89,8 +90,9 @@ func NewStore(img images.Store, provider InfoProvider, platform platforms.MatchC provider: provider, platform: platform, store: &store{ - images: make(map[string]Image), - digestSet: digestset.NewSet(), + images: make(map[string]Image), + digestSet: digestset.NewSet(), + pinnedRefs: make(map[string]sets.Set[string]), }, } } @@ -130,7 +132,13 @@ func (s *Store) update(ref string, img *Image) error { } if oldExist { if oldID == img.ID { - return nil + if s.store.isPinned(img.ID, ref) == img.Pinned { + return nil + } + if img.Pinned { + return s.store.pin(img.ID, ref) + } + return s.store.unpin(img.ID, ref) } // Updated. Remove tag from old image. s.store.delete(oldID, ref) @@ -206,9 +214,10 @@ func (s *Store) List() []Image { } type store struct { - lock sync.RWMutex - images map[string]Image - digestSet *digestset.Set + lock sync.RWMutex + images map[string]Image + digestSet *digestset.Set + pinnedRefs map[string]sets.Set[string] } func (s *store) list() []Image { @@ -233,6 +242,14 @@ func (s *store) add(img Image) error { } } + if img.Pinned { + if refs := s.pinnedRefs[img.ID]; refs == nil { + s.pinnedRefs[img.ID] = sets.New(img.References...) + } else { + refs.Insert(img.References...) + } + } + i, ok := s.images[img.ID] if !ok { // If the image doesn't exist, add it. @@ -246,6 +263,73 @@ func (s *store) add(img Image) error { return nil } +func (s *store) isPinned(id, ref string) bool { + s.lock.RLock() + defer s.lock.RUnlock() + digest, err := s.digestSet.Lookup(id) + if err != nil { + return false + } + refs := s.pinnedRefs[digest.String()] + return refs != nil && refs.Has(ref) +} + +func (s *store) pin(id, ref string) error { + s.lock.Lock() + defer s.lock.Unlock() + digest, err := s.digestSet.Lookup(id) + if err != nil { + if err == digestset.ErrDigestNotFound { + err = errdefs.ErrNotFound + } + return err + } + i, ok := s.images[digest.String()] + if !ok { + return errdefs.ErrNotFound + } + + if refs := s.pinnedRefs[digest.String()]; refs == nil { + s.pinnedRefs[digest.String()] = sets.New(ref) + } else { + refs.Insert(ref) + } + i.Pinned = true + s.images[digest.String()] = i + return nil +} + +func (s *store) unpin(id, ref string) error { + s.lock.Lock() + defer s.lock.Unlock() + digest, err := s.digestSet.Lookup(id) + if err != nil { + if err == digestset.ErrDigestNotFound { + err = errdefs.ErrNotFound + } + return err + } + i, ok := s.images[digest.String()] + if !ok { + return errdefs.ErrNotFound + } + + refs := s.pinnedRefs[digest.String()] + if refs == nil { + return nil + } + if refs.Delete(ref); len(refs) > 0 { + return nil + } + + // delete unpinned image, we only need to keep the pinned + // entries in the map + delete(s.pinnedRefs, digest.String()) + i.Pinned = false + s.images[digest.String()] = i + return nil +} + func (s *store) get(id string) (Image, error) { s.lock.RLock() defer s.lock.RUnlock() @@ -277,10 +361,20 @@ func (s *store) delete(id, ref string) { } i.References = util.SubtractStringSlice(i.References, ref) if len(i.References) != 0 { + if refs := s.pinnedRefs[digest.String()]; refs != nil { + if refs.Delete(ref); len(refs) == 0 { + i.Pinned = false + // delete unpinned image, we only need to keep the pinned + // entries in the map + delete(s.pinnedRefs, digest.String()) + } + } + s.images[digest.String()] = i return } // Remove the image if it is not referenced any more. s.digestSet.Remove(digest) delete(s.images, digest.String()) + delete(s.pinnedRefs, digest.String()) } diff --git a/pkg/cri/store/image/image_test.go b/pkg/cri/store/image/image_test.go index 7ee72cb5e..f9e23d3ac 100644 --- a/pkg/cri/store/image/image_test.go +++ b/pkg/cri/store/image/image_test.go @@ -22,6 +22,7 @@ import ( "testing" "github.com/containerd/containerd/v2/errdefs" + "k8s.io/apimachinery/pkg/util/sets" "github.com/opencontainers/go-digest/digestset" assertlib "github.com/stretchr/testify/assert" @@ -58,8 +59,9 @@ func TestInternalStore(t *testing.T) { genTruncIndex := func(normalName string) string { return normalName[:(len(normalName)+1)/2] } s := &store{ - images: make(map[string]Image), - digestSet: digestset.NewSet(), + images: make(map[string]Image), + digestSet: digestset.NewSet(), + pinnedRefs: make(map[string]sets.Set[string]), } t.Logf("should be able to add image") @@ -137,6 +139,73 @@ func TestInternalStore(t *testing.T) { } } +func TestInternalStorePinnedImage(t *testing.T) { + assert := assertlib.New(t) + s := &store{ + images: make(map[string]Image), + digestSet: digestset.NewSet(), + pinnedRefs: make(map[string]sets.Set[string]), + } + + ref1 := "containerd.io/ref-1" + image := Image{ + ID: "sha256:1123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + ChainID: "test-chain-id-1", + References: []string{ref1}, + Size: 10, + } + + t.Logf("add unpinned image ref, image should be unpinned") + assert.NoError(s.add(image)) + i, err := s.get(image.ID) + assert.NoError(err) + assert.False(i.Pinned) + assert.False(s.isPinned(image.ID, ref1)) + + t.Logf("add pinned image ref, image should be pinned") + ref2 := "containerd.io/ref-2" + image.References = []string{ref2} + image.Pinned = true + assert.NoError(s.add(image)) + i, err = s.get(image.ID) + assert.NoError(err) + assert.True(i.Pinned) + assert.False(s.isPinned(image.ID, ref1)) + assert.True(s.isPinned(image.ID, ref2)) + + t.Logf("pin unpinned image ref, image should be pinned, all refs should be pinned") + assert.NoError(s.pin(image.ID, ref1)) + i, err = s.get(image.ID) + assert.NoError(err) + assert.True(i.Pinned) + assert.True(s.isPinned(image.ID, ref1)) + assert.True(s.isPinned(image.ID, ref2)) + + t.Logf("unpin one of image refs, image should be pinned") + assert.NoError(s.unpin(image.ID, ref2)) + i, err = s.get(image.ID) + assert.NoError(err) + assert.True(i.Pinned) + assert.True(s.isPinned(image.ID, ref1)) + assert.False(s.isPinned(image.ID, ref2)) + + t.Logf("unpin the remaining one image ref, image should be unpinned") + assert.NoError(s.unpin(image.ID, ref1)) + i, err = s.get(image.ID) + assert.NoError(err) + assert.False(i.Pinned) + assert.False(s.isPinned(image.ID, ref1)) + assert.False(s.isPinned(image.ID, ref2)) + + t.Logf("pin one of image refs, then delete this, image should be unpinned") + assert.NoError(s.pin(image.ID, ref1)) + s.delete(image.ID, ref1) + i, err = s.get(image.ID) + assert.NoError(err) + assert.False(i.Pinned) + assert.False(s.isPinned(image.ID, ref2)) +} + func TestImageStore(t *testing.T) { id := "sha256:1123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" newID := "sha256:9923456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"