Support k8s format for snapshot

go.mod

@@ -21,6 +21,7 @@ require (
 	golang.org/x/sync v0.1.0
 	golang.org/x/sys v0.5.0
 	golang.org/x/term v0.5.0
+	gopkg.in/yaml.v3 v3.0.1
 	k8s.io/api v0.26.0
 	k8s.io/apimachinery v0.26.0
 	k8s.io/apiserver v0.26.0
@@ -82,7 +83,6 @@ require (
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/component-base v0.26.0 // indirect
 	k8s.io/gengo v0.0.0-20220902162205-c0856e24416d // indirect
 	k8s.io/klog/v2 v2.80.1 // indirect

pkg/kwokctl/cmd/snapshot/restore/restore.go

@@ -25,14 +25,16 @@ import (
 
 	"sigs.k8s.io/kwok/pkg/config"
 	"sigs.k8s.io/kwok/pkg/kwokctl/runtime"
+	"sigs.k8s.io/kwok/pkg/kwokctl/snapshot"
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/path"
 )
 
 type flagpole struct {
-	Name   string
-	Path   string
-	Format string
+	Name    string
+	Path    string
+	Format  string
+	Filters []string
 }
 
 // NewCommand returns a new cobra.Command to save the cluster as a snapshot.
@@ -49,7 +51,8 @@ func NewCommand(ctx context.Context) *cobra.Command {
 		},
 	}
 	cmd.Flags().StringVar(&flags.Path, "path", "", "Path to the snapshot")
-	cmd.Flags().StringVar(&flags.Format, "format", "etcd", "Format of the snapshot file (etcd)")
+	cmd.Flags().StringVar(&flags.Format, "format", "etcd", "Format of the snapshot file (etcd, k8s)")
+	cmd.Flags().StringSliceVar(&flags.Filters, "filter", snapshot.Resources, "Filter the resources to restore, only support for k8s format")
 	return cmd
 }
 
@@ -69,12 +72,18 @@ func runE(ctx context.Context, flags *flagpole) error {
 		return err
 	}
 
-	if flags.Format == "etcd" {
+	switch flags.Format {
+	case "etcd":
 		err = rt.SnapshotRestore(ctx, flags.Path)
 		if err != nil {
 			return err
 		}
-	} else {
+	case "k8s":
+		err = rt.SnapshotRestoreWithYAML(ctx, flags.Path, flags.Filters)
+		if err != nil {
+			return err
+		}
+	default:
 		return fmt.Errorf("unsupport format %q", flags.Format)
 	}
 	return nil

pkg/kwokctl/cmd/snapshot/save/save.go

@@ -25,14 +25,16 @@ import (
 
 	"sigs.k8s.io/kwok/pkg/config"
 	"sigs.k8s.io/kwok/pkg/kwokctl/runtime"
+	"sigs.k8s.io/kwok/pkg/kwokctl/snapshot"
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/path"
 )
 
 type flagpole struct {
-	Name   string
-	Path   string
-	Format string
+	Name    string
+	Path    string
+	Format  string
+	Filters []string
 }
 
 // NewCommand returns a new cobra.Command for cluster snapshotting.
@@ -49,7 +51,8 @@ func NewCommand(ctx context.Context) *cobra.Command {
 		},
 	}
 	cmd.Flags().StringVar(&flags.Path, "path", "", "Path to the snapshot")
-	cmd.Flags().StringVar(&flags.Format, "format", "etcd", "Format of the snapshot file (etcd)")
+	cmd.Flags().StringVar(&flags.Format, "format", "etcd", "Format of the snapshot file (etcd, k8s)")
+	cmd.Flags().StringSliceVar(&flags.Filters, "filter", snapshot.Resources, "Filter the resources to save, only support for k8s format")
 	return cmd
 }
 
@@ -69,12 +72,18 @@ func runE(ctx context.Context, flags *flagpole) error {
 		return err
 	}
 
-	if flags.Format == "etcd" {
+	switch flags.Format {
+	case "etcd":
 		err = rt.SnapshotSave(ctx, flags.Path)
 		if err != nil {
 			return err
 		}
-	} else {
+	case "k8s":
+		err = rt.SnapshotSaveWithYAML(ctx, flags.Path, flags.Filters)
+		if err != nil {
+			return err
+		}
+	default:
 		return fmt.Errorf("unsupport format %q", flags.Format)
 	}
 	return nil

pkg/kwokctl/runtime/binary/cluster_snapshot.go

@@ -100,3 +100,33 @@ func (c *Cluster) SnapshotRestore(ctx context.Context, path string) error {
 	}
 	return nil
 }
