Merge pull request #947 from wzshiming/feat/compress-yaml

Support compress snapshot yaml

pkg/kwokctl/cmd/snapshot/export/export.go

@@ -94,14 +94,6 @@ func runE(ctx context.Context, flags *flagpole) error {
 		return err
 	}
 
-	f, err := file.Open(flags.Path)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		_ = f.Close()
-	}()
-
 	logger := log.FromContext(ctx)
 
 	pagerConfig := &snapshot.PagerConfig{
@@ -130,7 +122,20 @@ func runE(ctx context.Context, flags *flagpole) error {
 		return err
 	}
 
-	var writer io.Writer = f
+	f, err := file.Open(flags.Path)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = f.Close()
+	}()
+
+	press := file.Compress(flags.Path, f)
+	defer func() {
+		_ = press.Close()
+	}()
+
+	var writer io.Writer = press
 
 	startTime := time.Now()
 	writer = recording.NewWriteHook(writer, func(b []byte) []byte {

pkg/kwokctl/cmd/snapshot/record/record.go

@@ -116,7 +116,12 @@ func runE(ctx context.Context, flags *flagpole) error {
 		_ = f.Close()
 	}()
 
-	var writer io.Writer = f
+	press := file.Compress(flags.Path, f)
+	defer func() {
+		_ = press.Close()
+	}()
+
+	var writer io.Writer = press
 
 	startTime := time.Now()
 	writer = recording.NewWriteHook(writer, func(bytes []byte) []byte {

pkg/kwokctl/cmd/snapshot/replay/replay.go

@@ -148,7 +148,15 @@ func runE(ctx context.Context, flags *flagpole) error {
 		_ = f.Close()
 	}()
 
-	var reader io.Reader = f
+	press, err := file.Decompress(flags.Path, f)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = press.Close()
+	}()
+
+	var reader io.Reader = press
 
 	startTime := time.Now()
 	reader = recording.NewReadHook(reader, func(bytes []byte) []byte {

pkg/utils/file/compress.go

@@ -0,0 +1,74 @@
+/*
+Copyright 2024 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 file
+
+import (
+	"bufio"
+	"compress/gzip"
+	"io"
+	"path/filepath"
+)
+
+// Compress compresses a writer
+func Compress(name string, w io.Writer) io.WriteCloser {
+	ext := filepath.Ext(name)
+	switch ext {
+	case ".gz", ".tgz":
+		z, _ := gzip.NewWriterLevel(w, gzip.BestCompression)
+		return z
+	}
+
+	return struct {
+		io.Writer
+		io.Closer
+	}{
+		Writer: w,
+		Closer: nopCloser{},
+	}
+}
+
+type nopCloser struct{}
+
+func (nopCloser) Close() error { return nil }
+
+const (
+	gzipID1 = 0x1f
+	gzipID2 = 0x8b
+)
+
+// Decompress decompresses a reader
+// It will return a raw reader if the reader is not compressed
+func Decompress(name string, r io.Reader) (io.ReadCloser, error) {
+	ext := filepath.Ext(name)
+	switch ext {
+	case ".gz", ".tgz":
+		bufReader := bufio.NewReader(r)
+
+		prefix, err := bufReader.Peek(2)
+		if err != nil {
+			return nil, err
+		}
+
+		if prefix[0] == gzipID1 && prefix[1] == gzipID2 {
+			return gzip.NewReader(bufReader)
+		}
+
+		return io.NopCloser(bufReader), nil
+	}
+
+	return io.NopCloser(r), nil
+}