Move CRI from pkg/ to internal/

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>

internal/cri/store/image/fake_image.go

@@ -0,0 +1,38 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package image
+
+import (
+	"fmt"
+
+	"github.com/containerd/platforms"
+)
+
+// NewFakeStore returns an image store with predefined images.
+// Update is not allowed for this fake store.
+func NewFakeStore(images []Image) (*Store, error) {
+	s := NewStore(nil, nil, platforms.Default())
+	for _, i := range images {
+		for _, ref := range i.References {
+			s.refCache[ref] = i.ID
+		}
+		if err := s.store.add(i); err != nil {
+			return nil, fmt.Errorf("add image %+v: %w", i, err)
+		}
+	}
+	return s, nil
+}

internal/cri/store/image/image.go

@@ -0,0 +1,378 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package image
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+
+	"github.com/containerd/containerd/v2/core/content"
+	"github.com/containerd/containerd/v2/core/images"
+	"github.com/containerd/containerd/v2/core/images/usage"
+	"github.com/containerd/containerd/v2/internal/cri/labels"
+	"github.com/containerd/containerd/v2/internal/cri/util"
+	"github.com/containerd/errdefs"
+	"github.com/containerd/platforms"
+	docker "github.com/distribution/reference"
+	"k8s.io/apimachinery/pkg/util/sets"
+
+	imagedigest "github.com/opencontainers/go-digest"
+	"github.com/opencontainers/go-digest/digestset"
+	imageidentity "github.com/opencontainers/image-spec/identity"
+	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// Image contains all resources associated with the image. All fields
+// MUST not be mutated directly after created.
+type Image struct {
+	// Id of the image. Normally the digest of image config.
+	ID string
+	// References are references to the image, e.g. RepoTag and RepoDigest.
+	References []string
+	// ChainID is the chainID of the image.
+	ChainID string
+	// Size is the compressed size of the image.
+	Size int64
+	// ImageSpec is the oci image structure which describes basic information about the image.
+	ImageSpec imagespec.Image
+	// Pinned image to prevent it from garbage collection
+	Pinned bool
+}
+
+// Getter is used to get images but does not make changes
+type Getter interface {
+	Get(ctx context.Context, name string) (images.Image, error)
+}
+
+// Store stores all images.
+type Store struct {
+	lock sync.RWMutex
+	// refCache is a containerd image reference to image id cache.
+	refCache map[string]string
+
+	// images is the local image store
+	images Getter
+
+	// content provider
+	provider content.InfoReaderProvider
+
+	// platform represents the currently supported platform for images
+	// TODO: Make this store multi-platform
+	platform platforms.MatchComparer
+
+	// store is the internal image store indexed by image id.
+	store *store
+}
+
+// NewStore creates an image store.
+func NewStore(img Getter, provider content.InfoReaderProvider, platform platforms.MatchComparer) *Store {
+	return &Store{
+		refCache: make(map[string]string),
+		images:   img,
+		provider: provider,
+		platform: platform,
+		store: &store{
+			images:     make(map[string]Image),
+			digestSet:  digestset.NewSet(),
+			pinnedRefs: make(map[string]sets.Set[string]),
+		},
+	}
+}
+
+// Update updates cache for a reference.
+func (s *Store) Update(ctx context.Context, ref string) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	i, err := s.images.Get(ctx, ref)
+	if err != nil && !errdefs.IsNotFound(err) {
+		return fmt.Errorf("get image from containerd: %w", err)
+	}
+
+	var img *Image
+	if err == nil {
+		img, err = s.getImage(ctx, i)
+		if err != nil {
+			return fmt.Errorf("get image info from containerd: %w", err)
+		}
+	}
+	return s.update(ref, img)
+}
+
+// update updates the internal cache. img == nil means that
+// the image does not exist in containerd.
+func (s *Store) update(ref string, img *Image) error {
+	oldID, oldExist := s.refCache[ref]
+	if img == nil {
+		// The image reference doesn't exist in containerd.
+		if oldExist {
+			// Remove the reference from the store.
+			s.store.delete(oldID, ref)
+			delete(s.refCache, ref)
+		}
+		return nil
+	}
+	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)
+	}
+	// New image. Add new image.
+	s.refCache[ref] = img.ID
+	return s.store.add(*img)
+}
+
+// getImage gets image information from containerd for current platform.
+func (s *Store) getImage(ctx context.Context, i images.Image) (*Image, error) {
+	diffIDs, err := i.RootFS(ctx, s.provider, s.platform)
+	if err != nil {
+		return nil, fmt.Errorf("get image diffIDs: %w", err)
+	}
+	chainID := imageidentity.ChainID(diffIDs)
+
+	size, err := usage.CalculateImageUsage(ctx, i, s.provider, usage.WithManifestLimit(s.platform, 1), usage.WithManifestUsage())
+	if err != nil {
+		return nil, fmt.Errorf("get image compressed resource size: %w", err)
+	}
+
+	desc, err := i.Config(ctx, s.provider, s.platform)
+	if err != nil {
+		return nil, fmt.Errorf("get image config descriptor: %w", err)
+	}
+	id := desc.Digest.String()
+
+	blob, err := content.ReadBlob(ctx, s.provider, desc)
+	if err != nil {
+		return nil, fmt.Errorf("read image config from content store: %w", err)
+	}
+
+	var spec imagespec.Image
+	if err := json.Unmarshal(blob, &spec); err != nil {
+		return nil, fmt.Errorf("unmarshal image config %s: %w", blob, err)
+	}
+
+	pinned := i.Labels[labels.PinnedImageLabelKey] == labels.PinnedImageLabelValue
+
+	return &Image{
+		ID:         id,
+		References: []string{i.Name},
+		ChainID:    chainID.String(),
+		Size:       size,
+		ImageSpec:  spec,
+		Pinned:     pinned,
+	}, nil
+
+}
+
+// Resolve resolves a image reference to image id.
+func (s *Store) Resolve(ref string) (string, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	id, ok := s.refCache[ref]
+	if !ok {
+		return "", errdefs.ErrNotFound
+	}
+	return id, nil
+}
+
+// Get gets image metadata by image id. The id can be truncated.
+// Returns various validation errors if the image id is invalid.
+// Returns errdefs.ErrNotFound if the image doesn't exist.
+func (s *Store) Get(id string) (Image, error) {
+	return s.store.get(id)
+}
+
+// List lists all images.
+func (s *Store) List() []Image {
+	return s.store.list()
+}
+
+type store struct {
+	lock       sync.RWMutex
+	images     map[string]Image
+	digestSet  *digestset.Set
+	pinnedRefs map[string]sets.Set[string]
+}
+
+func (s *store) list() []Image {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	var images []Image
+	for _, i := range s.images {
+		images = append(images, i)
+	}
+	return images
+}
+
+func (s *store) add(img Image) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	if _, err := s.digestSet.Lookup(img.ID); err != nil {
+		if err != digestset.ErrDigestNotFound {
+			return err
+		}
+		if err := s.digestSet.Add(imagedigest.Digest(img.ID)); err != nil {
+			return err
+		}
+	}
+
+	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.
+		s.images[img.ID] = img
+		return nil
+	}
+	// 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()
+	digest, err := s.digestSet.Lookup(id)
+	if err != nil {
+		if err == digestset.ErrDigestNotFound {
+			err = errdefs.ErrNotFound
+		}
+		return Image{}, err
+	}
+	if i, ok := s.images[digest.String()]; ok {
+		return i, nil
+	}
+	return Image{}, errdefs.ErrNotFound
+}
+
+func (s *store) delete(id, ref string) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	digest, err := s.digestSet.Lookup(id)
+	if err != nil {
+		// Note: The idIndex.Delete and delete doesn't handle truncated index.
+		// So we need to return if there are error.
+		return
+	}
+	i, ok := s.images[digest.String()]
+	if !ok {
+		return
+	}
+	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())
+}