+
+// SnapshotSaveWithYAML save the snapshot of cluster
+func (c *Cluster) SnapshotSaveWithYAML(ctx context.Context, path string, filters []string) error {
+	err := c.Cluster.SnapshotSaveWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// SnapshotRestoreWithYAML restore the snapshot of cluster
+func (c *Cluster) SnapshotRestoreWithYAML(ctx context.Context, path string, filters []string) error {
+	logger := log.FromContext(ctx)
+	err := c.StopComponent(ctx, "kube-controller-manager")
+	if err != nil {
+		logger.Error("Failed to stop kube-controller-manager", err)
+	}
+	defer func() {
+		err = c.StartComponent(ctx, "kube-controller-manager")
+		if err != nil {
+			logger.Error("Failed to start kube-controller-manager", err)
+		}
+	}()
+
+	err = c.Cluster.SnapshotRestoreWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/kwokctl/runtime/cluster_snapshot.go

@@ -0,0 +1,54 @@
+/*
+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 runtime
+
+import (
+	"bytes"
+	"context"
+	"os"
+
+	"sigs.k8s.io/kwok/pkg/kwokctl/snapshot"
+)
+
+// SnapshotSaveWithYAML save the snapshot of cluster
+func (c *Cluster) SnapshotSaveWithYAML(ctx context.Context, path string, filters []string) error {
+	file, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = file.Close()
+	}()
+	err = snapshot.Save(ctx, c, file, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// SnapshotRestoreWithYAML restore the snapshot of cluster
+func (c *Cluster) SnapshotRestoreWithYAML(ctx context.Context, path string, filters []string) error {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return err
+	}
+	err = snapshot.Load(ctx, c, bytes.NewBuffer(data), filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/kwokctl/runtime/compose/cluster_snapshot.go

@@ -138,3 +138,33 @@ func (c *Cluster) SnapshotRestore(ctx context.Context, path string) error {
 
 	return nil
 }
+
+// SnapshotSaveWithYAML save the snapshot of cluster
+func (c *Cluster) SnapshotSaveWithYAML(ctx context.Context, path string, filters []string) error {
+	err := c.Cluster.SnapshotSaveWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// SnapshotRestoreWithYAML restore the snapshot of cluster
+func (c *Cluster) SnapshotRestoreWithYAML(ctx context.Context, path string, filters []string) error {
+	logger := log.FromContext(ctx)
+	err := c.StopComponent(ctx, "kube-controller-manager")
+	if err != nil {
+		logger.Error("Failed to stop kube-controller-manager", err)
+	}
+	defer func() {
+		err = c.StartComponent(ctx, "kube-controller-manager")
+		if err != nil {
+			logger.Error("Failed to start kube-controller-manager", err)
+		}
+	}()
+
+	err = c.Cluster.SnapshotRestoreWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/kwokctl/runtime/config.go

@@ -106,4 +106,10 @@ type Runtime interface {
 
 	// SnapshotRestore restore the snapshot of cluster
 	SnapshotRestore(ctx context.Context, path string) error
+
+	// SnapshotSaveWithYAML save the snapshot of cluster
+	SnapshotSaveWithYAML(ctx context.Context, path string, filters []string) error
+
+	// SnapshotRestoreWithYAML restore the snapshot of cluster
+	SnapshotRestoreWithYAML(ctx context.Context, path string, filters []string) error
 }

pkg/kwokctl/runtime/kind/cluster_snapshot.go

@@ -19,6 +19,7 @@ package kind
 import (
 	"context"
 	"os"
+	"time"
 
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/exec"
@@ -107,3 +108,40 @@ func (c *Cluster) SnapshotRestore(ctx context.Context, path string) error {
 
 	return nil
 }
+
+// SnapshotSaveWithYAML save the snapshot of cluster
+func (c *Cluster) SnapshotSaveWithYAML(ctx context.Context, path string, filters []string) error {
+	err := c.Cluster.SnapshotSaveWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// SnapshotRestoreWithYAML restore the snapshot of cluster
+func (c *Cluster) SnapshotRestoreWithYAML(ctx context.Context, path string, filters []string) error {
+	logger := log.FromContext(ctx)
+	err := c.StopComponent(ctx, "kube-controller-manager")
+	if err != nil {
+		logger.Error("Failed to stop kube-controller-manager", err)
+	}
+	defer func() {
+		err = c.StartComponent(ctx, "kube-controller-manager")
+		if err != nil {
+			logger.Error("Failed to start kube-controller-manager", err)
+		}
+	}()
+
+	// Since `kube-controller-manager` is stopping,
+	// we have to wait for it to finish.
+	err = c.WaitReady(ctx, 30*time.Second)
+	if err != nil {
+		logger.Error("Failed to wait cluster ready", err)
+	}
+
+	err = c.Cluster.SnapshotRestoreWithYAML(ctx, path, filters)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/kwokctl/snapshot/doc.go

@@ -0,0 +1,18 @@
+/*
+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 snapshot is the snapshot of cluster
+package snapshot