Skip to content

Commit

Permalink
cri: fix update of pinned label for images
Browse files Browse the repository at this point in the history 
Signed-off-by: Iceber Gu <caiwei95@hotmail.com>
  • Loading branch information
@Iceber
Iceber committed Nov 10, 2023
1 parent bd2db42 commit 2e014fa
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 8 deletions.
106 changes: 100 additions & 6 deletions pkg/cri/store/image/image.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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]),
},
}
}
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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.
Expand All @@ -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()
Expand Down Expand Up @@ -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())
}
73 changes: 71 additions & 2 deletions pkg/cri/store/image/image_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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"
Expand Down

0 comments on commit 2e014fa

Please sign in to comment.