Imagepull per runtime class alpha release changes

This commit does the following:
1. Add RuntimeClassInImageCriApi feature gate
2. Extend pkg/kubelet/container Image struct
3. Adds runtimeHandler string in the following CRI calls
   i.   ImageStatus
   ii.  PullImageRequest
   iii.  RemoveImage

Signed-off-by: kiashok <kiashok@microsoft.com>

pkg/features/kube_features.go

@@ -710,6 +710,13 @@ const (
 	// certificate as expiration approaches.
 	RotateKubeletServerCertificate featuregate.Feature = "RotateKubeletServerCertificate"
 
+	// owner: @kiashok
+	// kep: https://kep.k8s.io/4216
+	// alpha: v1.29
+	//
+	// Adds support to pull images based on the runtime class specified.
+	RuntimeClassInImageCriAPI featuregate.Feature = "RuntimeClassInImageCriApi"
+
 	// owner: @danielvegamyhre
 	// kep: https://kep.k8s.io/2413
 	// beta: v1.27
@@ -1139,6 +1146,8 @@ var defaultKubernetesFeatureGates = map[featuregate.Feature]featuregate.FeatureS
 
 	RotateKubeletServerCertificate: {Default: true, PreRelease: featuregate.Beta},
 
+	RuntimeClassInImageCriAPI: {Default: false, PreRelease: featuregate.Alpha},
+
 	ElasticIndexedJob: {Default: true, PreRelease: featuregate.Beta},
 
 	SchedulerQueueingHints: {Default: true, PreRelease: featuregate.Beta},

pkg/kubelet/container/helpers.go

@@ -273,6 +273,7 @@ func ConvertPodStatusToRunningPod(runtimeName string, podStatus *PodStatus) Pod
 			Name:                 containerStatus.Name,
 			Image:                containerStatus.Image,
 			ImageID:              containerStatus.ImageID,
+			ImageRuntimeHandler:  containerStatus.ImageRuntimeHandler,
 			Hash:                 containerStatus.Hash,
 			HashWithoutResources: containerStatus.HashWithoutResources,
 			State:                containerStatus.State,

pkg/kubelet/container/runtime.go

@@ -52,6 +52,8 @@ type Version interface {
 type ImageSpec struct {
 	// ID of the image.
 	Image string
+	// Runtime handler used to pull this image
+	RuntimeHandler string
 	// The annotations for the image.
 	// This should be passed to CRI during image pulls and returned when images are listed.
 	Annotations []Annotation
@@ -282,6 +284,8 @@ type Container struct {
 	Image string
 	// The id of the image used by the container.
 	ImageID string
+	// Runtime handler used to pull the image if any.
+	ImageRuntimeHandler string
 	// Hash of the container, used for comparison. Optional for containers
 	// not managed by kubelet.
 	Hash uint64
@@ -347,6 +351,8 @@ type Status struct {
 	Image string
 	// ID of the image.
 	ImageID string
+	// Runtime handler used to pull the image if any.
+	ImageRuntimeHandler string
 	// Hash of the container, used for comparison.
 	Hash uint64
 	// Hash of the container over fields with Resources field zero'd out.

pkg/kubelet/images/image_gc_manager.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
@@ -32,8 +33,10 @@ import (
 	"k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apimachinery/pkg/util/wait"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/tools/record"
 	statsapi "k8s.io/kubelet/pkg/apis/stats/v1alpha1"
+	"k8s.io/kubernetes/pkg/features"
 	"k8s.io/kubernetes/pkg/kubelet/container"
 	"k8s.io/kubernetes/pkg/kubelet/events"
 	"k8s.io/kubernetes/pkg/kubelet/metrics"
@@ -43,6 +46,10 @@ import (
 // instrumentationScope is OpenTelemetry instrumentation scope name
 const instrumentationScope = "k8s.io/kubernetes/pkg/kubelet/images"
 
+// When RuntimeClassInImageCriAPI feature gate is enabled, imageRecord is
+// indexed as imageId-RuntimeHandler
+const imageIndexTupleFormat = "%s,%s"
+
 // StatsProvider is an interface for fetching stats used during image garbage
 // collection.
 type StatsProvider interface {
@@ -90,7 +97,12 @@ type realImageGCManager struct {
 	// Container runtime
 	runtime container.Runtime
 
-	// Records of images and their use.
+	// Records of images and their use. Indexed by ImageId.
+	// If RuntimeClassInImageCriAPI feature gate is enabled, imageRecords
+	// are identified by a tuple of (imageId,runtimeHandler) that is passed
+	// from ListImages() call. If no runtimehandler is specified in response
+	// to ListImages() by the container runtime, only imageID will be used as
+	// the index of this map.
 	imageRecords     map[string]*imageRecord
 	imageRecordsLock sync.Mutex
 
@@ -149,6 +161,8 @@ func (i *imageCache) get() []container.Image {
 
 // Information about the images we track.
 type imageRecord struct {
+	// runtime handler used to pull this image
+	runtimeHandlerUsedToPullImage string
 	// Time when this image was first detected.
 	firstDetected time.Time
 
@@ -223,6 +237,7 @@ func (im *realImageGCManager) GetImageList() ([]container.Image, error) {
 }
 
 func (im *realImageGCManager) detectImages(ctx context.Context, detectTime time.Time) (sets.String, error) {
+	isRuntimeClassInImageCriAPIEnabled := utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI)
 	imagesInUse := sets.NewString()
 
 	images, err := im.runtime.ListImages(ctx)
@@ -237,8 +252,14 @@ func (im *realImageGCManager) detectImages(ctx context.Context, detectTime time.
 	// Make a set of images in use by containers.
 	for _, pod := range pods {
 		for _, container := range pod.Containers {
-			klog.V(5).InfoS("Container uses image", "pod", klog.KRef(pod.Namespace, pod.Name), "containerName", container.Name, "containerImage", container.Image, "imageID", container.ImageID)
+			if !isRuntimeClassInImageCriAPIEnabled {
-			imagesInUse.Insert(container.ImageID)
+				klog.V(5).InfoS("Container uses image", "pod", klog.KRef(pod.Namespace, pod.Name), "containerName", container.Name, "containerImage", container.Image, "imageID", container.ImageID)
+				imagesInUse.Insert(container.ImageID)
+			} else {
+				imageKey := getImageTuple(container.ImageID, container.ImageRuntimeHandler)
+				klog.V(5).InfoS("Container uses image", "pod", klog.KRef(pod.Namespace, pod.Name), "containerName", container.Name, "containerImage", container.Image, "imageID", container.ImageID, "imageKey", imageKey)
+				imagesInUse.Insert(imageKey)
+			}
 		}
 	}
 
@@ -248,28 +269,36 @@ func (im *realImageGCManager) detectImages(ctx context.Context, detectTime time.
 	im.imageRecordsLock.Lock()
 	defer im.imageRecordsLock.Unlock()
 	for _, image := range images {
-		klog.V(5).InfoS("Adding image ID to currentImages", "imageID", image.ID)
+		imageKey := image.ID
-		currentImages.Insert(image.ID)
+		if !isRuntimeClassInImageCriAPIEnabled {
+			klog.V(5).InfoS("Adding image ID to currentImages", "imageID", imageKey)
+		} else {
+			imageKey = getImageTuple(image.ID, image.Spec.RuntimeHandler)
+			klog.V(5).InfoS("Adding image ID with runtime class to currentImages", "imageKey", imageKey, "runtimeHandler", image.Spec.RuntimeHandler)
+		}
+
+		currentImages.Insert(imageKey)
 
 		// New image, set it as detected now.
-		if _, ok := im.imageRecords[image.ID]; !ok {
+		if _, ok := im.imageRecords[imageKey]; !ok {
-			klog.V(5).InfoS("Image ID is new", "imageID", image.ID)
+			klog.V(5).InfoS("Image ID is new", "imageID", imageKey, "runtimeHandler", image.Spec.RuntimeHandler)
-			im.imageRecords[image.ID] = &imageRecord{
+			im.imageRecords[imageKey] = &imageRecord{
-				firstDetected: detectTime,
+				firstDetected:                 detectTime,
+				runtimeHandlerUsedToPullImage: image.Spec.RuntimeHandler,
 			}
 		}
 
 		// Set last used time to now if the image is being used.
-		if isImageUsed(image.ID, imagesInUse) {
+		if isImageUsed(imageKey, imagesInUse) {
-			klog.V(5).InfoS("Setting Image ID lastUsed", "imageID", image.ID, "lastUsed", now)
+			klog.V(5).InfoS("Setting Image ID lastUsed", "imageID", imageKey, "lastUsed", now)
-			im.imageRecords[image.ID].lastUsed = now
+			im.imageRecords[imageKey].lastUsed = now
 		}
 
-		klog.V(5).InfoS("Image ID has size", "imageID", image.ID, "size", image.Size)
+		klog.V(5).InfoS("Image ID has size", "imageID", imageKey, "size", image.Size)
-		im.imageRecords[image.ID].size = image.Size
+		im.imageRecords[imageKey].size = image.Size
 
-		klog.V(5).InfoS("Image ID is pinned", "imageID", image.ID, "pinned", image.Pinned)
+		klog.V(5).InfoS("Image ID is pinned", "imageID", imageKey, "pinned", image.Pinned)
-		im.imageRecords[image.ID].pinned = image.Pinned
+		im.imageRecords[imageKey].pinned = image.Pinned
 	}
 
 	// Remove old images from our records.
@@ -391,7 +420,7 @@ func (im *realImageGCManager) freeSpace(ctx context.Context, bytesToFree int64,
 	var deletionErrors []error
 	spaceFreed := int64(0)
 	for _, image := range images {
-		klog.V(5).InfoS("Evaluating image ID for possible garbage collection based on disk usage", "imageID", image.id)
+		klog.V(5).InfoS("Evaluating image ID for possible garbage collection based on disk usage", "imageID", image.id, "runtimeHandler", image.imageRecord.runtimeHandlerUsedToPullImage)
 		// Images that are currently in used were given a newer lastUsed.
 		if image.lastUsed.Equal(freeTime) || image.lastUsed.After(freeTime) {
 			klog.V(5).InfoS("Image ID was used too recently, not eligible for garbage collection", "imageID", image.id, "lastUsed", image.lastUsed, "freeTime", freeTime)
@@ -423,19 +452,28 @@ func (im *realImageGCManager) freeSpace(ctx context.Context, bytesToFree int64,
 }
 
 func (im *realImageGCManager) freeImage(ctx context.Context, image evictionInfo) error {
+	isRuntimeClassInImageCriAPIEnabled := utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI)
 	// Remove image. Continue despite errors.
-	klog.InfoS("Removing image to free bytes", "imageID", image.id, "size", image.size)
+	var err error
-	err := im.runtime.RemoveImage(ctx, container.ImageSpec{Image: image.id})
+	klog.InfoS("Removing image to free bytes", "imageID", image.id, "size", image.size, "runtimeHandler", image.runtimeHandlerUsedToPullImage)
+	err = im.runtime.RemoveImage(ctx, container.ImageSpec{Image: image.id, RuntimeHandler: image.runtimeHandlerUsedToPullImage})
 	if err != nil {
 		return err
 	}
-	delete(im.imageRecords, image.id)
+
+	imageKey := image.id
+	if isRuntimeClassInImageCriAPIEnabled {
+		imageKey = getImageTuple(image.id, image.runtimeHandlerUsedToPullImage)
+	}
+	delete(im.imageRecords, imageKey)
+
 	metrics.ImageGarbageCollectedTotal.Inc()
 	return err
 }
 
 // Queries all of the image records and arranges them in a slice of evictionInfo, sorted based on last time used, ignoring images pinned by the runtime.
 func (im *realImageGCManager) imagesInEvictionOrder(ctx context.Context, freeTime time.Time) ([]evictionInfo, error) {
+	isRuntimeClassInImageCriAPIEnabled := utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI)
 	imagesInUse, err := im.detectImages(ctx, freeTime)
 	if err != nil {
 		return nil, err
@@ -457,15 +495,46 @@ func (im *realImageGCManager) imagesInEvictionOrder(ctx context.Context, freeTim
 			continue
 
 		}
-		images = append(images, evictionInfo{
+		if !isRuntimeClassInImageCriAPIEnabled {
-			id:          image,
+			images = append(images, evictionInfo{
-			imageRecord: *record,
+				id:          image,
-		})
+				imageRecord: *record,
+			})
+		} else {
+			imageID := getImageIDFromTuple(image)
+			// Ensure imageID is valid or else continue
+			if imageID == "" {
+				im.recorder.Eventf(im.nodeRef, v1.EventTypeWarning, "ImageID is not valid, skipping, ImageID: %v", imageID)
+				continue
+			}
+			images = append(images, evictionInfo{
+				id:          imageID,
+				imageRecord: *record,
+			})
+		}
 	}
 	sort.Sort(byLastUsedAndDetected(images))
 	return images, nil
 }
 
+// If RuntimeClassInImageCriAPI feature gate is enabled, imageRecords
+// are identified by a tuple of (imageId,runtimeHandler) that is passed
+// from ListImages() call. If no runtimehandler is specified in response
+// to ListImages() by the container runtime, only imageID will be will
+// be returned.
+func getImageTuple(imageID, runtimeHandler string) string {
+	if runtimeHandler == "" {
+		return imageID
+	}
+	return fmt.Sprintf(imageIndexTupleFormat, imageID, runtimeHandler)
+}
+
+// get imageID from the imageTuple
+func getImageIDFromTuple(image string) string {
+	imageTuples := strings.Split(image, ",")
+	return imageTuples[0]
+}
+
 type evictionInfo struct {
 	id string
 	imageRecord

pkg/kubelet/images/image_gc_manager_test.go

@@ -28,8 +28,11 @@ import (
 	"github.com/stretchr/testify/require"
 
 	oteltrace "go.opentelemetry.io/otel/trace"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/tools/record"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
 	statsapi "k8s.io/kubelet/pkg/apis/stats/v1alpha1"
+	"k8s.io/kubernetes/pkg/features"
 	"k8s.io/kubernetes/pkg/kubelet/container"
 	containertest "k8s.io/kubernetes/pkg/kubelet/container/testing"
 	stats "k8s.io/kubernetes/pkg/kubelet/server/stats"
@@ -66,6 +69,15 @@ func (im *realImageGCManager) getImageRecord(name string) (*imageRecord, bool) {
 	return &vCopy, ok
 }
 
+func (im *realImageGCManager) getImageRecordWithRuntimeHandlerInImageCriAPIFeatureGate(name, runtimeHandler string) (*imageRecord, bool) {
+	im.imageRecordsLock.Lock()
+	defer im.imageRecordsLock.Unlock()
+	imageKey := getImageTuple(name, runtimeHandler)
+	v, ok := im.imageRecords[imageKey]
+	vCopy := *v
+	return &vCopy, ok
+}
+
 // Returns the id of the image with the given ID.
 func imageID(id int) string {
 	return fmt.Sprintf("image-%d", id)
@@ -84,6 +96,24 @@ func makeImage(id int, size int64) container.Image {
 	}
 }
 
+// Make an image with the specified ID.
+func makeImageWithRuntimeHandler(id int, size int64, runtimeHandler string) container.Image {
+	if runtimeHandler == "" {
+		return container.Image{
+			ID:   imageID(id),
+			Size: size,
+		}
+	} else {
+		return container.Image{
+			ID:   imageID(id),
+			Size: size,
+			Spec: container.ImageSpec{
+				RuntimeHandler: runtimeHandler,
+			},
+		}
+	}
+}
+
 // Make a container with the specified ID. It will use the image with the same ID.
 func makeContainer(id int) *container.Container {
 	return &container.Container{
@@ -141,6 +171,64 @@ func TestDetectImagesInitialDetect(t *testing.T) {
 	assert.True(withContainer.lastUsed.After(startTime))
 }
 
+func TestDetectImagesInitialDetectWithRuntimeHandlerInImageCriAPIFeatureGate(t *testing.T) {
+	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClassInImageCriAPI, true)()
+	testRuntimeHandler := "test-runtimeHandler"
+	ctx := context.Background()
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mockStatsProvider := statstest.NewMockProvider(mockCtrl)
+
+	manager, fakeRuntime := newRealImageGCManager(ImageGCPolicy{}, mockStatsProvider)
+	fakeRuntime.ImageList = []container.Image{
+		makeImageWithRuntimeHandler(0, 1024, testRuntimeHandler),
+		makeImageWithRuntimeHandler(1, 2048, testRuntimeHandler),
+		makeImageWithRuntimeHandler(2, 2048, ""),
+	}
+	fakeRuntime.AllPodList = []*containertest.FakePod{
+		{Pod: &container.Pod{
+			Containers: []*container.Container{
+				{
+					ID:      container.ContainerID{Type: "test", ID: fmt.Sprintf("container-%d", 1)},
+					ImageID: imageID(1),
+					// The image field is not set to simulate a no-name image
+					ImageRuntimeHandler: testRuntimeHandler,
+				},
+				{
+					ID:      container.ContainerID{Type: "test", ID: fmt.Sprintf("container-%d", 2)},
+					Image:   imageName(2),
+					ImageID: imageID(2),
+					// The runtime handler field is not set to simulate the case when
+					// the feature gate "RuntimeHandlerInImageCriApi" is on and container runtime has not implemented
+					// KEP 4216, which means that runtimeHandler string is not set in the
+					// responses from the container runtime.
+				},
+			},
+		}},
+	}
+
+	startTime := time.Now().Add(-time.Millisecond)
+	_, err := manager.detectImages(ctx, zero)
+	assert := assert.New(t)
+	require.NoError(t, err)
+	assert.Equal(manager.imageRecordsLen(), 3)
+	noContainer, ok := manager.getImageRecordWithRuntimeHandlerInImageCriAPIFeatureGate(imageID(0), testRuntimeHandler)
+	require.True(t, ok)
+	assert.Equal(zero, noContainer.firstDetected)
+	assert.Equal(testRuntimeHandler, noContainer.runtimeHandlerUsedToPullImage)
+	assert.Equal(zero, noContainer.lastUsed)
+	withContainerUsingNoNameImage, ok := manager.getImageRecordWithRuntimeHandlerInImageCriAPIFeatureGate(imageID(1), testRuntimeHandler)
+	require.True(t, ok)
+	assert.Equal(zero, withContainerUsingNoNameImage.firstDetected)
+	assert.True(withContainerUsingNoNameImage.lastUsed.After(startTime))
+	assert.Equal(testRuntimeHandler, withContainerUsingNoNameImage.runtimeHandlerUsedToPullImage)
+	withContainer, ok := manager.getImageRecordWithRuntimeHandlerInImageCriAPIFeatureGate(imageID(2), "")
+	require.True(t, ok)
+	assert.Equal(zero, withContainer.firstDetected)
+	assert.True(withContainer.lastUsed.After(startTime))
+	assert.Equal("", withContainer.runtimeHandlerUsedToPullImage)
+}
+
 func TestDetectImagesWithNewImage(t *testing.T) {
 	ctx := context.Background()
 	mockCtrl := gomock.NewController(t)
@@ -182,14 +270,17 @@ func TestDetectImagesWithNewImage(t *testing.T) {
 	require.True(t, ok)
 	assert.Equal(zero, noContainer.firstDetected)
 	assert.Equal(zero, noContainer.lastUsed)
+	assert.Equal("", noContainer.runtimeHandlerUsedToPullImage)
 	withContainer, ok := manager.getImageRecord(imageID(1))
 	require.True(t, ok)
 	assert.Equal(zero, withContainer.firstDetected)
 	assert.True(withContainer.lastUsed.After(startTime))
+	assert.Equal("", noContainer.runtimeHandlerUsedToPullImage)
 	newContainer, ok := manager.getImageRecord(imageID(2))
 	require.True(t, ok)
 	assert.Equal(detectedTime, newContainer.firstDetected)
 	assert.Equal(zero, noContainer.lastUsed)
+	assert.Equal("", noContainer.runtimeHandlerUsedToPullImage)
 }
 
 func TestDeleteUnusedImagesExemptSandboxImage(t *testing.T) {

pkg/kubelet/images/image_manager.go

@@ -98,7 +98,7 @@ func (m *imageManager) logIt(ref *v1.ObjectReference, eventtype, event, prefix,
 
 // EnsureImageExists pulls the image for the specified pod and container, and returns
 // (imageRef, error message, error).
-func (m *imageManager) EnsureImageExists(ctx context.Context, pod *v1.Pod, container *v1.Container, pullSecrets []v1.Secret, podSandboxConfig *runtimeapi.PodSandboxConfig) (string, string, error) {
+func (m *imageManager) EnsureImageExists(ctx context.Context, pod *v1.Pod, container *v1.Container, pullSecrets []v1.Secret, podSandboxConfig *runtimeapi.PodSandboxConfig, podRuntimeHandler string) (string, string, error) {
 	logPrefix := fmt.Sprintf("%s/%s/%s", pod.Namespace, pod.Name, container.Image)
 	ref, err := kubecontainer.GenerateContainerRef(pod, container)
 	if err != nil {
@@ -122,9 +122,11 @@ func (m *imageManager) EnsureImageExists(ctx context.Context, pod *v1.Pod, conta
 	}
 
 	spec := kubecontainer.ImageSpec{
-		Image:       image,
+		Image:          image,
-		Annotations: podAnnotations,
+		Annotations:    podAnnotations,
+		RuntimeHandler: podRuntimeHandler,
 	}
+
 	imageRef, err := m.imageService.GetImageRef(ctx, spec)
 	if err != nil {
 		msg := fmt.Sprintf("Failed to inspect image %q: %v", container.Image, err)

pkg/kubelet/images/image_manager_test.go

@@ -28,9 +28,12 @@ import (
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/tools/record"
 	"k8s.io/client-go/util/flowcontrol"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
 	crierrors "k8s.io/cri-api/pkg/errors"
+	"k8s.io/kubernetes/pkg/features"
 	. "k8s.io/kubernetes/pkg/kubelet/container"
 	ctest "k8s.io/kubernetes/pkg/kubelet/container/testing"
 	testingclock "k8s.io/utils/clock/testing"
@@ -269,7 +272,7 @@ func TestParallelPuller(t *testing.T) {
 				fakeRuntime.CalledFunctions = nil
 				fakeClock.Step(time.Second)
 
-				_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil)
+				_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, "")
 				fakeRuntime.AssertCalls(expected.calls)
 				assert.Equal(t, expected.err, err)
 				assert.Equal(t, expected.shouldRecordStartedPullingTime, fakePodPullingTimeRecorder.startedPullingRecorded)
@@ -301,7 +304,7 @@ func TestSerializedPuller(t *testing.T) {
 				fakeRuntime.CalledFunctions = nil
 				fakeClock.Step(time.Second)
 
-				_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil)
+				_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, "")
 				fakeRuntime.AssertCalls(expected.calls)
 				assert.Equal(t, expected.err, err)
 				assert.Equal(t, expected.shouldRecordStartedPullingTime, fakePodPullingTimeRecorder.startedPullingRecorded)
@@ -364,7 +367,7 @@ func TestPullAndListImageWithPodAnnotations(t *testing.T) {
 		fakeRuntime.ImageList = []Image{}
 		fakeClock.Step(time.Second)
 
-		_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil)
+		_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, "")
 		fakeRuntime.AssertCalls(c.expected[0].calls)
 		assert.Equal(t, c.expected[0].err, err, "tick=%d", 0)
 		assert.Equal(t, c.expected[0].shouldRecordStartedPullingTime, fakePodPullingTimeRecorder.startedPullingRecorded)
@@ -375,6 +378,67 @@ func TestPullAndListImageWithPodAnnotations(t *testing.T) {
 
 		image := images[0]
 		assert.Equal(t, "missing_image:latest", image.ID, "Image ID")
+		assert.Equal(t, "", image.Spec.RuntimeHandler, "image.Spec.RuntimeHandler not empty", "ImageID", image.ID)
+
+		expectedAnnotations := []Annotation{
+			{
+				Name:  "kubernetes.io/runtimehandler",
+				Value: "handler_name",
+			}}
+		assert.Equal(t, expectedAnnotations, image.Spec.Annotations, "image spec annotations")
+	})
+}
+
+func TestPullAndListImageWithRuntimeHandlerInImageCriAPIFeatureGate(t *testing.T) {
+	runtimeHandler := "handler_name"
+	pod := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            "test_pod",
+			Namespace:       "test-ns",
+			UID:             "bar",
+			ResourceVersion: "42",
+			Annotations: map[string]string{
+				"kubernetes.io/runtimehandler": runtimeHandler,
+			},
+		},
+		Spec: v1.PodSpec{
+			RuntimeClassName: &runtimeHandler,
+		},
+	}
+	c := pullerTestCase{ // pull missing image
+		testName:       "test pull and list image with pod annotations",
+		containerImage: "missing_image",
+		policy:         v1.PullIfNotPresent,
+		inspectErr:     nil,
+		pullerErr:      nil,
+		expected: []pullerExpects{
+			{[]string{"GetImageRef", "PullImage"}, nil, true, true},
+		}}
+
+	useSerializedEnv := true
+	t.Run(c.testName, func(t *testing.T) {
+		defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClassInImageCriAPI, true)()
+		ctx := context.Background()
+		puller, fakeClock, fakeRuntime, container, fakePodPullingTimeRecorder := pullerTestEnv(t, c, useSerializedEnv, nil)
+		fakeRuntime.CalledFunctions = nil
+		fakeRuntime.ImageList = []Image{}
+		fakeClock.Step(time.Second)
+
+		_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, runtimeHandler)
+		fakeRuntime.AssertCalls(c.expected[0].calls)
+		assert.Equal(t, c.expected[0].err, err, "tick=%d", 0)
+		assert.Equal(t, c.expected[0].shouldRecordStartedPullingTime, fakePodPullingTimeRecorder.startedPullingRecorded)
+		assert.Equal(t, c.expected[0].shouldRecordFinishedPullingTime, fakePodPullingTimeRecorder.finishedPullingRecorded)
+
+		images, _ := fakeRuntime.ListImages(ctx)
+		assert.Equal(t, 1, len(images), "ListImages() count")
+
+		image := images[0]
+		assert.Equal(t, "missing_image:latest", image.ID, "Image ID")
+
+		// when RuntimeClassInImageCriAPI feature gate is enabled, check runtime
+		// handler information for every image in the ListImages() response
+		assert.Equal(t, runtimeHandler, image.Spec.RuntimeHandler, "runtime handler returned not as expected", "Image ID", image)
 
 		expectedAnnotations := []Annotation{
 			{
@@ -419,7 +483,7 @@ func TestMaxParallelImagePullsLimit(t *testing.T) {
 	for i := 0; i < maxParallelImagePulls; i++ {
 		wg.Add(1)
 		go func() {
-			_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil)
+			_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, "")
 			assert.Nil(t, err)
 			wg.Done()
 		}()
@@ -431,7 +495,7 @@ func TestMaxParallelImagePullsLimit(t *testing.T) {
 	for i := 0; i < 2; i++ {
 		wg.Add(1)
 		go func() {
-			_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil)
+			_, _, err := puller.EnsureImageExists(ctx, pod, container, nil, nil, "")
 			assert.Nil(t, err)
 			wg.Done()
 		}()

pkg/kubelet/images/types.go

@@ -48,7 +48,7 @@ var (
 // Implementations are expected to be thread safe.
 type ImageManager interface {
 	// EnsureImageExists ensures that image specified in `container` exists.
-	EnsureImageExists(ctx context.Context, pod *v1.Pod, container *v1.Container, pullSecrets []v1.Secret, podSandboxConfig *runtimeapi.PodSandboxConfig) (string, string, error)
+	EnsureImageExists(ctx context.Context, pod *v1.Pod, container *v1.Container, pullSecrets []v1.Secret, podSandboxConfig *runtimeapi.PodSandboxConfig, podRuntimeHandler string) (string, string, error)
 
 	// TODO(ronl): consolidating image managing and deleting operation in this interface
 }

pkg/kubelet/kuberuntime/convert.go

@@ -19,7 +19,9 @@ package kuberuntime
 import (
 	"sort"
 
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
+	"k8s.io/kubernetes/pkg/features"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 )
 
@@ -42,10 +44,20 @@ func toKubeContainerImageSpec(image *runtimeapi.Image) kubecontainer.ImageSpec {
 		}
 	}
 
-	return kubecontainer.ImageSpec{
+	spec := kubecontainer.ImageSpec{
 		Image:       image.Id,
 		Annotations: annotations,
 	}
+	// if RuntimeClassInImageCriAPI feature gate is enabled, set runtimeHandler CRI field
+	if utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI) {
+		runtimeHandler := ""
+		if image.Spec != nil {
+			runtimeHandler = image.Spec.RuntimeHandler
+		}
+		spec.RuntimeHandler = runtimeHandler
+	}
+
+	return spec
 }
 
 func toRuntimeAPIImageSpec(imageSpec kubecontainer.ImageSpec) *runtimeapi.ImageSpec {
@@ -55,8 +67,15 @@ func toRuntimeAPIImageSpec(imageSpec kubecontainer.ImageSpec) *runtimeapi.ImageS
 			annotations[a.Name] = a.Value
 		}
 	}
-	return &runtimeapi.ImageSpec{
+
+	spec := runtimeapi.ImageSpec{
 		Image:       imageSpec.Image,
 		Annotations: annotations,
 	}
+	// if RuntimeClassInImageCriAPI feature gate is enabled, set runtimeHandler CRI field
+	if utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI) {
+		spec.RuntimeHandler = imageSpec.RuntimeHandler
+	}
+
+	return &spec
 }

pkg/kubelet/kuberuntime/convert_test.go

@@ -21,7 +21,10 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
+	"k8s.io/kubernetes/pkg/features"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 )
 
@@ -103,27 +106,32 @@ func TestConvertToRuntimeAPIImageSpec(t *testing.T) {
 	}{
 		{
 			input: kubecontainer.ImageSpec{
-				Image:       "test",
+				Image:          "test",
-				Annotations: nil,
+				RuntimeHandler: "",
+				Annotations:    nil,
 			},
 			expected: &runtimeapi.ImageSpec{
-				Image:       "test",
+				Image:          "test",
-				Annotations: map[string]string{},
+				RuntimeHandler: "",
+				Annotations:    map[string]string{},
 			},
 		},
 		{
 			input: kubecontainer.ImageSpec{
-				Image:       "test",
+				Image:          "test",
-				Annotations: []kubecontainer.Annotation{},
+				RuntimeHandler: "",
+				Annotations:    []kubecontainer.Annotation{},
 			},
 			expected: &runtimeapi.ImageSpec{
-				Image:       "test",
+				Image:          "test",
-				Annotations: map[string]string{},
+				RuntimeHandler: "",
+				Annotations:    map[string]string{},
 			},
 		},
 		{
 			input: kubecontainer.ImageSpec{
-				Image: "test",
+				Image:          "test",
+				RuntimeHandler: "",
 				Annotations: []kubecontainer.Annotation{
 					{
 						Name:  "kubernetes.io/os",
@@ -136,7 +144,8 @@ func TestConvertToRuntimeAPIImageSpec(t *testing.T) {
 				},
 			},
 			expected: &runtimeapi.ImageSpec{
-				Image: "test",
+				Image:          "test",
+				RuntimeHandler: "",
 				Annotations: map[string]string{
 					"kubernetes.io/os":             "linux",
 					"kubernetes.io/runtimehandler": "handler",
@@ -150,3 +159,142 @@ func TestConvertToRuntimeAPIImageSpec(t *testing.T) {
 		assert.Equal(t, test.expected, actual)
 	}
 }
+
+func TestConvertToKubeContainerImageSpecWithRuntimeHandlerInImageSpecCri(t *testing.T) {
+	testCases := []struct {
+		input    *runtimeapi.Image
+		expected kubecontainer.ImageSpec
+	}{
+		{
+			input: &runtimeapi.Image{
+				Id:   "test",
+				Spec: nil,
+			},
+			expected: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    []kubecontainer.Annotation(nil),
+			},
+		},
+		{
+			input: &runtimeapi.Image{
+				Id: "test",
+				Spec: &runtimeapi.ImageSpec{
+					Annotations: nil,
+				},
+			},
+			expected: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    []kubecontainer.Annotation(nil),
+			},
+		},
+		{
+			input: &runtimeapi.Image{
+				Id: "test",
+				Spec: &runtimeapi.ImageSpec{
+					Annotations: map[string]string{},
+				},
+			},
+			expected: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    []kubecontainer.Annotation(nil),
+			},
+		},
+		{
+			input: &runtimeapi.Image{
+				Id: "test",
+				Spec: &runtimeapi.ImageSpec{
+					RuntimeHandler: "test-runtimeHandler",
+					Annotations: map[string]string{
+						"kubernetes.io/os":             "linux",
+						"kubernetes.io/runtimehandler": "handler",
+					},
+				},
+			},
+			expected: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "test-runtimeHandler",
+				Annotations: []kubecontainer.Annotation{
+					{
+						Name:  "kubernetes.io/os",
+						Value: "linux",
+					},
+					{
+						Name:  "kubernetes.io/runtimehandler",
+						Value: "handler",
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClassInImageCriAPI, true)()
+		actual := toKubeContainerImageSpec(test.input)
+		assert.Equal(t, test.expected, actual)
+	}
+}
+
+func TestConvertToRuntimeAPIImageSpecWithRuntimeHandlerInImageSpecCri(t *testing.T) {
+	testCases := []struct {
+		input    kubecontainer.ImageSpec
+		expected *runtimeapi.ImageSpec
+	}{
+		{
+			input: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    nil,
+			},
+			expected: &runtimeapi.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    map[string]string{},
+			},
+		},
+		{
+			input: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    []kubecontainer.Annotation{},
+			},
+			expected: &runtimeapi.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "",
+				Annotations:    map[string]string{},
+			},
+		},
+		{
+			input: kubecontainer.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "test-runtimeHandler",
+				Annotations: []kubecontainer.Annotation{
+					{
+						Name:  "kubernetes.io/os",
+						Value: "linux",
+					},
+					{
+						Name:  "kubernetes.io/runtimehandler",
+						Value: "handler",
+					},
+				},
+			},
+			expected: &runtimeapi.ImageSpec{
+				Image:          "test",
+				RuntimeHandler: "test-runtimeHandler",
+				Annotations: map[string]string{
+					"kubernetes.io/os":             "linux",
+					"kubernetes.io/runtimehandler": "handler",
+				},
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClassInImageCriAPI, true)()
+		actual := toRuntimeAPIImageSpec(test.input)
+		assert.Equal(t, test.expected, actual)
+	}
+}

pkg/kubelet/kuberuntime/helpers.go

@@ -97,6 +97,7 @@ func (m *kubeGenericRuntimeManager) toKubeContainer(c *runtimeapi.Container) (*k
 		ID:                   kubecontainer.ContainerID{Type: m.runtimeName, ID: c.Id},
 		Name:                 c.GetMetadata().GetName(),
 		ImageID:              c.ImageRef,
+		ImageRuntimeHandler:  c.Image.RuntimeHandler,
 		Image:                c.Image.Image,
 		Hash:                 annotatedInfo.Hash,
 		HashWithoutResources: annotatedInfo.HashWithoutResources,

pkg/kubelet/kuberuntime/helpers_test.go

@@ -25,8 +25,11 @@ import (
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
 	runtimetesting "k8s.io/cri-api/pkg/apis/testing"
+	"k8s.io/kubernetes/pkg/features"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 )
 
@@ -136,11 +139,53 @@ func TestToKubeContainer(t *testing.T) {
 			Type: runtimetesting.FakeRuntimeName,
 			ID:   "test-id",
 		},
-		Name:    "test-name",
+		Name:                "test-name",
-		ImageID: "test-image-ref",
+		ImageID:             "test-image-ref",
-		Image:   "test-image",
+		Image:               "test-image",
-		Hash:    uint64(0x1234),
+		ImageRuntimeHandler: "",
-		State:   kubecontainer.ContainerStateRunning,
+		Hash:                uint64(0x1234),
+		State:               kubecontainer.ContainerStateRunning,
+	}
+
+	_, _, m, err := createTestRuntimeManager()
+	assert.NoError(t, err)
+	got, err := m.toKubeContainer(c)
+	assert.NoError(t, err)
+	assert.Equal(t, expect, got)
+
+	// unable to convert a nil pointer to a runtime container
+	_, err = m.toKubeContainer(nil)
+	assert.Error(t, err)
+	_, err = m.sandboxToKubeContainer(nil)
+	assert.Error(t, err)
+}
+
+func TestToKubeContainerWithRuntimeHandlerInImageSpecCri(t *testing.T) {
+	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClassInImageCriAPI, true)()
+	c := &runtimeapi.Container{
+		Id: "test-id",
+		Metadata: &runtimeapi.ContainerMetadata{
+			Name:    "test-name",
+			Attempt: 1,
+		},
+		Image:    &runtimeapi.ImageSpec{Image: "test-image", RuntimeHandler: "test-runtimeHandler"},
+		ImageRef: "test-image-ref",
+		State:    runtimeapi.ContainerState_CONTAINER_RUNNING,
+		Annotations: map[string]string{
+			containerHashLabel: "1234",
+		},
+	}
+	expect := &kubecontainer.Container{
+		ID: kubecontainer.ContainerID{
+			Type: runtimetesting.FakeRuntimeName,
+			ID:   "test-id",
+		},
+		Name:                "test-name",
+		ImageID:             "test-image-ref",
+		Image:               "test-image",
+		ImageRuntimeHandler: "test-runtimeHandler",
+		Hash:                uint64(0x1234),
+		State:               kubecontainer.ContainerStateRunning,
 	}
 
 	_, _, m, err := createTestRuntimeManager()

pkg/kubelet/kuberuntime/kuberuntime_container.go

@@ -179,7 +179,23 @@ func (m *kubeGenericRuntimeManager) startContainer(ctx context.Context, podSandb
 	container := spec.container
 
 	// Step 1: pull the image.
-	imageRef, msg, err := m.imagePuller.EnsureImageExists(ctx, pod, container, pullSecrets, podSandboxConfig)
+
+	// If RuntimeClassInImageCriAPI feature gate is enabled, pass runtimehandler
+	// information for the runtime class specified. If not runtime class is
+	// specified, then pass ""
+	podRuntimeHandler := ""
+	var err error
+	if utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI) {
+		if pod.Spec.RuntimeClassName != nil && *pod.Spec.RuntimeClassName != "" {
+			podRuntimeHandler, err = m.runtimeClassManager.LookupRuntimeHandler(pod.Spec.RuntimeClassName)
+			if err != nil {
+				msg := fmt.Sprintf("Failed to lookup runtimeHandler for runtimeClassName %v", pod.Spec.RuntimeClassName)
+				return msg, err
+			}
+		}
+	}
+
+	imageRef, msg, err := m.imagePuller.EnsureImageExists(ctx, pod, container, pullSecrets, podSandboxConfig, podRuntimeHandler)
 	if err != nil {
 		s, _ := grpcstatus.FromError(err)
 		m.recordContainerEvent(pod, container, "", v1.EventTypeWarning, events.FailedToCreateContainer, "Error: %v", s.Message())
@@ -601,6 +617,7 @@ func toKubeContainerStatus(status *runtimeapi.ContainerStatus, runtimeName strin
 		Name:                 labeledInfo.ContainerName,
 		Image:                status.Image.Image,
 		ImageID:              status.ImageRef,
+		ImageRuntimeHandler:  status.Image.RuntimeHandler,
 		Hash:                 annotatedInfo.Hash,
 		HashWithoutResources: annotatedInfo.HashWithoutResources,
 		RestartCount:         annotatedInfo.RestartCount,

pkg/kubelet/kuberuntime/kuberuntime_image.go

@@ -21,9 +21,11 @@ import (
 
 	v1 "k8s.io/api/core/v1"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
 	"k8s.io/klog/v2"
 	credentialprovidersecrets "k8s.io/kubernetes/pkg/credentialprovider/secrets"
+	"k8s.io/kubernetes/pkg/features"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 	"k8s.io/kubernetes/pkg/util/parsers"
 )
@@ -105,6 +107,17 @@ func (m *kubeGenericRuntimeManager) ListImages(ctx context.Context) ([]kubeconta
 	}
 
 	for _, img := range allImages {
+		// Container runtimes may choose not to implement changes needed for KEP 4216. If
+		// the changes are not implemented by a container runtime, the exisiting behavior
+		// of not populating the runtimeHandler CRI field in ImageSpec struct is preserved.
+		// Therefore, when RuntimeClassInImageCriAPI feature gate is set, check to see if this
+		// field is empty and log a warning message.
+		if utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClassInImageCriAPI) {
+			if img.Spec == nil || (img.Spec != nil && img.Spec.RuntimeHandler == "") {
+				klog.V(2).InfoS("WARNING: RuntimeHandler is empty", "ImageID", img.Id)
+			}
+		}
+
 		images = append(images, kubecontainer.Image{
 			ID:          img.Id,
 			Size:        int64(img.Size_),