Move CRI from pkg/ to internal/

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

internal/cri/store/container/container.go

@@ -0,0 +1,213 @@
+/*
+   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 container
+
+import (
+	"sync"
+
+	containerd "github.com/containerd/containerd/v2/client"
+	"github.com/containerd/containerd/v2/internal/truncindex"
+	cio "github.com/containerd/containerd/v2/internal/cri/io"
+	"github.com/containerd/containerd/v2/internal/cri/store"
+	"github.com/containerd/containerd/v2/internal/cri/store/label"
+	"github.com/containerd/containerd/v2/internal/cri/store/stats"
+	"github.com/containerd/errdefs"
+
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+// Container contains all resources associated with the container. All methods to
+// mutate the internal state are thread-safe.
+type Container struct {
+	// Metadata is the metadata of the container, it is **immutable** after created.
+	Metadata
+	// Status stores the status of the container.
+	Status StatusStorage
+	// Container is the containerd container client.
+	Container containerd.Container
+	// Container IO.
+	// IO could only be nil when the container is in unknown state.
+	IO *cio.ContainerIO
+	// StopCh is used to propagate the stop information of the container.
+	*store.StopCh
+	// IsStopSignaledWithTimeout the default is 0, and it is set to 1 after sending
+	// the signal once to avoid repeated sending of the signal.
+	IsStopSignaledWithTimeout *uint32
+	// Stats contains (mutable) stats for the container
+	Stats *stats.ContainerStats
+}
+
+// Opts sets specific information to newly created Container.
+type Opts func(*Container) error
+
+// WithContainer adds the containerd Container to the internal data store.
+func WithContainer(cntr containerd.Container) Opts {
+	return func(c *Container) error {
+		c.Container = cntr
+		return nil
+	}
+}
+
+// WithContainerIO adds IO into the container.
+func WithContainerIO(io *cio.ContainerIO) Opts {
+	return func(c *Container) error {
+		c.IO = io
+		return nil
+	}
+}
+
+// WithStatus adds status to the container.
+func WithStatus(status Status, root string) Opts {
+	return func(c *Container) error {
+		s, err := StoreStatus(root, c.ID, status)
+		if err != nil {
+			return err
+		}
+		c.Status = s
+		if s.Get().State() == runtime.ContainerState_CONTAINER_EXITED {
+			c.Stop()
+		}
+		return nil
+	}
+}
+
+// NewContainer creates an internally used container type.
+func NewContainer(metadata Metadata, opts ...Opts) (Container, error) {
+	c := Container{
+		Metadata:                  metadata,
+		StopCh:                    store.NewStopCh(),
+		IsStopSignaledWithTimeout: new(uint32),
+	}
+	for _, o := range opts {
+		if err := o(&c); err != nil {
+			return Container{}, err
+		}
+	}
+	return c, nil
+}
+
+// Delete deletes checkpoint for the container.
+func (c *Container) Delete() error {
+	return c.Status.Delete()
+}
+
+// Store stores all Containers.
+type Store struct {
+	lock       sync.RWMutex
+	containers map[string]Container
+	idIndex    *truncindex.TruncIndex
+	labels     *label.Store
+}
+
+// NewStore creates a container store.
+func NewStore(labels *label.Store) *Store {
+	return &Store{
+		containers: make(map[string]Container),
+		idIndex:    truncindex.NewTruncIndex([]string{}),
+		labels:     labels,
+	}
+}
+
+// Add a container into the store. Returns errdefs.ErrAlreadyExists if the
+// container already exists.
+func (s *Store) Add(c Container) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	if _, ok := s.containers[c.ID]; ok {
+		return errdefs.ErrAlreadyExists
+	}
+	if err := s.labels.Reserve(c.ProcessLabel); err != nil {
+		return err
+	}
+	if err := s.idIndex.Add(c.ID); err != nil {
+		return err
+	}
+	s.containers[c.ID] = c
+	return nil
+}
+
+// Get returns the container with specified id. Returns errdefs.ErrNotFound
+// if the container doesn't exist.
+func (s *Store) Get(id string) (Container, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = errdefs.ErrNotFound
+		}
+		return Container{}, err
+	}
+	if c, ok := s.containers[id]; ok {
+		return c, nil
+	}
+	return Container{}, errdefs.ErrNotFound
+}
+
+// List lists all containers.
+func (s *Store) List() []Container {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	var containers []Container
+	for _, c := range s.containers {
+		containers = append(containers, c)
+	}
+	return containers
+}
+
+// UpdateContainerStats updates the container specified by ID with the
+// stats present in 'newContainerStats'. Returns errdefs.ErrNotFound
+// if the container does not exist in the store.
+func (s *Store) UpdateContainerStats(id string, newContainerStats *stats.ContainerStats) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = errdefs.ErrNotFound
+		}
+		return err
+	}
+
+	if _, ok := s.containers[id]; !ok {
+		return errdefs.ErrNotFound
+	}
+
+	c := s.containers[id]
+	c.Stats = newContainerStats
+	s.containers[id] = c
+	return nil
+}
+
+// Delete deletes the container from store with specified id.
+func (s *Store) Delete(id string) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(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
+	}
+	c := s.containers[id]
+	if c.IO != nil {
+		c.IO.Close()
+	}
+	s.labels.Release(c.ProcessLabel)
+	s.idIndex.Delete(id)
+	delete(s.containers, id)
+}

internal/cri/store/container/container_test.go

@@ -0,0 +1,281 @@
+/*
+   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 container
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	cio "github.com/containerd/containerd/v2/internal/cri/io"
+	"github.com/containerd/containerd/v2/internal/cri/store/label"
+	"github.com/containerd/containerd/v2/internal/cri/store/stats"
+	"github.com/containerd/errdefs"
+
+	"github.com/opencontainers/selinux/go-selinux"
+	assertlib "github.com/stretchr/testify/assert"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+func TestContainerStore(t *testing.T) {
+	metadatas := map[string]Metadata{
+		"1": {
+			ID:        "1",
+			Name:      "Container-1",
+			SandboxID: "Sandbox-1",
+			Config: &runtime.ContainerConfig{
+				Metadata: &runtime.ContainerMetadata{
+					Name:    "TestPod-1",
+					Attempt: 1,
+				},
+			},
+			ImageRef:     "TestImage-1",
+			StopSignal:   "SIGTERM",
+			LogPath:      "/test/log/path/1",
+			ProcessLabel: "junk:junk:junk:c1,c2",
+		},
+		"2abcd": {
+			ID:        "2abcd",
+			Name:      "Container-2abcd",
+			SandboxID: "Sandbox-2abcd",
+			Config: &runtime.ContainerConfig{
+				Metadata: &runtime.ContainerMetadata{
+					Name:    "TestPod-2abcd",
+					Attempt: 2,
+				},
+			},
+			StopSignal:   "SIGTERM",
+			ImageRef:     "TestImage-2",
+			LogPath:      "/test/log/path/2",
+			ProcessLabel: "junk:junk:junk:c1,c2",
+		},
+		"4a333": {
+			ID:        "4a333",
+			Name:      "Container-4a333",
+			SandboxID: "Sandbox-4a333",
+			Config: &runtime.ContainerConfig{
+				Metadata: &runtime.ContainerMetadata{
+					Name:    "TestPod-4a333",
+					Attempt: 3,
+				},
+			},
+			StopSignal:   "SIGTERM",
+			ImageRef:     "TestImage-3",
+			LogPath:      "/test/log/path/3",
+			ProcessLabel: "junk:junk:junk:c1,c3",
+		},
+		"4abcd": {
+			ID:        "4abcd",
+			Name:      "Container-4abcd",
+			SandboxID: "Sandbox-4abcd",
+			Config: &runtime.ContainerConfig{
+				Metadata: &runtime.ContainerMetadata{
+					Name:    "TestPod-4abcd",
+					Attempt: 1,
+				},
+			},
+			StopSignal:   "SIGTERM",
+			ImageRef:     "TestImage-4abcd",
+			ProcessLabel: "junk:junk:junk:c1,c4",
+		},
+	}
+	statuses := map[string]Status{
+		"1": {
+			Pid:        1,
+			CreatedAt:  time.Now().UnixNano(),
+			StartedAt:  time.Now().UnixNano(),
+			FinishedAt: time.Now().UnixNano(),
+			ExitCode:   1,
+			Reason:     "TestReason-1",
+			Message:    "TestMessage-1",
+		},
+		"2abcd": {
+			Pid:        2,
+			CreatedAt:  time.Now().UnixNano(),
+			StartedAt:  time.Now().UnixNano(),
+			FinishedAt: time.Now().UnixNano(),
+			ExitCode:   2,
+			Reason:     "TestReason-2abcd",
+			Message:    "TestMessage-2abcd",
+		},
+		"4a333": {
+			Pid:        3,
+			CreatedAt:  time.Now().UnixNano(),
+			StartedAt:  time.Now().UnixNano(),
+			FinishedAt: time.Now().UnixNano(),
+			ExitCode:   3,
+			Reason:     "TestReason-4a333",
+			Message:    "TestMessage-4a333",
+			Starting:   true,
+		},
+		"4abcd": {
+			Pid:        4,
+			CreatedAt:  time.Now().UnixNano(),
+			StartedAt:  time.Now().UnixNano(),
+			FinishedAt: time.Now().UnixNano(),
+			ExitCode:   4,
+			Reason:     "TestReason-4abcd",
+			Message:    "TestMessage-4abcd",
+			Removing:   true,
+		},
+	}
+
+	stats := map[string]*stats.ContainerStats{
+		"1": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 1,
+		},
+		"2abcd": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 2,
+		},
+		"4a333": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 3,
+		},
+		"4abcd": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 4,
+		},
+	}
+	assert := assertlib.New(t)
+	containers := map[string]Container{}
+	for id := range metadatas {
+		container, err := NewContainer(
+			metadatas[id],
+			WithFakeStatus(statuses[id]),
+		)
+		assert.NoError(err)
+		containers[id] = container
+	}
+
+	s := NewStore(label.NewStore())
+	reserved := map[string]bool{}
+	s.labels.Reserver = func(label string) {
+		reserved[strings.SplitN(label, ":", 4)[3]] = true
+	}
+	s.labels.Releaser = func(label string) {
+		reserved[strings.SplitN(label, ":", 4)[3]] = false
+	}
+
+	t.Logf("should be able to add container")
+	for _, c := range containers {
+		assert.NoError(s.Add(c))
+	}
+
+	t.Logf("should be able to get container")
+	genTruncIndex := func(normalName string) string { return normalName[:(len(normalName)+1)/2] }
+	for id, c := range containers {
+		got, err := s.Get(genTruncIndex(id))
+		assert.NoError(err)
+		assert.Equal(c, got)
+		assert.Nil(c.Stats)
+	}
+
+	t.Logf("should be able to list containers")
+	cs := s.List()
+	assert.Len(cs, len(containers))
+
+	t.Logf("should be able to update stats on container")
+	for id := range containers {
+		err := s.UpdateContainerStats(id, stats[id])
+		assert.NoError(err)
+	}
+
+	// Validate stats were updated
+	cs = s.List()
+	assert.Len(cs, len(containers))
+	for _, c := range cs {
+		assert.Equal(stats[c.ID], c.Stats)
+	}
+
+	if selinux.GetEnabled() {
+		t.Logf("should have reserved labels (requires -tag selinux)")
+		assert.Equal(map[string]bool{
+			"c1,c2": true,
+			"c1,c3": true,
+			"c1,c4": true,
+		}, reserved)
+	}
+
+	cntrNum := len(containers)
+	for testID, v := range containers {
+		truncID := genTruncIndex(testID)
+
+		t.Logf("add should return already exists error for duplicated container")
+		assert.Equal(errdefs.ErrAlreadyExists, s.Add(v))
+
+		t.Logf("should be able to delete container")
+		s.Delete(truncID)
+		cntrNum--
+		cs = s.List()
+		assert.Len(cs, cntrNum)
+
+		t.Logf("get should return not exist error after deletion")
+		c, err := s.Get(truncID)
+		assert.Equal(Container{}, c)
+		assert.Equal(errdefs.ErrNotFound, err)
+	}
+
+	if selinux.GetEnabled() {
+		t.Logf("should have released all labels (requires -tag selinux)")
+		assert.Equal(map[string]bool{
+			"c1,c2": false,
+			"c1,c3": false,
+			"c1,c4": false,
+		}, reserved)
+	}
+}
+
+func TestWithContainerIO(t *testing.T) {
+	meta := Metadata{
+		ID:        "1",
+		Name:      "Container-1",
+		SandboxID: "Sandbox-1",
+		Config: &runtime.ContainerConfig{
+			Metadata: &runtime.ContainerMetadata{
+				Name:    "TestPod-1",
+				Attempt: 1,
+			},
+		},
+		ImageRef:   "TestImage-1",
+		StopSignal: "SIGTERM",
+		LogPath:    "/test/log/path",
+	}
+	status := Status{
+		Pid:        1,
+		CreatedAt:  time.Now().UnixNano(),
+		StartedAt:  time.Now().UnixNano(),
+		FinishedAt: time.Now().UnixNano(),
+		ExitCode:   1,
+		Reason:     "TestReason-1",
+		Message:    "TestMessage-1",
+	}
+	assert := assertlib.New(t)
+
+	c, err := NewContainer(meta, WithFakeStatus(status))
+	assert.NoError(err)
+	assert.Nil(c.IO)
+
+	c, err = NewContainer(
+		meta,
+		WithFakeStatus(status),
+		WithContainerIO(&cio.ContainerIO{}),
+	)
+	assert.NoError(err)
+	assert.NotNil(c.IO)
+}

internal/cri/store/container/fake_status.go

@@ -0,0 +1,62 @@
+/*
+   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 container
+
+import "sync"
+
+// WithFakeStatus adds fake status to the container.
+func WithFakeStatus(status Status) Opts {
+	return func(c *Container) error {
+		c.Status = &fakeStatusStorage{status: status}
+		if status.FinishedAt != 0 {
+			// Fake the TaskExit event
+			c.Stop()
+		}
+		return nil
+	}
+}
+
+// fakeStatusStorage is a fake status storage for testing.
+type fakeStatusStorage struct {
+	sync.RWMutex
+	status Status
+}
+
+func (f *fakeStatusStorage) Get() Status {
+	f.RLock()
+	defer f.RUnlock()
+	return f.status
+}
+
+func (f *fakeStatusStorage) UpdateSync(u UpdateFunc) error {
+	return f.Update(u)
+}
+
+func (f *fakeStatusStorage) Update(u UpdateFunc) error {
+	f.Lock()
+	defer f.Unlock()
+	newStatus, err := u(f.status)
+	if err != nil {
+		return err
+	}
+	f.status = newStatus
+	return nil
+}
+
+func (f *fakeStatusStorage) Delete() error {
+	return nil
+}

internal/cri/store/container/metadata.go

@@ -0,0 +1,88 @@
+/*
+   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 container
+
+import (
+	"encoding/json"
+	"fmt"
+
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+// NOTE(random-liu):
+// 1) Metadata is immutable after created.
+// 2) Metadata is checkpointed as containerd container label.
+
+// metadataVersion is current version of container metadata.
+const metadataVersion = "v1"
+
+// versionedMetadata is the internal versioned container metadata.
+type versionedMetadata struct {
+	// Version indicates the version of the versioned container metadata.
+	Version string
+	// Metadata's type is metadataInternal. If not there will be a recursive call in MarshalJSON.
+	Metadata metadataInternal
+}
+
+// metadataInternal is for internal use.
+type metadataInternal Metadata
+
+// Metadata is the unversioned container metadata.
+type Metadata struct {
+	// ID is the container id.
+	ID string
+	// Name is the container name.
+	Name string
+	// SandboxID is the sandbox id the container belongs to.
+	SandboxID string
+	// Config is the CRI container config.
+	// NOTE(random-liu): Resource limits are updatable, the source
+	// of truth for resource limits are in containerd.
+	Config *runtime.ContainerConfig
+	// ImageRef is the reference of image used by the container.
+	ImageRef string
+	// LogPath is the container log path.
+	LogPath string
+	// StopSignal is the system call signal that will be sent to the container to exit.
+	// TODO(random-liu): Add integration test for stop signal.
+	StopSignal string
+	// ProcessLabel is the SELinux process label for the container
+	ProcessLabel string
+}
+
+// MarshalJSON encodes Metadata into bytes in json format.
+func (c *Metadata) MarshalJSON() ([]byte, error) {
+	return json.Marshal(&versionedMetadata{
+		Version:  metadataVersion,
+		Metadata: metadataInternal(*c),
+	})
+}
+
+// UnmarshalJSON decodes Metadata from bytes.
+func (c *Metadata) UnmarshalJSON(data []byte) error {
+	versioned := &versionedMetadata{}
+	if err := json.Unmarshal(data, versioned); err != nil {
+		return err
+	}
+	// Handle old version after upgrade.
+	switch versioned.Version {
+	case metadataVersion:
+		*c = Metadata(versioned.Metadata)
+		return nil
+	}
+	return fmt.Errorf("unsupported version: %q", versioned.Version)
+}

internal/cri/store/container/metadata_test.go

@@ -0,0 +1,81 @@
+/*
+   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 container
+
+import (
+	"encoding/json"
+	"testing"
+
+	assertlib "github.com/stretchr/testify/assert"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+func TestMetadataMarshalUnmarshal(t *testing.T) {
+	meta := &Metadata{
+		ID:        "test-id",
+		Name:      "test-name",
+		SandboxID: "test-sandbox-id",
+		Config: &runtime.ContainerConfig{
+			Metadata: &runtime.ContainerMetadata{
+				Name:    "test-name",
+				Attempt: 1,
+			},
+		},
+		ImageRef: "test-image-ref",
+		LogPath:  "/test/log/path",
+	}
+
+	assert := assertlib.New(t)
+	newMeta := &Metadata{}
+	newVerMeta := &versionedMetadata{}
+
+	t.Logf("should be able to do json.marshal")
+	data, err := json.Marshal(meta)
+	assert.NoError(err)
+	data1, err := json.Marshal(&versionedMetadata{
+		Version:  metadataVersion,
+		Metadata: metadataInternal(*meta),
+	})
+	assert.NoError(err)
+	assert.Equal(data, data1)
+
+	t.Logf("should be able to do MarshalJSON")
+	data, err = meta.MarshalJSON()
+	assert.NoError(err)
+	assert.NoError(newMeta.UnmarshalJSON(data))
+	assert.Equal(meta, newMeta)
+
+	t.Logf("should be able to do MarshalJSON and json.Unmarshal")
+	data, err = meta.MarshalJSON()
+	assert.NoError(err)
+	assert.NoError(json.Unmarshal(data, newVerMeta))
+	assert.Equal(meta, (*Metadata)(&newVerMeta.Metadata))
+
+	t.Logf("should be able to do json.Marshal and UnmarshalJSON")
+	data, err = json.Marshal(meta)
+	assert.NoError(err)
+	assert.NoError(newMeta.UnmarshalJSON(data))
+	assert.Equal(meta, newMeta)
+
+	t.Logf("should json.Unmarshal fail for unsupported version")
+	unsupported, err := json.Marshal(&versionedMetadata{
+		Version:  "random-test-version",
+		Metadata: metadataInternal(*meta),
+	})
+	assert.NoError(err)
+	assert.Error(json.Unmarshal(unsupported, &newMeta))
+}

internal/cri/store/container/status.go

@@ -0,0 +1,299 @@
+/*
+   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 container
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sync"
+
+	"github.com/containerd/continuity"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+// The container state machine in the CRI plugin:
+//
+//                         +              +
+//                         |              |
+//                         | Create       | Load
+//                         |              |
+//                    +----v----+         |
+//                    |         |         |
+//                    | CREATED <---------+-----------+
+//                    |         |         |           |
+//                    +----+-----         |           |
+//                         |              |           |
+//                         | Start        |           |
+//                         |              |           |
+//                    +----v----+         |           |
+//   Exec    +--------+         |         |           |
+//  Attach   |        | RUNNING <---------+           |
+// LogReopen +-------->         |         |           |
+//                    +----+----+         |           |
+//                         |              |           |
+//                         | Stop/Exit    |           |
+//                         |              |           |
+//                    +----v----+         |           |
+//                    |         <---------+      +----v----+
+//                    |  EXITED |                |         |
+//                    |         <----------------+ UNKNOWN |
+//                    +----+----+       Stop     |         |
+//                         |                     +---------+
+//                         | Remove
+//                         v
+//                      DELETED
+
+// statusVersion is current version of container status.
+const statusVersion = "v1"
+
+// versionedStatus is the internal used versioned container status.
+type versionedStatus struct {
+	// Version indicates the version of the versioned container status.
+	Version string
+	Status
+}
+
+// Status is the status of a container.
+type Status struct {
+	// Pid is the init process id of the container.
+	Pid uint32
+	// CreatedAt is the created timestamp.
+	CreatedAt int64
+	// StartedAt is the started timestamp.
+	StartedAt int64
+	// FinishedAt is the finished timestamp.
+	FinishedAt int64
+	// ExitCode is the container exit code.
+	ExitCode int32
+	// CamelCase string explaining why container is in its current state.
+	Reason string
+	// Human-readable message indicating details about why container is in its
+	// current state.
+	Message string
+	// Starting indicates that the container is in starting state.
+	// This field doesn't need to be checkpointed.
+	Starting bool `json:"-"`
+	// Removing indicates that the container is in removing state.
+	// This field doesn't need to be checkpointed.
+	Removing bool `json:"-"`
+	// Unknown indicates that the container status is not fully loaded.
+	// This field doesn't need to be checkpointed.
+	Unknown bool `json:"-"`
+	// Resources has container runtime resource constraints
+	Resources *runtime.ContainerResources
+}
+
+// State returns current state of the container based on the container status.
+func (s Status) State() runtime.ContainerState {
+	if s.Unknown {
+		return runtime.ContainerState_CONTAINER_UNKNOWN
+	}
+	if s.FinishedAt != 0 {
+		return runtime.ContainerState_CONTAINER_EXITED
+	}
+	if s.StartedAt != 0 {
+		return runtime.ContainerState_CONTAINER_RUNNING
+	}
+	if s.CreatedAt != 0 {
+		return runtime.ContainerState_CONTAINER_CREATED
+	}
+	return runtime.ContainerState_CONTAINER_UNKNOWN
+}
+
+// encode encodes Status into bytes in json format.
+func (s *Status) encode() ([]byte, error) {
+	return json.Marshal(&versionedStatus{
+		Version: statusVersion,
+		Status:  *s,
+	})
+}
+
+// decode decodes Status from bytes.
+func (s *Status) decode(data []byte) error {
+	versioned := &versionedStatus{}
+	if err := json.Unmarshal(data, versioned); err != nil {
+		return err
+	}
+	// Handle old version after upgrade.
+	switch versioned.Version {
+	case statusVersion:
+		*s = versioned.Status
+		return nil
+	}
+	return errors.New("unsupported version")
+}
+
+// UpdateFunc is function used to update the container status. If there
+// is an error, the update will be rolled back.
+type UpdateFunc func(Status) (Status, error)
+
+// StatusStorage manages the container status with a storage backend.
+type StatusStorage interface {
+	// Get a container status.
+	Get() Status
+	// UpdateSync updates the container status and the on disk checkpoint.
+	// Note that the update MUST be applied in one transaction.
+	UpdateSync(UpdateFunc) error
+	// Update the container status. Note that the update MUST be applied
+	// in one transaction.
+	Update(UpdateFunc) error
+	// Delete the container status.
+	// Note:
+	// * Delete should be idempotent.
+	// * The status must be deleted in one transaction.
+	Delete() error
+}
+
+// StoreStatus creates the storage containing the passed in container status with the
+// specified id.
+// The status MUST be created in one transaction.
+func StoreStatus(root, id string, status Status) (StatusStorage, error) {
+	data, err := status.encode()
+	if err != nil {
+		return nil, fmt.Errorf("failed to encode status: %w", err)
+	}
+	path := filepath.Join(root, "status")
+	if err := continuity.AtomicWriteFile(path, data, 0600); err != nil {
+		return nil, fmt.Errorf("failed to checkpoint status to %q: %w", path, err)
+	}
+	return &statusStorage{
+		path:   path,
+		status: status,
+	}, nil
+}
+
+// LoadStatus loads container status from checkpoint. There shouldn't be threads
+// writing to the file during loading.
+func LoadStatus(root, id string) (Status, error) {
+	path := filepath.Join(root, "status")
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return Status{}, fmt.Errorf("failed to read status from %q: %w", path, err)
+	}
+	var status Status
+	if err := status.decode(data); err != nil {
+		return Status{}, fmt.Errorf("failed to decode status %q: %w", data, err)
+	}
+	return status, nil
+}
+
+type statusStorage struct {
+	sync.RWMutex
+	path   string
+	status Status
+}
+
+// Get a copy of container status.
+func (s *statusStorage) Get() Status {
+	s.RLock()
+	defer s.RUnlock()
+	// Deep copy is needed in case some fields in Status are updated after Get()
+	// is called.
+	return deepCopyOf(s.status)
+}
+
+func deepCopyOf(s Status) Status {
+	copy := s
+	// Resources is the only field that is a pointer, and therefore needs
+	// a manual deep copy.
+	// This will need updates when new fields are added to ContainerResources.
+	if s.Resources == nil {
+		return copy
+	}
+	copy.Resources = &runtime.ContainerResources{}
+	if s.Resources != nil && s.Resources.Linux != nil {
+		hugepageLimits := make([]*runtime.HugepageLimit, 0, len(s.Resources.Linux.HugepageLimits))
+		for _, l := range s.Resources.Linux.HugepageLimits {
+			if l != nil {
+				hugepageLimits = append(hugepageLimits, &runtime.HugepageLimit{
+					PageSize: l.PageSize,
+					Limit:    l.Limit,
+				})
+			}
+		}
+		copy.Resources = &runtime.ContainerResources{
+			Linux: &runtime.LinuxContainerResources{
+				CpuPeriod:              s.Resources.Linux.CpuPeriod,
+				CpuQuota:               s.Resources.Linux.CpuQuota,
+				CpuShares:              s.Resources.Linux.CpuShares,
+				CpusetCpus:             s.Resources.Linux.CpusetCpus,
+				CpusetMems:             s.Resources.Linux.CpusetMems,
+				MemoryLimitInBytes:     s.Resources.Linux.MemoryLimitInBytes,
+				MemorySwapLimitInBytes: s.Resources.Linux.MemorySwapLimitInBytes,
+				OomScoreAdj:            s.Resources.Linux.OomScoreAdj,
+				Unified:                s.Resources.Linux.Unified,
+				HugepageLimits:         hugepageLimits,
+			},
+		}
+	}
+
+	if s.Resources != nil && s.Resources.Windows != nil {
+		copy.Resources = &runtime.ContainerResources{
+			Windows: &runtime.WindowsContainerResources{
+				CpuShares:          s.Resources.Windows.CpuShares,
+				CpuCount:           s.Resources.Windows.CpuCount,
+				CpuMaximum:         s.Resources.Windows.CpuMaximum,
+				MemoryLimitInBytes: s.Resources.Windows.MemoryLimitInBytes,
+				RootfsSizeInBytes:  s.Resources.Windows.RootfsSizeInBytes,
+			},
+		}
+	}
+	return copy
+}
+
+// UpdateSync updates the container status and the on disk checkpoint.
+func (s *statusStorage) UpdateSync(u UpdateFunc) error {
+	s.Lock()
+	defer s.Unlock()
+	newStatus, err := u(s.status)
+	if err != nil {
+		return err
+	}
+	data, err := newStatus.encode()
+	if err != nil {
+		return fmt.Errorf("failed to encode status: %w", err)
+	}
+	if err := continuity.AtomicWriteFile(s.path, data, 0600); err != nil {
+		return fmt.Errorf("failed to checkpoint status to %q: %w", s.path, err)
+	}
+	s.status = newStatus
+	return nil
+}
+
+// Update the container status.
+func (s *statusStorage) Update(u UpdateFunc) error {
+	s.Lock()
+	defer s.Unlock()
+	newStatus, err := u(s.status)
+	if err != nil {
+		return err
+	}
+	s.status = newStatus
+	return nil
+}
+
+// Delete deletes the container status from disk atomically.
+func (s *statusStorage) Delete() error {
+	temp := filepath.Dir(s.path) + ".del-" + filepath.Base(s.path)
+	if err := os.Rename(s.path, temp); err != nil && !os.IsNotExist(err) {
+		return err
+	}
+	return os.RemoveAll(temp)
+}

internal/cri/store/container/status_test.go

@@ -0,0 +1,188 @@
+/*
+   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 container
+
+import (
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	assertlib "github.com/stretchr/testify/assert"
+	requirelib "github.com/stretchr/testify/require"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+func TestContainerState(t *testing.T) {
+	for c, test := range map[string]struct {
+		status Status
+		state  runtime.ContainerState
+	}{
+		"unknown state": {
+			status: Status{
+				Unknown: true,
+			},
+			state: runtime.ContainerState_CONTAINER_UNKNOWN,
+		},
+		"unknown state because there is no timestamp set": {
+			status: Status{},
+			state:  runtime.ContainerState_CONTAINER_UNKNOWN,
+		},
+		"created state": {
+			status: Status{
+				CreatedAt: time.Now().UnixNano(),
+			},
+			state: runtime.ContainerState_CONTAINER_CREATED,
+		},
+		"running state": {
+			status: Status{
+				CreatedAt: time.Now().UnixNano(),
+				StartedAt: time.Now().UnixNano(),
+			},
+			state: runtime.ContainerState_CONTAINER_RUNNING,
+		},
+		"exited state": {
+			status: Status{
+				CreatedAt:  time.Now().UnixNano(),
+				FinishedAt: time.Now().UnixNano(),
+			},
+			state: runtime.ContainerState_CONTAINER_EXITED,
+		},
+	} {
+		t.Run(c, func(t *testing.T) {
+			assertlib.Equal(t, test.state, test.status.State())
+		})
+	}
+}
+
+func TestStatusEncodeDecode(t *testing.T) {
+	s := &Status{
+		Pid:        1234,
+		CreatedAt:  time.Now().UnixNano(),
+		StartedAt:  time.Now().UnixNano(),
+		FinishedAt: time.Now().UnixNano(),
+		ExitCode:   1,
+		Reason:     "test-reason",
+		Message:    "test-message",
+		Removing:   true,
+		Starting:   true,
+		Unknown:    true,
+	}
+	assert := assertlib.New(t)
+	data, err := s.encode()
+	assert.NoError(err)
+	newS := &Status{}
+	assert.NoError(newS.decode(data))
+	s.Removing = false // Removing should not be encoded.
+	s.Starting = false // Starting should not be encoded.
+	s.Unknown = false  // Unknown should not be encoded.
+	assert.Equal(s, newS)
+
+	unsupported, err := json.Marshal(&versionedStatus{
+		Version: "random-test-version",
+		Status:  *s,
+	})
+	assert.NoError(err)
+	assert.Error(newS.decode(unsupported))
+}
+
+func TestStatus(t *testing.T) {
+	testID := "test-id"
+	testStatus := Status{
+		CreatedAt: time.Now().UnixNano(),
+	}
+	updateStatus := Status{
+		CreatedAt: time.Now().UnixNano(),
+		StartedAt: time.Now().UnixNano(),
+	}
+	updateErr := errors.New("update error")
+	assert := assertlib.New(t)
+	require := requirelib.New(t)
+
+	tempDir := t.TempDir()
+	statusFile := filepath.Join(tempDir, "status")
+
+	t.Logf("simple store and get")
+	s, err := StoreStatus(tempDir, testID, testStatus)
+	assert.NoError(err)
+	old := s.Get()
+	assert.Equal(testStatus, old)
+	_, err = os.Stat(statusFile)
+	assert.NoError(err)
+	loaded, err := LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+
+	t.Logf("failed update should not take effect")
+	err = s.Update(func(o Status) (Status, error) {
+		return updateStatus, updateErr
+	})
+	assert.Equal(updateErr, err)
+	assert.Equal(testStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+
+	t.Logf("successful update should take effect but not checkpoint")
+	err = s.Update(func(o Status) (Status, error) {
+		return updateStatus, nil
+	})
+	assert.NoError(err)
+	assert.Equal(updateStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+	// Recover status.
+	assert.NoError(s.Update(func(o Status) (Status, error) {
+		return testStatus, nil
+	}))
+
+	t.Logf("failed update sync should not take effect")
+	err = s.UpdateSync(func(o Status) (Status, error) {
+		return updateStatus, updateErr
+	})
+	assert.Equal(updateErr, err)
+	assert.Equal(testStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+
+	t.Logf("successful update sync should take effect and checkpoint")
+	err = s.UpdateSync(func(o Status) (Status, error) {
+		return updateStatus, nil
+	})
+	assert.NoError(err)
+	assert.Equal(updateStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(updateStatus, loaded)
+
+	t.Logf("successful update should not affect existing snapshot")
+	assert.Equal(testStatus, old)
+
+	t.Logf("delete status")
+	assert.NoError(s.Delete())
+	_, err = LoadStatus(tempDir, testID)
+	assert.Error(err)
+	_, err = os.Stat(statusFile)
+	assert.True(os.IsNotExist(err))
+
+	t.Logf("delete status should be idempotent")
+	assert.NoError(s.Delete())
+}

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)
+			}
+		})
+	}
+}

internal/cri/store/label/label.go

@@ -0,0 +1,97 @@
+/*
+   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 label
+
+import (
+	"sync"
+
+	"github.com/opencontainers/selinux/go-selinux"
+)
+
+// Store is used to store SELinux process labels
+type Store struct {
+	sync.Mutex
+	levels   map[string]int
+	Releaser func(string)
+	Reserver func(string)
+}
+
+// NewStore creates a new SELinux process label store
+func NewStore() *Store {
+	return &Store{
+		levels:   map[string]int{},
+		Releaser: selinux.ReleaseLabel,
+		Reserver: selinux.ReserveLabel,
+	}
+}
+
+// Reserve reserves the MLS/MCS level component of the specified label
+// and prevents multiple reserves for the same level
+func (s *Store) Reserve(label string) error {
+	s.Lock()
+	defer s.Unlock()
+
+	context, err := selinux.NewContext(label)
+	if err != nil {
+		return err
+	}
+
+	level := context["level"]
+	// no reason to count empty
+	if level == "" {
+		return nil
+	}
+
+	if _, ok := s.levels[level]; !ok {
+		s.Reserver(label)
+	}
+
+	s.levels[level]++
+	return nil
+}
+
+// Release un-reserves the MLS/MCS level component of the specified label,
+// allowing it to be used by another process once labels with the same
+// level have been released.
+func (s *Store) Release(label string) {
+	s.Lock()
+	defer s.Unlock()
+
+	context, err := selinux.NewContext(label)
+	if err != nil {
+		return
+	}
+
+	level := context["level"]
+	if level == "" {
+		return
+	}
+
+	count, ok := s.levels[level]
+	if !ok {
+		return
+	}
+	switch {
+	case count == 1:
+		s.Releaser(label)
+		delete(s.levels, level)
+	case count < 1:
+		delete(s.levels, level)
+	case count > 1:
+		s.levels[level] = count - 1
+	}
+}

internal/cri/store/label/label_test.go

@@ -0,0 +1,116 @@
+/*
+   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 label
+
+import (
+	"testing"
+
+	"github.com/opencontainers/selinux/go-selinux"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddThenRemove(t *testing.T) {
+	if !selinux.GetEnabled() {
+		t.Skip("selinux is not enabled")
+	}
+
+	assert := assert.New(t)
+	store := NewStore()
+	releaseCount := 0
+	reserveCount := 0
+	store.Releaser = func(label string) {
+		assert.Contains(label, ":c1,c2")
+		releaseCount++
+		assert.Equal(1, releaseCount)
+	}
+	store.Reserver = func(label string) {
+		assert.Contains(label, ":c1,c2")
+		reserveCount++
+		assert.Equal(1, reserveCount)
+	}
+
+	t.Log("should count to two level")
+	assert.NoError(store.Reserve("junk:junk:junk:c1,c2"))
+	assert.NoError(store.Reserve("junk2:junk2:junk2:c1,c2"))
+
+	t.Log("should have one item")
+	assert.Equal(1, len(store.levels))
+
+	t.Log("c1,c2 count should be 2")
+	assert.Equal(2, store.levels["c1,c2"])
+
+	store.Release("junk:junk:junk:c1,c2")
+	store.Release("junk2:junk2:junk2:c1,c2")
+
+	t.Log("should have 0 items")
+	assert.Equal(0, len(store.levels))
+
+	t.Log("should have reserved")
+	assert.Equal(1, reserveCount)
+
+	t.Log("should have released")
+	assert.Equal(1, releaseCount)
+}
+
+func TestJunkData(t *testing.T) {
+	if !selinux.GetEnabled() {
+		t.Skip("selinux is not enabled")
+	}
+
+	assert := assert.New(t)
+	store := NewStore()
+	releaseCount := 0
+	store.Releaser = func(label string) {
+		releaseCount++
+	}
+	reserveCount := 0
+	store.Reserver = func(label string) {
+		reserveCount++
+	}
+
+	t.Log("should ignore empty label")
+	assert.NoError(store.Reserve(""))
+	assert.Equal(0, len(store.levels))
+	store.Release("")
+	assert.Equal(0, len(store.levels))
+	assert.Equal(0, releaseCount)
+	assert.Equal(0, reserveCount)
+
+	t.Log("should fail on bad label")
+	assert.Error(store.Reserve("junkjunkjunkc1c2"))
+	assert.Equal(0, len(store.levels))
+	store.Release("junkjunkjunkc1c2")
+	assert.Equal(0, len(store.levels))
+	assert.Equal(0, releaseCount)
+	assert.Equal(0, reserveCount)
+
+	t.Log("should not release unknown label")
+	store.Release("junk2:junk2:junk2:c1,c2")
+	assert.Equal(0, len(store.levels))
+	assert.Equal(0, releaseCount)
+	assert.Equal(0, reserveCount)
+
+	t.Log("should release once even if too many deletes")
+	assert.NoError(store.Reserve("junk2:junk2:junk2:c1,c2"))
+	assert.Equal(1, len(store.levels))
+	assert.Equal(1, store.levels["c1,c2"])
+	store.Release("junk2:junk2:junk2:c1,c2")
+	store.Release("junk2:junk2:junk2:c1,c2")
+	assert.Equal(0, len(store.levels))
+	assert.Equal(1, releaseCount)
+	assert.Equal(1, reserveCount)
+}

internal/cri/store/sandbox/metadata.go

@@ -0,0 +1,88 @@
+/*
+   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 sandbox
+
+import (
+	"encoding/json"
+	"fmt"
+
+	cni "github.com/containerd/go-cni"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+// NOTE(random-liu):
+// 1) Metadata is immutable after created.
+// 2) Metadata is checkpointed as containerd container label.
+
+// metadataVersion is current version of sandbox metadata.
+const metadataVersion = "v1"
+
+// versionedMetadata is the internal versioned sandbox metadata.
+type versionedMetadata struct {
+	// Version indicates the version of the versioned sandbox metadata.
+	Version string
+	// Metadata's type is metadataInternal. If not there will be a recursive call in MarshalJSON.
+	Metadata metadataInternal
+}
+
+// metadataInternal is for internal use.
+type metadataInternal Metadata
+
+// Metadata is the unversioned sandbox metadata.
+type Metadata struct {
+	// ID is the sandbox id.
+	ID string
+	// Name is the sandbox name.
+	Name string
+	// Config is the CRI sandbox config.
+	Config *runtime.PodSandboxConfig
+	// NetNSPath is the network namespace used by the sandbox.
+	NetNSPath string
+	// IP of Pod if it is attached to non host network
+	IP string
+	// AdditionalIPs of the Pod if it is attached to non host network
+	AdditionalIPs []string
+	// RuntimeHandler is the runtime handler name of the pod.
+	RuntimeHandler string
+	// CNIresult resulting configuration for attached network namespace interfaces
+	CNIResult *cni.Result
+	// ProcessLabel is the SELinux process label for the container
+	ProcessLabel string
+}
+
+// MarshalJSON encodes Metadata into bytes in json format.
+func (c *Metadata) MarshalJSON() ([]byte, error) {
+	return json.Marshal(&versionedMetadata{
+		Version:  metadataVersion,
+		Metadata: metadataInternal(*c),
+	})
+}
+
+// UnmarshalJSON decodes Metadata from bytes.
+func (c *Metadata) UnmarshalJSON(data []byte) error {
+	versioned := &versionedMetadata{}
+	if err := json.Unmarshal(data, versioned); err != nil {
+		return err
+	}
+	// Handle old version after upgrade.
+	switch versioned.Version {
+	case metadataVersion:
+		*c = Metadata(versioned.Metadata)
+		return nil
+	}
+	return fmt.Errorf("unsupported version: %q", versioned.Version)
+}

internal/cri/store/sandbox/metadata_test.go

@@ -0,0 +1,79 @@
+/*
+   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 sandbox
+
+import (
+	"encoding/json"
+	"testing"
+
+	assertlib "github.com/stretchr/testify/assert"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+func TestMetadataMarshalUnmarshal(t *testing.T) {
+	meta := &Metadata{
+		ID:   "test-id",
+		Name: "test-name",
+		Config: &runtime.PodSandboxConfig{
+			Metadata: &runtime.PodSandboxMetadata{
+				Name:      "test-name",
+				Uid:       "test-uid",
+				Namespace: "test-namespace",
+				Attempt:   1,
+			},
+		},
+	}
+	assert := assertlib.New(t)
+	newMeta := &Metadata{}
+	newVerMeta := &versionedMetadata{}
+
+	t.Logf("should be able to do json.marshal")
+	data, err := json.Marshal(meta)
+	assert.NoError(err)
+	data1, err := json.Marshal(&versionedMetadata{
+		Version:  metadataVersion,
+		Metadata: metadataInternal(*meta),
+	})
+	assert.NoError(err)
+	assert.Equal(data, data1)
+
+	t.Logf("should be able to do MarshalJSON")
+	data, err = meta.MarshalJSON()
+	assert.NoError(err)
+	assert.NoError(newMeta.UnmarshalJSON(data))
+	assert.Equal(meta, newMeta)
+
+	t.Logf("should be able to do MarshalJSON and json.Unmarshal")
+	data, err = meta.MarshalJSON()
+	assert.NoError(err)
+	assert.NoError(json.Unmarshal(data, newVerMeta))
+	assert.Equal(meta, (*Metadata)(&newVerMeta.Metadata))
+
+	t.Logf("should be able to do json.Marshal and UnmarshalJSON")
+	data, err = json.Marshal(meta)
+	assert.NoError(err)
+	assert.NoError(newMeta.UnmarshalJSON(data))
+	assert.Equal(meta, newMeta)
+
+	t.Logf("should json.Unmarshal fail for unsupported version")
+	unsupported, err := json.Marshal(&versionedMetadata{
+		Version:  "random-test-version",
+		Metadata: metadataInternal(*meta),
+	})
+	assert.NoError(err)
+	assert.Error(json.Unmarshal(unsupported, &newMeta))
+}

internal/cri/store/sandbox/sandbox.go

@@ -0,0 +1,165 @@
+/*
+   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 sandbox
+
+import (
+	"sync"
+
+	containerd "github.com/containerd/containerd/v2/client"
+	"github.com/containerd/containerd/v2/internal/cri/store"
+	"github.com/containerd/containerd/v2/internal/cri/store/label"
+	"github.com/containerd/containerd/v2/internal/cri/store/stats"
+	"github.com/containerd/containerd/v2/internal/truncindex"
+	"github.com/containerd/containerd/v2/pkg/netns"
+	"github.com/containerd/errdefs"
+)
+
+// Sandbox contains all resources associated with the sandbox. All methods to
+// mutate the internal state are thread safe.
+type Sandbox struct {
+	// Metadata is the metadata of the sandbox, it is immutable after created.
+	Metadata
+	// Status stores the status of the sandbox.
+	Status StatusStorage
+	// Container is the containerd sandbox container client.
+	Container containerd.Container
+	// CNI network namespace client.
+	// For hostnetwork pod, this is always nil;
+	// For non hostnetwork pod, this should never be nil.
+	NetNS *netns.NetNS
+	// StopCh is used to propagate the stop information of the sandbox.
+	*store.StopCh
+	// Stats contains (mutable) stats for the (pause) sandbox container
+	Stats *stats.ContainerStats
+}
+
+// NewSandbox creates an internally used sandbox type. This functions reminds
+// the caller that a sandbox must have a status.
+func NewSandbox(metadata Metadata, status Status) Sandbox {
+	s := Sandbox{
+		Metadata: metadata,
+		Status:   StoreStatus(status),
+		StopCh:   store.NewStopCh(),
+	}
+	if status.State == StateNotReady {
+		s.Stop()
+	}
+	return s
+}
+
+// Store stores all sandboxes.
+type Store struct {
+	lock      sync.RWMutex
+	sandboxes map[string]Sandbox
+	idIndex   *truncindex.TruncIndex
+	labels    *label.Store
+}
+
+// NewStore creates a sandbox store.
+func NewStore(labels *label.Store) *Store {
+	return &Store{
+		sandboxes: make(map[string]Sandbox),
+		idIndex:   truncindex.NewTruncIndex([]string{}),
+		labels:    labels,
+	}
+}
+
+// Add a sandbox into the store. Returns errdefs.ErrAlreadyExists if the sandbox is
+// already stored.
+func (s *Store) Add(sb Sandbox) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	if _, ok := s.sandboxes[sb.ID]; ok {
+		return errdefs.ErrAlreadyExists
+	}
+	if err := s.labels.Reserve(sb.ProcessLabel); err != nil {
+		return err
+	}
+	if err := s.idIndex.Add(sb.ID); err != nil {
+		return err
+	}
+	s.sandboxes[sb.ID] = sb
+	return nil
+}
+
+// Get returns the sandbox with specified id.
+// Returns errdefs.ErrNotFound if the sandbox doesn't exist.
+func (s *Store) Get(id string) (Sandbox, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = errdefs.ErrNotFound
+		}
+		return Sandbox{}, err
+	}
+	if sb, ok := s.sandboxes[id]; ok {
+		return sb, nil
+	}
+	return Sandbox{}, errdefs.ErrNotFound
+}
+
+// List lists all sandboxes.
+func (s *Store) List() []Sandbox {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	var sandboxes []Sandbox
+	for _, sb := range s.sandboxes {
+		sandboxes = append(sandboxes, sb)
+	}
+	return sandboxes
+}
+
+// UpdateContainerStats updates the sandbox specified by ID with the
+// stats present in 'newContainerStats'. Returns errdefs.ErrNotFound
+// if the sandbox does not exist in the store.
+func (s *Store) UpdateContainerStats(id string, newContainerStats *stats.ContainerStats) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = errdefs.ErrNotFound
+		}
+		return err
+	}
+
+	if _, ok := s.sandboxes[id]; !ok {
+		return errdefs.ErrNotFound
+	}
+
+	c := s.sandboxes[id]
+	c.Stats = newContainerStats
+	s.sandboxes[id] = c
+	return nil
+}
+
+// Delete deletes the sandbox with specified id.
+func (s *Store) Delete(id string) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(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
+	}
+	s.labels.Release(s.sandboxes[id].ProcessLabel)
+	s.idIndex.Delete(id)
+	delete(s.sandboxes, id)
+}

internal/cri/store/sandbox/sandbox_test.go

@@ -0,0 +1,189 @@
+/*
+   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 sandbox
+
+import (
+	"testing"
+	"time"
+
+	"github.com/containerd/containerd/v2/internal/cri/store/label"
+	"github.com/containerd/containerd/v2/internal/cri/store/stats"
+	"github.com/containerd/errdefs"
+
+	assertlib "github.com/stretchr/testify/assert"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+func TestSandboxStore(t *testing.T) {
+	sandboxes := map[string]Sandbox{
+		"1": NewSandbox(
+			Metadata{
+				ID:   "1",
+				Name: "Sandbox-1",
+				Config: &runtime.PodSandboxConfig{
+					Metadata: &runtime.PodSandboxMetadata{
+						Name:      "TestPod-1",
+						Uid:       "TestUid-1",
+						Namespace: "TestNamespace-1",
+						Attempt:   1,
+					},
+				},
+				NetNSPath: "TestNetNS-1",
+			},
+			Status{State: StateReady},
+		),
+		"2abcd": NewSandbox(
+			Metadata{
+				ID:   "2abcd",
+				Name: "Sandbox-2abcd",
+				Config: &runtime.PodSandboxConfig{
+					Metadata: &runtime.PodSandboxMetadata{
+						Name:      "TestPod-2abcd",
+						Uid:       "TestUid-2abcd",
+						Namespace: "TestNamespace-2abcd",
+						Attempt:   2,
+					},
+				},
+				NetNSPath: "TestNetNS-2",
+			},
+			Status{State: StateNotReady},
+		),
+		"4a333": NewSandbox(
+			Metadata{
+				ID:   "4a333",
+				Name: "Sandbox-4a333",
+				Config: &runtime.PodSandboxConfig{
+					Metadata: &runtime.PodSandboxMetadata{
+						Name:      "TestPod-4a333",
+						Uid:       "TestUid-4a333",
+						Namespace: "TestNamespace-4a333",
+						Attempt:   3,
+					},
+				},
+				NetNSPath: "TestNetNS-3",
+			},
+			Status{State: StateNotReady},
+		),
+		"4abcd": NewSandbox(
+			Metadata{
+				ID:   "4abcd",
+				Name: "Sandbox-4abcd",
+				Config: &runtime.PodSandboxConfig{
+					Metadata: &runtime.PodSandboxMetadata{
+						Name:      "TestPod-4abcd",
+						Uid:       "TestUid-4abcd",
+						Namespace: "TestNamespace-4abcd",
+						Attempt:   1,
+					},
+				},
+				NetNSPath: "TestNetNS-4abcd",
+			},
+			Status{State: StateReady},
+		),
+	}
+	unknown := NewSandbox(
+		Metadata{
+			ID:   "3defg",
+			Name: "Sandbox-3defg",
+			Config: &runtime.PodSandboxConfig{
+				Metadata: &runtime.PodSandboxMetadata{
+					Name:      "TestPod-3defg",
+					Uid:       "TestUid-3defg",
+					Namespace: "TestNamespace-3defg",
+					Attempt:   1,
+				},
+			},
+			NetNSPath: "TestNetNS-3defg",
+		},
+		Status{State: StateUnknown},
+	)
+	stats := map[string]*stats.ContainerStats{
+		"1": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 1,
+		},
+		"2abcd": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 2,
+		},
+		"4a333": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 3,
+		},
+		"4abcd": {
+			Timestamp:            time.Now(),
+			UsageCoreNanoSeconds: 4,
+		},
+	}
+	assert := assertlib.New(t)
+	s := NewStore(label.NewStore())
+
+	t.Logf("should be able to add sandbox")
+	for _, sb := range sandboxes {
+		assert.NoError(s.Add(sb))
+	}
+	assert.NoError(s.Add(unknown))
+
+	t.Logf("should be able to get sandbox")
+	genTruncIndex := func(normalName string) string { return normalName[:(len(normalName)+1)/2] }
+	for id, sb := range sandboxes {
+		got, err := s.Get(genTruncIndex(id))
+		assert.NoError(err)
+		assert.Equal(sb, got)
+	}
+
+	t.Logf("should be able to get sandbox in unknown state with Get")
+	got, err := s.Get(unknown.ID)
+	assert.NoError(err)
+	assert.Equal(unknown, got)
+
+	t.Logf("should be able to list sandboxes")
+	sbNum := len(sandboxes) + 1
+	sbs := s.List()
+	assert.Len(sbs, sbNum)
+
+	t.Logf("should be able to update stats on container")
+	for id := range sandboxes {
+		err := s.UpdateContainerStats(id, stats[id])
+		assert.NoError(err)
+	}
+
+	// Validate stats were updated
+	sbs = s.List()
+	assert.Len(sbs, sbNum)
+	for _, sb := range sbs {
+		assert.Equal(stats[sb.ID], sb.Stats)
+	}
+
+	for testID, v := range sandboxes {
+		truncID := genTruncIndex(testID)
+
+		t.Logf("add should return already exists error for duplicated sandbox")
+		assert.Equal(errdefs.ErrAlreadyExists, s.Add(v))
+
+		t.Logf("should be able to delete sandbox")
+		s.Delete(truncID)
+		sbNum--
+		sbs = s.List()
+		assert.Len(sbs, sbNum)
+
+		t.Logf("get should return not exist error after deletion")
+		sb, err := s.Get(truncID)
+		assert.Equal(Sandbox{}, sb)
+		assert.Equal(errdefs.ErrNotFound, err)
+	}
+}

internal/cri/store/sandbox/status.go

@@ -0,0 +1,155 @@
+/*
+   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 sandbox
+
+import (
+	"strconv"
+	"sync"
+	"time"
+
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
+)
+
+// The sandbox state machine in the CRI plugin:
+//                    +              +
+//                    |              |
+//                    | Create(Run)  | Load
+//                    |              |
+//                    |              |
+//                    |              |    Start
+//                    |              |(failed and not cleaned)
+//      Start         |--------------|--------------+
+//(failed but cleaned)|              |              |
+// +------------------+              |-----------+  |
+// |                  | Start(Run)   |           |  |
+// |                  |              |           |  |
+// | PortForward +----v----+         |           |  |
+// |      +------+         |         |           |  |
+// |      |      |  READY  <---------+           |  |
+// |      +------>         |         |           |  |
+// |             +----+----+         |           |  |
+// |                  |              |           |  |
+// |                  | Stop/Exit    |           |  |
+// |                  |              |           |  |
+// |             +----v----+         |           |  |
+// |             |         <---------+      +----v--v-+
+// |             | NOTREADY|                |         |
+// |             |         <----------------+ UNKNOWN |
+// |             +----+----+       Stop     |         |
+// |                  |                     +---------+
+// |                  | Remove
+// |                  v
+// +-------------> DELETED
+
+// State is the sandbox state we use in containerd/cri.
+// It includes unknown, which is internal states not defined in CRI.
+// The state mapping from internal states to CRI states:
+// * ready -> ready
+// * not ready -> not ready
+// * unknown -> not ready
+type State uint32
+
+const (
+	// StateReady is ready state, it means sandbox container
+	// is running.
+	StateReady State = iota
+	// StateNotReady is notready state, it ONLY means sandbox
+	// container is not running.
+	// StopPodSandbox should still be called for NOTREADY sandbox to
+	// cleanup resources other than sandbox container, e.g. network namespace.
+	// This is an assumption made in CRI.
+	StateNotReady
+	// StateUnknown is unknown state. Sandbox only goes
+	// into unknown state when its status fails to be loaded.
+	StateUnknown
+)
+
+// String returns the string representation of the state
+func (s State) String() string {
+	switch s {
+	case StateReady:
+		return runtime.PodSandboxState_SANDBOX_READY.String()
+	case StateNotReady:
+		return runtime.PodSandboxState_SANDBOX_NOTREADY.String()
+	case StateUnknown:
+		// PodSandboxState doesn't have an unknown state, but State does, so return a string using the same convention
+		return "SANDBOX_UNKNOWN"
+	default:
+		return "invalid sandbox state value: " + strconv.Itoa(int(s))
+	}
+}
+
+// Status is the status of a sandbox.
+type Status struct {
+	// Pid is the init process id of the sandbox container.
+	Pid uint32
+	// CreatedAt is the created timestamp.
+	CreatedAt time.Time
+	// ExitedAt is the stop timestamp
+	ExitedAt time.Time
+	// ExitStatus is the stop sandbox status
+	ExitStatus uint32
+	// State is the state of the sandbox.
+	State State
+}
+
+// UpdateFunc is function used to update the sandbox status. If there
+// is an error, the update will be rolled back.
+type UpdateFunc func(Status) (Status, error)
+
+// StatusStorage manages the sandbox status.
+// The status storage for sandbox is different from container status storage,
+// because we don't checkpoint sandbox status. If we need checkpoint in the
+// future, we should combine this with container status storage.
+type StatusStorage interface {
+	// Get a sandbox status.
+	Get() Status
+	// Update the sandbox status. Note that the update MUST be applied
+	// in one transaction.
+	Update(UpdateFunc) error
+}
+
+// StoreStatus creates the storage containing the passed in sandbox status with the
+// specified id.
+// The status MUST be created in one transaction.
+func StoreStatus(status Status) StatusStorage {
+	return &statusStorage{status: status}
+}
+
+type statusStorage struct {
+	sync.RWMutex
+	status Status
+}
+
+// Get a copy of sandbox status.
+func (s *statusStorage) Get() Status {
+	s.RLock()
+	defer s.RUnlock()
+	return s.status
+}
+
+// Update the sandbox status.
+func (s *statusStorage) Update(u UpdateFunc) error {
+	s.Lock()
+	defer s.Unlock()
+	newStatus, err := u(s.status)
+	if err != nil {
+		return err
+	}
+	s.status = newStatus
+	return nil
+}

internal/cri/store/sandbox/status_test.go

@@ -0,0 +1,67 @@
+/*
+   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 sandbox
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	assertlib "github.com/stretchr/testify/assert"
+)
+
+func TestStatus(t *testing.T) {
+	testStatus := Status{
+		Pid:       123,
+		CreatedAt: time.Now(),
+		State:     StateUnknown,
+	}
+	updateStatus := Status{
+		Pid:       456,
+		CreatedAt: time.Now(),
+		State:     StateReady,
+	}
+	updateErr := errors.New("update error")
+	assert := assertlib.New(t)
+
+	t.Logf("simple store and get")
+	s := StoreStatus(testStatus)
+	old := s.Get()
+	assert.Equal(testStatus, old)
+
+	t.Logf("failed update should not take effect")
+	err := s.Update(func(o Status) (Status, error) {
+		return updateStatus, updateErr
+	})
+	assert.Equal(updateErr, err)
+	assert.Equal(testStatus, s.Get())
+
+	t.Logf("successful update should take effect but not checkpoint")
+	err = s.Update(func(o Status) (Status, error) {
+		return updateStatus, nil
+	})
+	assert.NoError(err)
+	assert.Equal(updateStatus, s.Get())
+}
+
+func TestStateStringConversion(t *testing.T) {
+	assert := assertlib.New(t)
+	assert.Equal("SANDBOX_READY", StateReady.String())
+	assert.Equal("SANDBOX_NOTREADY", StateNotReady.String())
+	assert.Equal("SANDBOX_UNKNOWN", StateUnknown.String())
+	assert.Equal("invalid sandbox state value: 123", State(123).String())
+}

internal/cri/store/snapshot/snapshot.go

@@ -0,0 +1,93 @@
+/*
+   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 snapshot
+
+import (
+	"sync"
+
+	snapshot "github.com/containerd/containerd/v2/core/snapshots"
+	"github.com/containerd/errdefs"
+)
+
+type Key struct {
+	// Key is the key of the snapshot
+	Key string
+	// Snapshotter is the name of the snapshotter managing the snapshot
+	Snapshotter string
+}
+
+// Snapshot contains the information about the snapshot.
+type Snapshot struct {
+	// Key is the key of the snapshot
+	Key Key
+	// Kind is the kind of the snapshot (active, committed, view)
+	Kind snapshot.Kind
+	// Size is the size of the snapshot in bytes.
+	Size uint64
+	// Inodes is the number of inodes used by the snapshot
+	Inodes uint64
+	// Timestamp is latest update time (in nanoseconds) of the snapshot
+	// information.
+	Timestamp int64
+}
+
+// Store stores all snapshots.
+type Store struct {
+	lock      sync.RWMutex
+	snapshots map[Key]Snapshot
+}
+
+// NewStore creates a snapshot store.
+func NewStore() *Store {
+	return &Store{snapshots: make(map[Key]Snapshot)}
+}
+
+// Add a snapshot into the store.
+func (s *Store) Add(snapshot Snapshot) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.snapshots[snapshot.Key] = snapshot
+}
+
+// Get returns the snapshot with specified key. Returns errdefs.ErrNotFound if the
+// snapshot doesn't exist.
+func (s *Store) Get(key Key) (Snapshot, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	if sn, ok := s.snapshots[key]; ok {
+		return sn, nil
+	}
+	return Snapshot{}, errdefs.ErrNotFound
+}
+
+// List lists all snapshots.
+func (s *Store) List() []Snapshot {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	var snapshots []Snapshot
+	for _, sn := range s.snapshots {
+		snapshots = append(snapshots, sn)
+	}
+	return snapshots
+}
+
+// Delete deletes the snapshot with specified key.
+func (s *Store) Delete(key Key) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	delete(s.snapshots, key)
+}

internal/cri/store/snapshot/snapshot_test.go

@@ -0,0 +1,108 @@
+/*
+   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 snapshot
+
+import (
+	"testing"
+	"time"
+
+	snapshot "github.com/containerd/containerd/v2/core/snapshots"
+	"github.com/containerd/errdefs"
+
+	assertlib "github.com/stretchr/testify/assert"
+)
+
+func TestSnapshotStore(t *testing.T) {
+	key1 := Key{
+		Key:         "key1",
+		Snapshotter: "snapshotter1",
+	}
+	key2 := Key{
+		Key:         "key2",
+		Snapshotter: "snapshotter1",
+	}
+	key3 := Key{
+		Key:         "key1",
+		Snapshotter: "snapshotter2",
+	}
+	snapshots := map[Key]Snapshot{
+		key1: {
+			Key:       key1,
+			Kind:      snapshot.KindActive,
+			Size:      10,
+			Inodes:    100,
+			Timestamp: time.Now().UnixNano(),
+		},
+		key2: {
+			Key:       key2,
+			Kind:      snapshot.KindCommitted,
+			Size:      20,
+			Inodes:    200,
+			Timestamp: time.Now().UnixNano(),
+		},
+		key3: {
+			Key:       key3,
+			Kind:      snapshot.KindView,
+			Size:      0,
+			Inodes:    0,
+			Timestamp: time.Now().UnixNano(),
+		},
+	}
+	assert := assertlib.New(t)
+
+	s := NewStore()
+
+	t.Logf("should be able to add snapshot")
+	for _, sn := range snapshots {
+		s.Add(sn)
+	}
+
+	t.Logf("should be able to get snapshot")
+	for id, sn := range snapshots {
+		got, err := s.Get(id)
+		assert.NoError(err)
+		assert.Equal(sn, got)
+	}
+
+	t.Logf("should be able to list snapshot")
+	sns := s.List()
+	assert.Len(sns, 3)
+
+	invalidKey := Key{
+		Key:         "key2",
+		Snapshotter: "snapshotter2",
+	}
+	t.Logf("should not delete snapshot with invalid key")
+	s.Delete(invalidKey)
+	sns = s.List()
+	assert.Len(sns, 3)
+
+	testKey := Key{
+		Key:         "key2",
+		Snapshotter: "snapshotter1",
+	}
+
+	t.Logf("should be able to delete snapshot")
+	s.Delete(testKey)
+	sns = s.List()
+	assert.Len(sns, 2)
+
+	t.Logf("get should return empty struct and ErrNotExist after deletion")
+	sn, err := s.Get(testKey)
+	assert.Equal(Snapshot{}, sn)
+	assert.Equal(errdefs.ErrNotFound, err)
+}

internal/cri/store/stats/stats.go

@@ -0,0 +1,27 @@
+/*
+   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 stats
+
+import "time"
+
+// ContainerStats contains the information about container stats.
+type ContainerStats struct {
+	// Timestamp of when stats were collected
+	Timestamp time.Time
+	// Cumulative CPU usage (sum across all cores) since object creation.
+	UsageCoreNanoSeconds uint64
+}

internal/cri/store/util.go

@@ -0,0 +1,42 @@
+/*
+   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 store
+
+import "sync"
+
+// StopCh is used to propagate the stop information of a container.
+type StopCh struct {
+	ch   chan struct{}
+	once sync.Once
+}
+
+// NewStopCh creates a stop channel. The channel is open by default.
+func NewStopCh() *StopCh {
+	return &StopCh{ch: make(chan struct{})}
+}
+
+// Stop close stopCh of the container.
+func (s *StopCh) Stop() {
+	s.once.Do(func() {
+		close(s.ch)
+	})
+}
+
+// Stopped return the stopCh of the container as a readonly channel.
+func (s *StopCh) Stopped() <-chan struct{} {
+	return s.ch
+}