feat(compress): add snappy (openGemini#808)

Signed-off-by: xkx <triumph_9431@qq.com>

Author: xkx9431

lib/compress/compression_pool.go

@@ -0,0 +1,104 @@
+// Copyright 2024 openGemini 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 compress
+
+import (
+	"compress/gzip"
+	"io"
+	"sync"
+
+	"github.com/golang/snappy"
+	"github.com/klauspost/compress/zstd"
+	"github.com/openGemini/openGemini/lib/cpu"
+	"github.com/openGemini/openGemini/lib/pool"
+)
+
+// #region Gzip Writer Pool
+var gzipWriterPool *pool.FixedPoolV2[*gzip.Writer] = pool.NewFixedPoolV2[*gzip.Writer](func() *gzip.Writer {
+	return gzip.NewWriter(nil)
+}, cpu.GetCpuNum()*2)
+
+func GetGzipWriter(w io.Writer) *gzip.Writer {
+	gz := gzipWriterPool.Get()
+	gz.Reset(w)
+	return gz
+}
+
+func PutGzipWriter(gz *gzip.Writer) {
+	gz.Close()
+	gzipWriterPool.Put(gz)
+}
+
+// #endregion
+
+// #region Gzip Reader Pool
+var gzipReaderPool sync.Pool
+
+func GetGzipReader(r io.Reader) (*gzip.Reader, error) {
+	v := gzipReaderPool.Get()
+	if v == nil {
+		return gzip.NewReader(r)
+	}
+	zr := v.(*gzip.Reader)
+	if err := zr.Reset(r); err != nil {
+		return nil, err
+	}
+	return zr, nil
+}
+
+// PutGzipReader returns back gzip reader obtained via GetGzipReader.
+func PutGzipReader(zr *gzip.Reader) {
+	_ = zr.Close()
+	gzipReaderPool.Put(zr)
+}
+
+// #endregion
+
+// #region Zstd Writer Pool
+var zstdWriterPool *pool.FixedPoolV2[*zstd.Encoder] = pool.NewFixedPoolV2[*zstd.Encoder](func() *zstd.Encoder {
+	encoder, _ := zstd.NewWriter(nil)
+	return encoder
+}, cpu.GetCpuNum()*2)
+
+func GetZstdWriter(w io.Writer) *zstd.Encoder {
+	zstdEncoder := zstdWriterPool.Get()
+	zstdEncoder.Reset(w)
+	return zstdEncoder
+}
+
+func PutZstdWriter(zstdEncoder *zstd.Encoder) {
+	zstdEncoder.Close()
+	zstdWriterPool.Put(zstdEncoder)
+}
+
+// #endregion
+
+// #region Zstd Reader Pool
+var snappyWriterPool = pool.NewFixedPoolV2[*snappy.Writer](func() *snappy.Writer {
+	return snappy.NewBufferedWriter(nil)
+}, cpu.GetCpuNum()*2)
+
+func GetSnappyWriter(w io.Writer) *snappy.Writer {
+	snappyWriter := snappyWriterPool.Get()
+	snappyWriter.Reset(w)
+	return snappyWriter
+}
+
+func PutSnappyWriter(snappyWriter *snappy.Writer) {
+	snappyWriter.Close()
+	snappyWriterPool.Put(snappyWriter)
+}
+
+// #endregion

lib/compress/compression_pool_test.go

@@ -0,0 +1,74 @@
+// Copyright 2024 openGemini 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 compress
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+	"testing"
+
+	"github.com/golang/snappy"
+	"github.com/klauspost/compress/zstd"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGzipWriterPool(t *testing.T) {
+	var buf bytes.Buffer
+	writer := GetGzipWriter(&buf)
+	_, err := writer.Write([]byte("test data"))
+	assert.NoError(t, err)
+	PutGzipWriter(writer)
+
+	reader, err := gzip.NewReader(&buf)
+	assert.NoError(t, err)
+	defer reader.Close()
+
+	result := new(bytes.Buffer)
+	_, err = io.Copy(result, reader)
+	assert.NoError(t, err)
+	assert.Equal(t, "test data", result.String())
+}
+
+func TestZstdWriterPool(t *testing.T) {
+	var buf bytes.Buffer
+	writer := GetZstdWriter(&buf)
+	_, err := writer.Write([]byte("test data"))
+	assert.NoError(t, err)
+	PutZstdWriter(writer)
+
+	reader, err := zstd.NewReader(&buf)
+	assert.NoError(t, err)
+	defer reader.Close()
+
+	result := new(bytes.Buffer)
+	_, err = io.Copy(result, reader)
+	assert.NoError(t, err)
+	assert.Equal(t, "test data", result.String())
+}
+
+func TestSnappyWriterPool(t *testing.T) {
+	var buf bytes.Buffer
+	writer := GetSnappyWriter(&buf)
+	_, err := writer.Write([]byte("test data"))
+	assert.NoError(t, err)
+	PutSnappyWriter(writer)
+
+	reader := snappy.NewReader(&buf)
+	result := new(bytes.Buffer)
+	_, err = io.Copy(result, reader)
+	assert.NoError(t, err)
+	assert.Equal(t, "test data", result.String())
+}

lib/pool/pool.go

@@ -191,6 +191,44 @@ func (p *FixedCachePool) Put(v interface{}) {
 	}
 }
 