internal/cri/store/image/image_test.go

@@ -0,0 +1,318 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package image
+
+import (
+	"sort"
+	"strings"
+	"testing"
+
+	"github.com/containerd/errdefs"
+	"k8s.io/apimachinery/pkg/util/sets"
+
+	"github.com/opencontainers/go-digest/digestset"
+	assertlib "github.com/stretchr/testify/assert"
+)
+
+func TestInternalStore(t *testing.T) {
+	images := []Image{
+		{
+			ID:         "sha256:1123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
+			ChainID:    "test-chain-id-1",
+			References: []string{"containerd.io/ref-1"},
+			Size:       10,
+		},
+		{
+			ID:         "sha256:2123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
+			ChainID:    "test-chain-id-2abcd",
+			References: []string{"containerd.io/ref-2abcd"},
+			Size:       20,
+		},
+		{
+			ID:         "sha256:3123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
+			References: []string{"containerd.io/ref-4a333"},
+			ChainID:    "test-chain-id-4a333",
+			Size:       30,
+		},
+		{
+			ID:         "sha256:4123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
+			References: []string{"containerd.io/ref-4abcd"},
+			ChainID:    "test-chain-id-4abcd",
+			Size:       40,
+		},
+	}
+	assert := assertlib.New(t)
+	genTruncIndex := func(normalName string) string { return normalName[:(len(normalName)+1)/2] }
+
+	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")
+	for _, img := range images {
+		err := s.add(img)
+		assert.NoError(err)
+	}
+
+	t.Logf("should be able to get image")
+	for _, v := range images {
+		truncID := genTruncIndex(v.ID)
+		got, err := s.get(truncID)
+		assert.NoError(err, "truncID:%s, fullID:%s", truncID, v.ID)
+		assert.Equal(v, got)
+	}
+
+	t.Logf("should be able to get image by truncated imageId without algorithm")
+	for _, v := range images {
+		truncID := genTruncIndex(v.ID[strings.Index(v.ID, ":")+1:])
+		got, err := s.get(truncID)
+		assert.NoError(err, "truncID:%s, fullID:%s", truncID, v.ID)
+		assert.Equal(v, got)
+	}
+
+	t.Logf("should not be able to get image by ambiguous prefix")
+	ambiguousPrefixs := []string{"sha256", "sha256:"}
+	for _, v := range ambiguousPrefixs {
+		_, err := s.get(v)
+		assert.NotEqual(nil, err)
+	}
+
+	t.Logf("should be able to list images")
+	imgs := s.list()
+	assert.Len(imgs, len(images))
+
+	imageNum := len(images)
+	for _, v := range images {
+		truncID := genTruncIndex(v.ID)
+		oldRef := v.References[0]
+		newRef := oldRef + "new"
+
+		t.Logf("should be able to add new references")
+		newImg := v
+		newImg.References = []string{newRef}
+		err := s.add(newImg)
+		assert.NoError(err)
+		got, err := s.get(truncID)
+		assert.NoError(err)
+		assert.Len(got.References, 2)
+		assert.Contains(got.References, oldRef, newRef)
+
+		t.Logf("should not be able to add duplicated references")
+		err = s.add(newImg)
+		assert.NoError(err)
+		got, err = s.get(truncID)
+		assert.NoError(err)
+		assert.Len(got.References, 2)
+		assert.Contains(got.References, oldRef, newRef)
+
+		t.Logf("should be able to delete image references")
+		s.delete(truncID, oldRef)
+		got, err = s.get(truncID)
+		assert.NoError(err)
+		assert.Equal([]string{newRef}, got.References)
+
+		t.Logf("should be able to delete image")
+		s.delete(truncID, newRef)
+		got, err = s.get(truncID)
+		assert.Equal(errdefs.ErrNotFound, err)
+		assert.Equal(Image{}, got)
+
+		imageNum--
+		imgs = s.list()
+		assert.Len(imgs, imageNum)
+	}
+}
+
+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"
+	image := Image{
+		ID:         id,
+		ChainID:    "test-chain-id-1",
+		References: []string{"containerd.io/ref-1"},
+		Size:       10,
+	}
+	assert := assertlib.New(t)
+
+	equal := func(i1, i2 Image) {
+		sort.Strings(i1.References)
+		sort.Strings(i2.References)
+		assert.Equal(i1, i2)
+	}
+	for desc, test := range map[string]struct {
+		ref      string
+		image    *Image
+		expected []Image
+	}{
+		"nothing should happen if a non-exist ref disappear": {
+			ref:      "containerd.io/ref-2",
+			image:    nil,
+			expected: []Image{image},
+		},
+		"new ref for an existing image": {
+			ref: "containerd.io/ref-2",
+			image: &Image{
+				ID:         id,
+				ChainID:    "test-chain-id-1",
+				References: []string{"containerd.io/ref-2"},
+				Size:       10,
+			},
+			expected: []Image{
+				{
+					ID:         id,
+					ChainID:    "test-chain-id-1",
+					References: []string{"containerd.io/ref-1", "containerd.io/ref-2"},
+					Size:       10,
+				},
+			},
+		},
+		"new ref for a new image": {
+			ref: "containerd.io/ref-2",
+			image: &Image{
+				ID:         newID,
+				ChainID:    "test-chain-id-2",
+				References: []string{"containerd.io/ref-2"},
+				Size:       20,
+			},
+			expected: []Image{
+				image,
+				{
+					ID:         newID,
+					ChainID:    "test-chain-id-2",
+					References: []string{"containerd.io/ref-2"},
+					Size:       20,
+				},
+			},
+		},
+		"existing ref point to a new image": {
+			ref: "containerd.io/ref-1",
+			image: &Image{
+				ID:         newID,
+				ChainID:    "test-chain-id-2",
+				References: []string{"containerd.io/ref-1"},
+				Size:       20,
+			},
+			expected: []Image{
+				{
+					ID:         newID,
+					ChainID:    "test-chain-id-2",
+					References: []string{"containerd.io/ref-1"},
+					Size:       20,
+				},
+			},
+		},
+		"existing ref disappear": {
+			ref:      "containerd.io/ref-1",
+			image:    nil,
+			expected: []Image{},
+		},
+	} {
+		t.Run(desc, func(t *testing.T) {
+			s, err := NewFakeStore([]Image{image})
+			assert.NoError(err)
+			assert.NoError(s.update(test.ref, test.image))
+
+			assert.Len(s.List(), len(test.expected))
+			for _, expect := range test.expected {
+				got, err := s.Get(expect.ID)
+				assert.NoError(err)
+				equal(got, expect)
+				for _, ref := range expect.References {
+					id, err := s.Resolve(ref)
+					assert.NoError(err)
+					assert.Equal(expect.ID, id)
+				}
+			}
+
+			if test.image == nil {
+				// Shouldn't be able to index by removed ref.
+				id, err := s.Resolve(test.ref)
+				assert.Equal(errdefs.ErrNotFound, err)
+				assert.Empty(id)
+			}
+		})
+	}
+}