Replace k/k checkpoint manager with file reads/writes

Author: nojnhuh

cmd/dra-example-kubeletplugin/checkpoint.go

@@ -1,37 +1,87 @@
+/*
+ * Copyright 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 main
 
 import (
 	"encoding/json"
-
-	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
-	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+	"fmt"
+	"os"
+	"path/filepath"
 
 	"sigs.k8s.io/dra-example-driver/internal/profiles"
 )
 
 type PreparedClaims map[string]profiles.PreparedDevices
 
 type Checkpoint struct {
-	Checksum checksum.Checksum `json:"checksum"`
-	V1       *CheckpointV1     `json:"v1,omitempty"`
+	V1 *CheckpointV1 `json:"v1,omitempty"`
 }
 
-var _ checkpointmanager.Checkpoint = &Checkpoint{}
-
 type CheckpointV1 struct {
 	PreparedClaims PreparedClaims `json:"preparedClaims,omitempty"`
 }
 
 func newCheckpoint() *Checkpoint {
 	pc := &Checkpoint{
-		Checksum: 0,
-		V1: &CheckpointV1{
-			PreparedClaims: make(PreparedClaims),
-		},
+		V1: &CheckpointV1{},
 	}
 	return pc
 }
 
+func readCheckpoint(path string) (*Checkpoint, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	checkpoint := new(Checkpoint)
+	err = json.Unmarshal(data, checkpoint)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal json from %s: %w", path, err)
+	}
+	return checkpoint, nil
+}
+
+func writeCheckpoint(path string, checkpoint *Checkpoint) (err error) {
+	data, err := json.Marshal(checkpoint)
+	if err != nil {
+		return fmt.Errorf("marshal json: %w", err)
+	}
+	dir := filepath.Dir(path)
+	tmp, err := os.CreateTemp(dir, "tmp-checkpoint-*")
+	if err != nil {
+		return fmt.Errorf("create temp file in %s: %w", dir, err)
+	}
+	defer func() {
+		if err1 := tmp.Close(); err1 != nil && err == nil {
+			err = fmt.Errorf("close temp file: %w", err1)
+		}
+	}()
+	if _, err := tmp.Write(data); err != nil {
+		return fmt.Errorf("write to temp file %s: %w", tmp.Name(), err)
+	}
+	if err := tmp.Sync(); err != nil {
+		return fmt.Errorf("sync temp file: %w", err)
+	}
+	if err := os.Rename(tmp.Name(), path); err != nil {
+		return fmt.Errorf("rename %s to %s: %w", tmp.Name(), path, err)
+	}
+	return nil
+}
+
 func (cp *Checkpoint) GetPreparedDevices(claimUID string) profiles.PreparedDevices {
 	if cp.V1 == nil {
 		return nil
@@ -47,6 +97,10 @@ func (cp *Checkpoint) AddPreparedDevices(claimUID string, pds profiles.PreparedD
 		return
 	}
 
+	if cp.V1.PreparedClaims == nil {
+		cp.V1.PreparedClaims = make(PreparedClaims)
+	}
+
 	cp.V1.PreparedClaims[claimUID] = pds
 }
 
@@ -57,30 +111,3 @@ func (cp *Checkpoint) RemovePreparedDevices(claimUID string) {
 
 	delete(cp.V1.PreparedClaims, claimUID)
 }
-
-func (cp *Checkpoint) MarshalCheckpoint() ([]byte, error) {
-	cp.Checksum = 0
-	out, err := json.Marshal(*cp)
-	if err != nil {
-		return nil, err
-	}
-	cp.Checksum = checksum.New(out)
-	return json.Marshal(*cp)
-}
-
-func (cp *Checkpoint) UnmarshalCheckpoint(data []byte) error {
-	return json.Unmarshal(data, cp)
-}
-
-func (cp *Checkpoint) VerifyChecksum() error {
-	ck := cp.Checksum
-	cp.Checksum = 0
-	defer func() {
-		cp.Checksum = ck
-	}()
-	out, err := json.Marshal(*cp)
-	if err != nil {
-		return err
-	}
-	return ck.Verify(out)
-}

cmd/dra-example-kubeletplugin/checkpoint_test.go

@@ -0,0 +1,108 @@
+/*
+ * Copyright 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 main
+
+import (
+	"io/fs"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+	"tags.cncf.io/container-device-interface/specs-go"
+
+	"sigs.k8s.io/dra-example-driver/internal/profiles"
+)
+
+func TestReadWriteCheckpointRoundtrip(t *testing.T) {
+	tests := map[string]struct {
+		checkpoint *Checkpoint
+	}{
+		"new checkpoint": {
+			checkpoint: newCheckpoint(),
+		},
+		"populated checkpoint": {
+			checkpoint: &Checkpoint{
+				V1: &CheckpointV1{
+					PreparedClaims{
+						"uid": profiles.PreparedDevices{
+							{
+								Device: drapbv1.Device{
+									RequestNames: []string{"req"},
+									PoolName:     "pool",
+									DeviceName:   "dev",
+									CdiDeviceIds: []string{"id"},
+								},
+								ContainerEdits: &cdi.ContainerEdits{
+									ContainerEdits: &specs.ContainerEdits{
+										Env: []string{"KEY=value"},
+									},
+								},
+								AdminAccess: true,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			dir := t.TempDir()
+			path := filepath.Join(dir, DriverPluginCheckpointFile)
+
+			checkpoint, err := readCheckpoint(path)
+			assert.Nil(t, checkpoint)
+			assert.ErrorIs(t, err, fs.ErrNotExist)
+
+			checkpoint = test.checkpoint
+			err = writeCheckpoint(path, checkpoint)
+			require.NoError(t, err)
+
+			read, err := readCheckpoint(path)
+			require.NoError(t, err)
+			assert.Equal(t, test.checkpoint, read)
+		})
+	}
+
+	// The checkpoint format used to contain a checksum. This test ensures that
+	// checkpoints written in the old format can still be read.
+	t.Run("old checkpoint format", func(t *testing.T) {
+		dir := t.TempDir()
+		path := filepath.Join(dir, DriverPluginCheckpointFile)
+
+		old := `{
+			"checksum": 1,
+			"v1": {
+				"preparedClaims": {
+					"uid": []
+				}
+			}
+		}`
+
+		err := os.WriteFile(path, []byte(old), 0o600)
+		require.NoError(t, err)
+
+		checkpoint, err := readCheckpoint(path)
+		assert.NoError(t, err)
+		assert.NotNil(t, checkpoint.V1.PreparedClaims)
+	})
+}