From bb4c31e97f730629024bc48c60403f0177748860 Mon Sep 17 00:00:00 2001
From: Bison Xu <1251604436@qq.com>
Date: Fri, 28 Oct 2022 11:22:58 +0800
Subject: [PATCH] add zstd compress (#327)

zstd compress

Co-authored-by: xubaisheng <>
---
 go.mod                               |  1 +
 go.sum                               |  2 ++
 pkg/compressor/zstd_compress.go      | 26 +++++++++++++++
 pkg/compressor/zstd_compress_test.go | 48 ++++++++++++++++++++++++++++
 4 files changed, 77 insertions(+)
 create mode 100644 pkg/compressor/zstd_compress_test.go

diff --git a/go.mod b/go.mod
index e756638c..7867ac89 100644
--- a/go.mod
+++ b/go.mod
@@ -86,6 +86,7 @@ require (
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/k0kubun/pp v3.0.1+incompatible // indirect
+	github.com/klauspost/compress v1.15.11
 	github.com/knadh/koanf v1.4.3 // indirect
 	github.com/leodido/go-urn v1.2.1 // indirect
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
diff --git a/go.sum b/go.sum
index 42e93fcf..08702f50 100644
--- a/go.sum
+++ b/go.sum
@@ -538,6 +538,8 @@ github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvW
 github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/klauspost/compress v1.15.11 h1:Lcadnb3RKGin4FYM/orgq0qde+nc15E5Cbqg4B9Sx9c=
+github.com/klauspost/compress v1.15.11/go.mod h1:QPwzmACJjUTFsnSHH934V6woptycfrDDJnH7hvFVbGM=
 github.com/knadh/koanf v1.4.3 h1:rSJcSH5LSFhvzBRsAYfT3k7eLP0I4UxeZqjtAatk+wc=
 github.com/knadh/koanf v1.4.3/go.mod h1:5FAkuykKXZvLqhAbP4peWgM5CTcZmn7L1d27k/a+kfg=
 github.com/koding/multiconfig v0.0.0-20171124222453-69c27309b2d7/go.mod h1:Y2SaZf2Rzd0pXkLVhLlCiAXFCLSXAIbTKDivVgff/AM=
diff --git a/pkg/compressor/zstd_compress.go b/pkg/compressor/zstd_compress.go
index d1eea987..715d0cc0 100644
--- a/pkg/compressor/zstd_compress.go
+++ b/pkg/compressor/zstd_compress.go
@@ -16,3 +16,29 @@
  */
 
 package compressor
+
+import (
+	"github.com/klauspost/compress/zstd"
+)
+
+type Zstd struct{}
+
+func (z Zstd) Compress(data []byte) ([]byte, error) {
+	var encoder, err = zstd.NewWriter(nil)
+	if err != nil {
+		return nil, err
+	}
+	return encoder.EncodeAll(data, make([]byte, 0, len(data))), nil
+}
+
+func (z Zstd) Decompress(data []byte) ([]byte, error) {
+	var decoder, err = zstd.NewReader(nil)
+	if err != nil {
+		return nil, err
+	}
+	return decoder.DecodeAll(data, nil)
+}
+
+func (z Zstd) GetCompressorType() CompressorType {
+	return CompressorZstd
+}
diff --git a/pkg/compressor/zstd_compress_test.go b/pkg/compressor/zstd_compress_test.go
new file mode 100644
index 00000000..1ae32611
--- /dev/null
+++ b/pkg/compressor/zstd_compress_test.go
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 compressor
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestZstdCompress(t *testing.T) {
+	ts := []struct {
+		text string
+	}{
+		{
+			text: strings.Repeat("Don't communicate by sharing memory, share memory by communicating.", 1000),
+		},
+		{
+			text: "88888msj0*&^^%$$#$@!~jjdjfjdlfjkhhdh//><|}{{|\"",
+		},
+	}
+
+	dc := &Zstd{}
+	assert.EqualValues(t, CompressorZstd, dc.GetCompressorType())
+
+	for _, s := range ts {
+		var data = []byte(s.text)
+		dataCompressed, _ := dc.Compress(data)
+		ret, _ := dc.Decompress(dataCompressed)
+		assert.EqualValues(t, s.text, string(ret))
+	}
+}