Encapsulate Checkpoint internal state

Makefile

@@ -102,15 +102,24 @@ coverage: test
 	cat $(COVERAGE_FILE) | grep -v "_mock.go" > $(COVERAGE_FILE).no-mocks
 	go tool cover -func=$(COVERAGE_FILE).no-mocks
 
-generate: generate-deepcopy
+generate: generate-deepcopy generate-conversion
 
 generate-deepcopy:
 	for api in $(APIS); do \
-		rm -f $(CURDIR)/api/$(VENDOR)/resource/$${api}/zz_generated.deepcopy.go; \
+		rm -f $${api}/zz_generated.deepcopy.go; \
 		controller-gen \
 			object:headerFile=$(CURDIR)/hack/boilerplate.generatego.txt \
-			paths=$(CURDIR)/api/$(VENDOR)/resource/$${api}/ \
-			output:object:dir=$(CURDIR)/api/$(VENDOR)/resource/$${api}; \
+			paths=$${api}/ \
+			output:object:dir=$${api}; \
+	done
+
+generate-conversion:
+	for api in $(APIS); do \
+		rm -f $${api}/zz_generated.conversion.go; \
+		conversion-gen \
+			--go-header-file=$(CURDIR)/hack/boilerplate.generatego.txt \
+			--output-file=zz_generated.conversion.go \
+			$${api}/; \
 	done
 
 setup-e2e:

cmd/dra-example-kubeletplugin/cdi.go

@@ -25,8 +25,6 @@ import (
 	cdiapi "tags.cncf.io/container-device-interface/pkg/cdi"
 	cdiparser "tags.cncf.io/container-device-interface/pkg/parser"
 	cdispec "tags.cncf.io/container-device-interface/specs-go"
-
-	"sigs.k8s.io/dra-example-driver/internal/profiles"
 )
 
 const cdiCommonDeviceName = "common"
@@ -85,7 +83,7 @@ func (cdi *CDIHandler) CreateCommonSpecFile() error {
 	return cdi.cache.WriteSpec(spec, specName)
 }
 
-func (cdi *CDIHandler) CreateClaimSpecFile(claimUID string, devices profiles.PreparedDevices) error {
+func (cdi *CDIHandler) CreateClaimSpecFile(claimUID string, devices PreparedDevices) error {
 	specName := cdiapi.GenerateTransientSpecName(cdi.vendor(), cdi.class, claimUID)
 
 	spec := &cdispec.Spec{

cmd/dra-example-kubeletplugin/checkpoint.go

@@ -1,53 +1,74 @@
+/*
+ * 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/checksum"
-)
+	"errors"
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
 
-type Checkpoint struct {
-	Checksum checksum.Checksum `json:"checksum"`
-	V1       *CheckpointV1     `json:"v1,omitempty"`
-}
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/utils/ptr"
 
-type CheckpointV1 struct {
-	PreparedClaims PreparedClaims `json:"preparedClaims,omitempty"`
-}
+	checkpointapi "sigs.k8s.io/dra-example-driver/internal/api/checkpoint"
+)
 
-func newCheckpoint() *Checkpoint {
-	pc := &Checkpoint{
-		Checksum: 0,
-		V1: &CheckpointV1{
-			PreparedClaims: make(PreparedClaims),
-		},
+// readCheckpoint returns the Checkpoint at the given path in the format
+// expected by the given decoder. If the path doesn't exist, returns an empty
+// Checkpoint and no error.
+func readCheckpoint(path string, decoder runtime.Decoder) (*checkpointapi.Checkpoint, error) {
+	data, err := os.ReadFile(path)
+	if err != nil && !errors.Is(err, fs.ErrNotExist) {
+		return nil, err
+	}
+	checkpoint := new(checkpointapi.Checkpoint)
+	if data != nil {
+		_, _, err = decoder.Decode(data, ptr.To(checkpointapi.SchemeGroupVersion.WithKind("Checkpoint")), checkpoint)
+		if err != nil {
+			return nil, fmt.Errorf("unmarshal JSON from %s: %w", path, err)
+		}
 	}
-	return pc
+	return checkpoint, nil
 }
 
-func (cp *Checkpoint) MarshalCheckpoint() ([]byte, error) {
-	cp.Checksum = 0
-	out, err := json.Marshal(*cp)
+// writeCheckpoint writes checkpoint to the file at path in
+// the format prescribed by encoder. The file is overwritten if it already
+// exists and is created if it does not already exist.
+func writeCheckpoint(path string, encoder runtime.Encoder, checkpoint *checkpointapi.Checkpoint) (err error) {
+	dir := filepath.Dir(path)
+	tmp, err := os.CreateTemp(dir, "tmp-checkpoint-*")
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("create temp file in %s: %w", dir, 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
+		if err1 := tmp.Close(); err1 != nil && err == nil {
+			err = fmt.Errorf("close temp file: %w", err1)
+		}
 	}()
-	out, err := json.Marshal(*cp)
-	if err != nil {
-		return err
+	if err := encoder.Encode(checkpoint, tmp); err != nil {
+		return fmt.Errorf("encode 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 ck.Verify(out)
+	return nil
 }

cmd/dra-example-kubeletplugin/checkpoint_test.go

@@ -0,0 +1,70 @@
+/*
+ * 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 (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/types"
+
+	checkpointapi "sigs.k8s.io/dra-example-driver/internal/api/checkpoint"
+)
+
+func TestReadWriteCheckpointRoundtrip(t *testing.T) {
+	tests := map[string]struct {
+		checkpoint *checkpointapi.Checkpoint
+	}{
+		"new checkpoint": {
+			checkpoint: new(checkpointapi.Checkpoint),
+		},
+		"populated checkpoint": {
+			checkpoint: &checkpointapi.Checkpoint{
+				PreparedClaims: []checkpointapi.PreparedClaim{
+					{UID: types.UID("123")},
+					{UID: types.UID("456")},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			dir := t.TempDir()
+			path := filepath.Join(dir, DriverPluginCheckpointFile)
+
+			decoder, encoder, err := checkpointSerializer()
+			if err != nil {
+				t.Fatal("failed to initialize checkpoint serializer:", err)
+			}
+
+			checkpoint, err := readCheckpoint(path, decoder)
+			assert.NoError(t, err)
+			assert.Equal(t, new(checkpointapi.Checkpoint), checkpoint)
+
+			checkpoint = test.checkpoint
+			err = writeCheckpoint(path, encoder, checkpoint)
+			require.NoError(t, err)
+
+			read, err := readCheckpoint(path, decoder)
+			require.NoError(t, err)
+			assert.Equal(t, test.checkpoint, read)
+		})
+	}
+}

cmd/dra-example-kubeletplugin/driver.go

@@ -132,7 +132,7 @@ func (d *driver) UnprepareResourceClaims(ctx context.Context, claims []kubeletpl
 }
 
 func (d *driver) unprepareResourceClaim(_ context.Context, claim kubeletplugin.NamespacedObject) error {
-	if err := d.state.Unprepare(string(claim.UID)); err != nil {
+	if err := d.state.Unprepare(claim.UID); err != nil {
 		return fmt.Errorf("error unpreparing devices for claim %v: %w", claim.UID, err)
 	}