DRA: refactor checkpointing

Author: bart0sh

pkg/kubelet/checkpointmanager/checkpoint_manager.go

@@ -85,9 +85,6 @@ func (manager *impl) GetCheckpoint(checkpointKey string, checkpoint Checkpoint)
 		return err
 	}
 	err = checkpoint.UnmarshalCheckpoint(blob)
-	if err == nil {
-		err = checkpoint.VerifyChecksum()
-	}
 	return err
 }
 

pkg/kubelet/cm/dra/state/checkpoint.go

@@ -0,0 +1,230 @@
+/*
+Copyright 2024 The Kubernetes 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 state
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+)
+
+// V1 data structures are copies of the correspondent structures from state_checkpoint.go (k/k 1.30)
+// This is done to avoid breaking changes
+//
+// NOTE: These structures must not be changed, new checkpoint version(s) should be added instead
+
+type CheckpointV1 struct {
+	CheckpointVersion
+	Data     []ClaimInfoStateV1 `json:"entries,omitempty"`
+	Checksum checksum.Checksum  `json:"checksum"`
+}
+
+// ClaimInfoState is used to store claim info state in a checkpoint
+type ClaimInfoStateV1 struct {
+	// Name of the DRA driver
+	DriverName string
+
+	// ClassName is a resource class of the claim
+	ClassName string
+
+	// ClaimUID is an UID of the resource claim
+	ClaimUID types.UID
+
+	// ClaimName is a name of the resource claim
+	ClaimName string
+
+	// Namespace is a claim namespace
+	Namespace string
+
+	// PodUIDs is a set of pod UIDs that reference a resource
+	PodUIDs sets.Set[string]
+
+	// ResourceHandles is a list of opaque resource data for processing by a specific kubelet plugin
+	ResourceHandles []ResourceHandleV1
+
+	// CDIDevices is a map of DriverName --> CDI devices returned by the
+	// GRPC API call NodePrepareResource
+	CDIDevices map[string][]string
+}
+
+// ResourceHandle holds opaque resource data for processing by a specific kubelet plugin.
+type ResourceHandleV1 struct {
+	// DriverName specifies the name of the resource driver whose kubelet
+	// plugin should be invoked to process this ResourceHandle's data once it
+	// lands on a node. This may differ from the DriverName set in
+	// ResourceClaimStatus this ResourceHandle is embedded in.
+	DriverName string `json:"driverName,omitempty" protobuf:"bytes,1,opt,name=driverName"`
+
+	// Data contains the opaque data associated with this ResourceHandle. It is
+	// set by the controller component of the resource driver whose name
+	// matches the DriverName set in the ResourceClaimStatus this
+	// ResourceHandle is embedded in. It is set at allocation time and is
+	// intended for processing by the kubelet plugin whose name matches
+	// the DriverName set in this ResourceHandle.
+	//
+	// The maximum size of this field is 16KiB. This may get increased in the
+	// future, but not reduced.
+	// +optional
+	Data string `json:"data,omitempty" protobuf:"bytes,2,opt,name=data"`
+
+	// If StructuredData is set, then it needs to be used instead of Data.
+	//
+	// +optional
+	StructuredData *StructuredResourceHandleV1 `json:"structuredData,omitempty" protobuf:"bytes,5,opt,name=structuredData"`
+}
+
+// StructuredResourceHandle is the in-tree representation of the allocation result.
+type StructuredResourceHandleV1 struct {
+	// VendorClassParameters are the per-claim configuration parameters
+	// from the resource class at the time that the claim was allocated.
+	//
+	// +optional
+	VendorClassParameters runtime.RawExtension `json:"vendorClassParameters,omitempty" protobuf:"bytes,1,opt,name=vendorClassParameters"`
+
+	// VendorClaimParameters are the per-claim configuration parameters
+	// from the resource claim parameters at the time that the claim was
+	// allocated.
+	//
+	// +optional
+	VendorClaimParameters runtime.RawExtension `json:"vendorClaimParameters,omitempty" protobuf:"bytes,2,opt,name=vendorClaimParameters"`
+
+	// NodeName is the name of the node providing the necessary resources
+	// if the resources are local to a node.
+	//
+	// +optional
+	NodeName string `json:"nodeName,omitempty" protobuf:"bytes,4,name=nodeName"`
+
+	// Results lists all allocated driver resources.
+	//
+	// +listType=atomic
+	Results []DriverAllocationResultV1 `json:"results" protobuf:"bytes,5,name=results"`
+}
+
+// DriverAllocationResult contains vendor parameters and the allocation result for
+// one request.
+type DriverAllocationResultV1 struct {
+	// VendorRequestParameters are the per-request configuration parameters
+	// from the time that the claim was allocated.
+	//
+	// +optional
+	VendorRequestParameters runtime.RawExtension `json:"vendorRequestParameters,omitempty" protobuf:"bytes,1,opt,name=vendorRequestParameters"`
+
+	AllocationResultModelV1 `json:",inline" protobuf:"bytes,2,name=allocationResultModel"`
+}
+
+// AllocationResultModel must have one and only one field set.
+type AllocationResultModelV1 struct {
+	// NamedResources describes the allocation result when using the named resources model.
+	//
+	// +optional
+	NamedResources *NamedResourcesAllocationResultV1 `json:"namedResources,omitempty" protobuf:"bytes,1,opt,name=namedResources"`
+}
+
+// NamedResourcesAllocationResult is used in AllocationResultModel.
+type NamedResourcesAllocationResultV1 struct {
+	// Name is the name of the selected resource instance.
+	Name string `json:"name" protobuf:"bytes,1,name=name"`
+}
+
+func NewCheckpointV1(data ClaimInfoStateList) *CheckpointV1 {
+	cp := &CheckpointV1{
+		CheckpointVersion: CheckpointVersion{Version: v1},
+	}
+	cp.SetData(data)
+	return cp
+}
+
+// SetData sets checkpoint data from ClaimInfoStateList
+func (cp *CheckpointV1) SetData(data ClaimInfoStateList) {
+	cp.Data = make([]ClaimInfoStateV1, len(data))
+	for i, entry := range data {
+		cp.Data[i] = ClaimInfoStateV1{
+			ClaimUID:   entry.ClaimUID,
+			ClaimName:  entry.ClaimName,
+			Namespace:  entry.Namespace,
+			PodUIDs:    entry.PodUIDs,
+			CDIDevices: map[string][]string{},
+		}
+		for driverName, driverState := range entry.DriverState {
+			cp.Data[i].CDIDevices[driverName] = make([]string, len(driverState.Devices))
+			for j, device := range driverState.Devices {
+				cp.Data[i].CDIDevices[driverName][j] = device.DeviceName
+			}
+		}
+	}
+	cp.Checksum = checksum.New(cp)
+}
+
+// GetData returns checkpoint data in a form of ClaimInfoStateList
+func (cp *CheckpointV1) GetData() ClaimInfoStateList {
+	result := make(ClaimInfoStateList, len(cp.Data))
+	for i, entry := range cp.Data {
+		result[i] = ClaimInfoState{
+			ClaimUID:    entry.ClaimUID,
+			ClaimName:   entry.ClaimName,
+			Namespace:   entry.Namespace,
+			PodUIDs:     entry.PodUIDs,
+			DriverState: map[string]DriverState{},
+		}
+		for driverName, devices := range entry.CDIDevices {
+			result[i].DriverState[driverName] = DriverState{
+				Devices: make([]Device, len(devices)),
+			}
+			for j, cdiID := range devices {
+				result[i].DriverState[driverName].Devices[j] = Device{
+					CDIDeviceIDs: []string{cdiID},
+				}
+			}
+		}
+	}
+	return result
+}
+
+// MarshalCheckpoint marshals checkpoint to JSON
+func (cp *CheckpointV1) MarshalCheckpoint() ([]byte, error) {
+	// make sure checksum wasn't set before so it doesn't affect output checksum
+	cp.Checksum = 0
+	cp.Checksum = checksum.New(cp)
+	return json.Marshal(*cp)
+}
+
+// UnmarshalCheckpoint tries to unmarshal passed bytes to checkpoint
+func (cp *CheckpointV1) UnmarshalCheckpoint(blob []byte) error {
+	version := &CheckpointVersion{}
+	if err := json.Unmarshal(blob, version); err != nil {
+		return err
+	}
+	if version.Version == v1 {
+		// NOTE: We can't verify V1 checksum due to the changes in structures names
+		// e.g. DRAManagerCheckpoint -> CheckpointV1
+		return json.Unmarshal(blob, cp)
+	}
+	return fmt.Errorf("unsupported checkpoint version %s", version.Version)
+}
+
+// VerifyChecksum verifies that current checksum of checkpoint is valid
+func (cp *CheckpointV1) VerifyChecksum() error {
+	ck := cp.Checksum
+	cp.Checksum = 0
+	err := ck.Verify(cp)
+	cp.Checksum = ck
+	return err
+}