[release/1.7] cri: fix using the pinned label to pin image

integration/containerd_image_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/integration/images"
 	"github.com/containerd/containerd/namespaces"
+	"github.com/containerd/containerd/pkg/cri/labels"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
@@ -47,7 +48,8 @@ func TestContainerdImage(t *testing.T) {
 	}
 
 	t.Logf("pull the image into containerd")
-	_, err = containerdClient.Pull(ctx, testImage, containerd.WithPullUnpack, containerd.WithPullLabel("foo", "bar"))
+	lbs := map[string]string{"foo": "bar", labels.PinnedImageLabelKey: labels.PinnedImageLabelValue}
+	_, err = containerdClient.Pull(ctx, testImage, containerd.WithPullUnpack, containerd.WithPullLabels(lbs))
 	assert.NoError(t, err)
 	defer func() {
 		// Make sure the image is cleaned up in any case.
@@ -128,6 +130,13 @@ func TestContainerdImage(t *testing.T) {
 	img, err := containerdClient.GetImage(ctx, testImage)
 	assert.NoError(t, err)
 	assert.Equal(t, img.Labels()["foo"], "bar")
+	assert.Equal(t, img.Labels()[labels.ImageLabelKey], labels.ImageLabelValue)
+
+	t.Logf("the image should be pinned")
+	i, err = imageService.ImageStatus(&runtime.ImageSpec{Image: testImage})
+	require.NoError(t, err)
+	require.NotNil(t, i)
+	assert.True(t, i.Pinned)
 
 	t.Logf("should be able to start container with the image")
 	sb, sbConfig := PodSandboxConfigWithCleanup(t, "sandbox", "containerd-image")

pkg/cri/store/image/image.go

@@ -26,6 +26,7 @@ import (
 	"github.com/containerd/containerd/pkg/cri/labels"
 	"github.com/containerd/containerd/pkg/cri/util"
 	"github.com/containerd/containerd/reference/docker"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	imagedigest "github.com/opencontainers/go-digest"
 	"github.com/opencontainers/go-digest/digestset"
@@ -67,8 +68,9 @@ func NewStore(client *containerd.Client) *Store {
 		refCache: make(map[string]string),
 		client:   client,
 		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]),
 		},
 	}
 }
@@ -106,7 +108,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)
@@ -178,9 +186,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 {
@@ -205,6 +214,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.
@@ -213,10 +230,78 @@ func (s *store) add(img Image) error {
 	}
 	// Or else, merge and sort the references.
 	i.References = docker.Sort(util.MergeStringSlices(i.References, img.References))
+	i.Pinned = i.Pinned || img.Pinned
 	s.images[img.ID] = i
 	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()
@@ -248,10 +333,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/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"