+// FixedPoolV2 is a fixed size pool with object generic type information
+type FixedPoolV2[T any] struct {
+	pool chan T
+	new  func() T
+}
+
+func NewFixedPoolV2[T any](new func() T, size int) *FixedPoolV2[T] {
+	return &FixedPoolV2[T]{
+		pool: make(chan T, size),
+		new:  new,
+	}
+}
+
+func (p *FixedPoolV2[T]) Get() T {
+	select {
+	case item := <-p.pool:
+		return item
+	default:
+		return p.new()
+	}
+}
+
+func (p *FixedPoolV2[T]) Put(item T) {
+	select {
+	case p.pool <- item:
+	default:
+	}
+}
+
+func (p *FixedPoolV2[T]) Reset(size int, new func() T) {
+	p.pool = make(chan T, size)
+	p.new = new
+}
+
+func (p *FixedPoolV2[T]) Len() int {
+	return len(p.pool)
+}
+
 var intSlicePool sync.Pool
 
 func GetIntSlice(size int) []int {

lib/pool/pool_test.go

@@ -73,3 +73,57 @@ func TestChunkMetaBuffer(t *testing.T) {
 	buf = pool.GetChunkMetaBuffer()
 	require.Equal(t, 10, len(buf.B))
 }
+
+func TestNewFixedPoolV2(t *testing.T) {
+	newFunc := func() int {
+		return 42
+	}
+	testPool := pool.NewFixedPoolV2(newFunc, 2)
+
+	require.NotNil(t, testPool)
+}
+
+func TestFixedPoolV2_Get(t *testing.T) {
+	newFunc := func() int {
+		return 42
+	}
+	pool := pool.NewFixedPoolV2(newFunc, 2)
+
+	// Test getting a new item when pool is empty
+	item := pool.Get()
+	require.Equal(t, 42, item)
+
+	// Test getting an item from the pool
+	pool.Put(100)
+	item = pool.Get()
+	require.Equal(t, 100, item)
+}
+
+func TestFixedPoolV2_Put(t *testing.T) {
+	newFunc := func() int {
+		return 42
+	}
+	pool := pool.NewFixedPoolV2(newFunc, 2)
+
+	// Test putting an item into the pool
+	pool.Put(100)
+	require.Equal(t, 1, pool.Len())
+
+	// Test putting an item into a full pool
+	pool.Put(200)
+	pool.Put(300)
+	require.Equal(t, 2, pool.Len())
+}
+
+func TestFixedPoolV2_Reset(t *testing.T) {
+	newFunc := func() int {
+		return 42
+	}
+	pool := pool.NewFixedPoolV2(newFunc, 2)
+
+	// Test resetting the pool
+	pool.Put(100)
+	pool.Put(200)
+	pool.Reset(2, newFunc)
+	require.Equal(t, 0, pool.Len())
+}

lib/util/lifted/influx/httpd/handler.go

@@ -27,6 +27,7 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/openGemini/openGemini/app"
 	"github.com/openGemini/openGemini/engine/hybridqp"
+	compression "github.com/openGemini/openGemini/lib/compress"
 	config2 "github.com/openGemini/openGemini/lib/config"
 	"github.com/openGemini/openGemini/lib/cpu"
 	"github.com/openGemini/openGemini/lib/errno"
@@ -1318,15 +1319,15 @@ func (h *Handler) serveWrite(w http.ResponseWriter, r *http.Request, user meta2.
 
 	// Handle gzip decoding of the body
 	if r.Header.Get("Content-Encoding") == "gzip" {
-		b, err := GetGzipReader(r.Body)
+		b, err := compression.GetGzipReader(r.Body)
 		if err != nil {
 			h.httpError(w, err.Error(), http.StatusBadRequest)
 			error := errno.NewError(errno.HttpBadRequest)
 			h.Logger.Error("write error:Handle gzip decoding of the body err", zap.Error(error), zap.String("db", database))
 			atomic.AddInt64(&statistics.HandlerStat.Write400ErrRequests, 1)
 			return
 		}
-		defer PutGzipReader(b)
+		defer compression.PutGzipReader(b)
 		body = b
 	}