cri: fix update of pinned label for images

Signed-off-by: Iceber Gu <caiwei95@hotmail.com>

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"
@@ -91,6 +92,7 @@ func NewStore(img images.Store, provider InfoProvider, platform platforms.MatchC
 		store: &store{
 			images:     make(map[string]Image),
 			digestSet:  digestset.NewSet(),
+			pinnedRefs: make(map[string]sets.Set[string]),
 		},
 	}
 }
@@ -130,8 +132,14 @@ func (s *Store) update(ref string, img *Image) error {
 	}
 	if oldExist {
 		if oldID == img.ID {
+			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)
 	}
@@ -209,6 +217,7 @@ type store struct {
 	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())
 }

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"
@@ -60,6 +61,7 @@ func TestInternalStore(t *testing.T) {
 	s := &store{
 		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"