Merge pull request kubernetes#126303 from bart0sh/PR150-dra-refactor-checkpoint-upstream

DRA: refactor checkpointing

Author: k8s-ci-robot

pkg/kubelet/cm/dra/claiminfo.go

@@ -40,8 +40,8 @@ type ClaimInfo struct {
 // claimInfoCache is a cache of processed resource claims keyed by namespace/claimname.
 type claimInfoCache struct {
 	sync.RWMutex
-	state     state.CheckpointState
-	claimInfo map[string]*ClaimInfo
+	checkpointer state.Checkpointer
+	claimInfo    map[string]*ClaimInfo
 }
 
 // newClaimInfoFromClaim creates a new claim info from a resource claim.
@@ -77,12 +77,12 @@ func newClaimInfoFromState(state *state.ClaimInfoState) *ClaimInfo {
 }
 
 // setCDIDevices adds a set of CDI devices to the claim info.
-func (info *ClaimInfo) addDevice(driverName string, device state.Device) {
+func (info *ClaimInfo) addDevice(driverName string, deviceState state.Device) {
 	if info.DriverState == nil {
 		info.DriverState = make(map[string]state.DriverState)
 	}
 	driverState := info.DriverState[driverName]
-	driverState.Devices = append(driverState.Devices, device)
+	driverState.Devices = append(driverState.Devices, deviceState)
 	info.DriverState[driverName] = driverState
 }
 
@@ -113,22 +113,27 @@ func (info *ClaimInfo) isPrepared() bool {
 
 // newClaimInfoCache creates a new claim info cache object, pre-populated from a checkpoint (if present).
 func newClaimInfoCache(stateDir, checkpointName string) (*claimInfoCache, error) {
-	stateImpl, err := state.NewCheckpointState(stateDir, checkpointName)
+	checkpointer, err := state.NewCheckpointer(stateDir, checkpointName)
 	if err != nil {
 		return nil, fmt.Errorf("could not initialize checkpoint manager, please drain node and remove dra state file, err: %+v", err)
 	}
 
-	curState, err := stateImpl.GetOrCreate()
+	checkpoint, err := checkpointer.GetOrCreate()
 	if err != nil {
 		return nil, fmt.Errorf("error calling GetOrCreate() on checkpoint state: %v", err)
 	}
 
 	cache := &claimInfoCache{
-		state:     stateImpl,
-		claimInfo: make(map[string]*ClaimInfo),
+		checkpointer: checkpointer,
+		claimInfo:    make(map[string]*ClaimInfo),
 	}
 
-	for _, entry := range curState {
+	entries, err := checkpoint.GetClaimInfoStateList()
+	if err != nil {
+		return nil, fmt.Errorf("error calling GetEntries() on checkpoint: %w", err)
+
+	}
+	for _, entry := range entries {
 		info := newClaimInfoFromState(&entry)
 		cache.claimInfo[info.Namespace+"/"+info.ClaimName] = info
 	}
@@ -192,7 +197,11 @@ func (cache *claimInfoCache) syncToCheckpoint() error {
 	for _, infoClaim := range cache.claimInfo {
 		claimInfoStateList = append(claimInfoStateList, infoClaim.ClaimInfoState)
 	}
-	return cache.state.Store(claimInfoStateList)
+	checkpoint, err := state.NewCheckpoint(claimInfoStateList)
+	if err != nil {
+		return err
+	}
+	return cache.checkpointer.Store(checkpoint)
 }
 
 // cdiDevicesAsList returns a list of CDIDevices from the provided claim info.

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

@@ -18,51 +18,90 @@ package state
 
 import (
 	"encoding/json"
+	"hash/crc32"
 
-	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
-	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/errors"
 )
 
-var _ checkpointmanager.Checkpoint = &DRAManagerCheckpoint{}
+const (
+	CheckpointAPIGroup   = "checkpoint.dra.kubelet.k8s.io"
+	CheckpointKind       = "DRACheckpoint"
+	CheckpointAPIVersion = CheckpointAPIGroup + "/v1"
+)
 
-const checkpointVersion = "v1"
+// Checkpoint represents a structure to store DRA checkpoint data
+type Checkpoint struct {
+	// Data is a JSON serialized checkpoint data
+	Data string
+	// Checksum is a checksum of Data
+	Checksum uint32
+}
 
-// DRAManagerCheckpoint struct is used to store pod dynamic resources assignments in a checkpoint
-type DRAManagerCheckpoint struct {
-	Version  string             `json:"version"`
-	Entries  ClaimInfoStateList `json:"entries,omitempty"`
-	Checksum checksum.Checksum  `json:"checksum"`
+type CheckpointData struct {
+	metav1.TypeMeta
+	ClaimInfoStateList ClaimInfoStateList
 }
 
-// List of claim info to store in checkpoint
-type ClaimInfoStateList []ClaimInfoState
+// NewCheckpoint creates a new checkpoint from a list of claim info states
+func NewCheckpoint(data ClaimInfoStateList) (*Checkpoint, error) {
+	cpData := &CheckpointData{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       CheckpointKind,
+			APIVersion: CheckpointAPIVersion,
+		},
+		ClaimInfoStateList: data,
+	}
+
+	cpDataBytes, err := json.Marshal(cpData)
+	if err != nil {
+		return nil, err
+	}
 
-// NewDRAManagerCheckpoint returns an instance of Checkpoint
-func NewDRAManagerCheckpoint() *DRAManagerCheckpoint {
-	return &DRAManagerCheckpoint{
-		Version: checkpointVersion,
-		Entries: ClaimInfoStateList{},
+	cp := &Checkpoint{
+		Data:     string(cpDataBytes),
+		Checksum: crc32.ChecksumIEEE(cpDataBytes),
 	}
+
+	return cp, nil
 }
 
-// MarshalCheckpoint returns marshalled checkpoint
-func (dc *DRAManagerCheckpoint) MarshalCheckpoint() ([]byte, error) {
-	// make sure checksum wasn't set before so it doesn't affect output checksum
-	dc.Checksum = 0
-	dc.Checksum = checksum.New(dc)
-	return json.Marshal(*dc)
+// MarshalCheckpoint marshals checkpoint to JSON
+func (cp *Checkpoint) MarshalCheckpoint() ([]byte, error) {
+	return json.Marshal(cp)
 }
 
-// UnmarshalCheckpoint tries to unmarshal passed bytes to checkpoint
-func (dc *DRAManagerCheckpoint) UnmarshalCheckpoint(blob []byte) error {
-	return json.Unmarshal(blob, dc)
+// UnmarshalCheckpoint unmarshals checkpoint from JSON
+// and verifies its data checksum
+func (cp *Checkpoint) UnmarshalCheckpoint(blob []byte) error {
+	if err := json.Unmarshal(blob, cp); err != nil {
+		return err
+	}
+
+	// verify checksum
+	if err := cp.VerifyChecksum(); err != nil {
+		return err
+	}
+
+	return nil
 }
 
-// VerifyChecksum verifies that current checksum of checkpoint is valid
-func (dc *DRAManagerCheckpoint) VerifyChecksum() error {
-	ck := dc.Checksum
-	dc.Checksum = 0
-	err := ck.Verify(dc)
-	dc.Checksum = ck
-	return err
+// VerifyChecksum verifies that current checksum
+// of checkpointed Data is valid
+func (cp *Checkpoint) VerifyChecksum() error {
+	expectedCS := crc32.ChecksumIEEE([]byte(cp.Data))
+	if expectedCS != cp.Checksum {
+		return &errors.CorruptCheckpointError{ActualCS: uint64(cp.Checksum), ExpectedCS: uint64(expectedCS)}
+	}
+	return nil
+}
+
+// GetClaimInfoStateList returns list of claim info states from checkpoint
+func (cp *Checkpoint) GetClaimInfoStateList() (ClaimInfoStateList, error) {
+	var data CheckpointData
+	if err := json.Unmarshal([]byte(cp.Data), &data); err != nil {
+		return nil, err
+	}
+
+	return data.ClaimInfoStateList, nil
 }

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

@@ -0,0 +1,98 @@
+/*
+Copyright 2023 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 (
+	"errors"
+	"fmt"
+	"sync"
+
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
+	checkpointerrors "k8s.io/kubernetes/pkg/kubelet/checkpointmanager/errors"
+)
+
+type Checkpointer interface {
+	GetOrCreate() (*Checkpoint, error)
+	Store(*Checkpoint) error
+}
+
+type checkpointer struct {
+	sync.RWMutex
+	checkpointManager checkpointmanager.CheckpointManager
+	checkpointName    string
+}
+
+// NewCheckpointer creates new checkpointer for keeping track of claim info  with checkpoint backend
+func NewCheckpointer(stateDir, checkpointName string) (Checkpointer, error) {
+	if len(checkpointName) == 0 {
+		return nil, fmt.Errorf("received empty string instead of checkpointName")
+	}
+
+	checkpointManager, err := checkpointmanager.NewCheckpointManager(stateDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to initialize checkpoint manager: %w", err)
+	}
+
+	checkpointer := &checkpointer{
+		checkpointManager: checkpointManager,
+		checkpointName:    checkpointName,
+	}
+
+	return checkpointer, nil
+}
+
+// GetOrCreate gets list of claim info states from a checkpoint
+// or creates empty list if checkpoint doesn't exist
+func (sc *checkpointer) GetOrCreate() (*Checkpoint, error) {
+	sc.Lock()
+	defer sc.Unlock()
+
+	checkpoint, err := NewCheckpoint(nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new checkpoint: %w", err)
+	}
+
+	err = sc.checkpointManager.GetCheckpoint(sc.checkpointName, checkpoint)
+	if errors.Is(err, checkpointerrors.ErrCheckpointNotFound) {
+		err = sc.store(checkpoint)
+		if err != nil {
+			return nil, fmt.Errorf("failed to store checkpoint %v: %w", sc.checkpointName, err)
+		}
+		return checkpoint, nil
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to get checkpoint %v: %w", sc.checkpointName, err)
+	}
+
+	return checkpoint, nil
+}
+
+// Store stores checkpoint to the file
+func (sc *checkpointer) Store(checkpoint *Checkpoint) error {
+	sc.Lock()
+	defer sc.Unlock()
+
+	return sc.store(checkpoint)
+}
+
+// store saves state to a checkpoint, caller is responsible for locking
+func (sc *checkpointer) store(checkpoint *Checkpoint) error {
+	if err := sc.checkpointManager.CreateCheckpoint(sc.checkpointName, checkpoint); err != nil {
+		return fmt.Errorf("could not save checkpoint %s: %w", sc.checkpointName, err)
+	}
+	return nil
+}