DRA: refactor checkpointing

Co-authored-by: Kevin Klues <klueska@gmail.com>

Author: bart0sh

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

@@ -0,0 +1,86 @@
+/*
+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 checkpoint
+
+import (
+	"encoding/json"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+	state "k8s.io/kubernetes/pkg/kubelet/cm/dra/state/v1"
+)
+
+type Checkpoint struct {
+	Data     string
+	Checksum checksum.Checksum
+}
+
+func NewCheckpoint(data state.ClaimInfoStateList) (*Checkpoint, error) {
+	cpData := &state.CheckpointData{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       state.CheckpointKind,
+			APIVersion: state.CheckpointAPIVersion,
+		},
+		Entries: data,
+	}
+
+	cpDataBytes, err := json.Marshal(cpData)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Checkpoint{
+		Data:     string(cpDataBytes),
+		Checksum: checksum.New(string(cpDataBytes)),
+	}, nil
+}
+
+// MarshalCheckpoint marshals checkpoint to JSON
+func (cp *Checkpoint) MarshalCheckpoint() ([]byte, error) {
+	return json.Marshal(cp)
+}
+
+// UnmarshalCheckpoint unmarshals checkpoint from JSON
+// and verifies its 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 checkpointed Data is valid
+func (cp *Checkpoint) VerifyChecksum() error {
+	return cp.Checksum.Verify(cp.Data)
+}
+
+// GetEntries returns list of claim info states from checkpoint
+func (cp *Checkpoint) GetEntries() (state.ClaimInfoStateList, error) {
+	var cpData state.CheckpointData
+	if err := json.Unmarshal([]byte(cp.Data), &cpData); err != nil {
+		return nil, err
+	}
+
+	return cpData.Entries, nil
+}

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

@@ -0,0 +1,97 @@
+/*
+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 checkpoint
+
+import (
+	"fmt"
+	"sync"
+
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/errors"
+	state "k8s.io/kubernetes/pkg/kubelet/cm/dra/state/v1"
+)
+
+type Checkpointer interface {
+	GetOrCreate() (state.ClaimInfoStateList, error)
+	Store(state.ClaimInfoStateList) 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: %v", err)
+	}
+	checkpointer := &checkpointer{
+		checkpointManager: checkpointManager,
+		checkpointName:    checkpointName,
+	}
+
+	return checkpointer, nil
+}
+
+// GetOrCreate gets list of claim info states from a checkpoint
+// or creates empty list it checkpoint doesn't exist yet
+func (sc *checkpointer) GetOrCreate() (state.ClaimInfoStateList, error) {
+	sc.Lock()
+	defer sc.Unlock()
+
+	checkpoint, err := NewCheckpoint(nil)
+	err = sc.checkpointManager.GetCheckpoint(sc.checkpointName, checkpoint)
+	if err == errors.ErrCheckpointNotFound {
+		sc.store(state.ClaimInfoStateList{})
+		return state.ClaimInfoStateList{}, nil
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to get checkpoint %v: %w", sc.checkpointName, err)
+	}
+
+	return checkpoint.GetEntries()
+}
+
+// saves state to a checkpoint
+func (sc *checkpointer) Store(claimInfoStateList state.ClaimInfoStateList) error {
+	sc.Lock()
+	defer sc.Unlock()
+
+	return sc.store(claimInfoStateList)
+}
+
+// saves state to a checkpoint, caller is responsible for locking
+func (sc *checkpointer) store(claimInfoStateList state.ClaimInfoStateList) error {
+	checkpoint, err := NewCheckpoint(claimInfoStateList)
+
+	if err != nil {
+		return err
+	}
+
+	err = sc.checkpointManager.CreateCheckpoint(sc.checkpointName, checkpoint)
+	if err != nil {
+		return fmt.Errorf("could not save checkpoint %s: %v", sc.checkpointName, err)
+	}
+	return nil
+}