Merge pull request #1663 from rancher/zstd-patch-engine

Add new zstd dictionary-based patch engine

pkg/config/v1beta1/gateway_config.go

@@ -223,6 +223,7 @@ const (
 
 const (
 	PatchEngineBsdiff PatchEngine = "bsdiff"
+	PatchEngineZstd   PatchEngine = "zstd"
 )
 
 type BinaryPluginsSpec struct {

pkg/update/patch/patcher.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/gabstv/go-bsdiff/pkg/bsdiff"
 	"github.com/gabstv/go-bsdiff/pkg/bspatch"
+	"github.com/klauspost/compress/zstd"
 )
 
 const (
@@ -39,6 +40,7 @@ func (BsdiffPatcher) CheckFormat(reader io.ReaderAt) bool {
 
 var allPatchEngines = []BinaryPatcher{
 	BsdiffPatcher{},
+	ZstdPatcher{},
 }
 
 func NewPatcherFromFormat(reader io.ReaderAt) (BinaryPatcher, bool) {
@@ -49,3 +51,63 @@ func NewPatcherFromFormat(reader io.ReaderAt) (BinaryPatcher, bool) {
 	}
 	return nil, false
 }
+
+type ZstdPatcher struct{}
+
+func (ZstdPatcher) GeneratePatch(old io.Reader, new io.Reader, patchOut io.Writer) (err error) {
+	dict, err := io.ReadAll(old)
+	if err != nil {
+		return err
+	}
+	enc, err := zstd.NewWriter(nil,
+		zstd.WithEncoderDictRaw(0, dict),
+		zstd.WithWindowSize(zstd.MaxWindowSize),
+		zstd.WithEncoderLevel(zstd.SpeedBestCompression),
+	)
+	defer enc.Close()
+	if err != nil {
+		return err
+	}
+	newBytes, err := io.ReadAll(new)
+	if err != nil {
+		return err
+	}
+	patch := enc.EncodeAll(newBytes, nil)
+	patchOut.Write(patch)
+	return nil
+}
+
+func (ZstdPatcher) ApplyPatch(old io.Reader, patch io.Reader, newOut io.Writer) (err error) {
+	dict, err := io.ReadAll(old)
+	if err != nil {
+		return err
+	}
+	patchBytes, err := io.ReadAll(patch)
+	if err != nil {
+		return err
+	}
+
+	dec, err := zstd.NewReader(nil,
+		zstd.WithDecoderDictRaw(0, dict),
+		zstd.WithDecoderLowmem(true),
+	)
+	if err != nil {
+		return err
+	}
+	defer dec.Close()
+	out, err := dec.DecodeAll(patchBytes, nil)
+	if err != nil {
+		return err
+	}
+	newOut.Write(out)
+	return nil
+}
+
+func (ZstdPatcher) CheckFormat(reader io.ReaderAt) bool {
+	header := "\x28\xb5\x2f\xfd"
+	buf := make([]byte, len(header))
+	if _, err := reader.ReadAt(buf, 0); err == nil {
+		return bytes.Equal(buf, []byte(header))
+	}
+	return false
+}

pkg/update/patch/server/server.go

@@ -72,6 +72,8 @@ func NewFilesystemPluginSyncServer(
 	switch cfg.Binary.Cache.PatchEngine {
 	case v1beta1.PatchEngineBsdiff:
 		patchEngine = patch.BsdiffPatcher{}
+	case v1beta1.PatchEngineZstd:
+		patchEngine = patch.ZstdPatcher{}
 	default:
 		return nil, fmt.Errorf("unknown patch engine: %s", cfg.Binary.Cache.PatchEngine)
 	}

pkg/update/patch/server/server_suite_test.go

@@ -12,6 +12,7 @@ import (
 	. "github.com/onsi/gomega"
 	"github.com/onsi/gomega/gexec"
 	controlv1 "github.com/rancher/opni/pkg/apis/control/v1"
+	"github.com/rancher/opni/pkg/config/v1beta1"
 	_ "github.com/rancher/opni/pkg/test/setup"
 	"github.com/rancher/opni/pkg/test/testutil"
 	"github.com/rancher/opni/pkg/update/patch"
@@ -41,8 +42,15 @@ var (
 		"test2": {},
 	}
 
-	test1v1tov2Patch = new(bytes.Buffer)
-	test2v1tov2Patch = new(bytes.Buffer)
+	test1v1tov2Patch = map[v1beta1.PatchEngine]*bytes.Buffer{
+		"bsdiff": new(bytes.Buffer),
+		"zstd":   new(bytes.Buffer),
+	}
+
+	test2v1tov2Patch = map[v1beta1.PatchEngine]*bytes.Buffer{
+		"bsdiff": new(bytes.Buffer),
+		"zstd":   new(bytes.Buffer),
+	}
 )
 
 var osfs = afero.Afero{Fs: afero.NewOsFs()}
@@ -113,24 +121,30 @@ var _ = BeforeSuite(func() {
 	})
 	Expect(eg.Wait()).To(Succeed())
 
-	patcher := patch.BsdiffPatcher{}
+	patchers := map[v1beta1.PatchEngine]patch.BinaryPatcher{
+		v1beta1.PatchEngineBsdiff: patch.BsdiffPatcher{},
+		v1beta1.PatchEngineZstd:   patch.ZstdPatcher{},
+	}
 	eg = errgroup.Group{}
-
-	eg.Go(func() error {
-		return patcher.GeneratePatch(
-			bytes.NewReader(testBinaries["test1"]["v1"]),
-			bytes.NewReader(testBinaries["test1"]["v2"]),
-			test1v1tov2Patch,
-		)
-	})
-
-	eg.Go(func() error {
-		return patcher.GeneratePatch(
-			bytes.NewReader(testBinaries["test2"]["v1"]),
-			bytes.NewReader(testBinaries["test2"]["v2"]),
-			test2v1tov2Patch,
-		)
-	})
+	for name, patcher := range patchers {
+		name, patcher := name, patcher
+
+		eg.Go(func() error {
+			return patcher.GeneratePatch(
+				bytes.NewReader(testBinaries["test1"]["v1"]),
+				bytes.NewReader(testBinaries["test1"]["v2"]),
+				test1v1tov2Patch[name],
+			)
+		})
+
+		eg.Go(func() error {
+			return patcher.GeneratePatch(
+				bytes.NewReader(testBinaries["test2"]["v1"]),
+				bytes.NewReader(testBinaries["test2"]["v2"]),
+				test2v1tov2Patch[name],
+			)
+		})
+	}
 
 	Expect(eg.Wait()).To(Succeed())