Pool for decoders

CHANGELOG.md

@@ -1,3 +1,7 @@
+
+* Supported pool of decoders
+
+
 ## v3.101.0
 * Added `table.Client.ReadRows` method with internal retries
 

internal/topic/topicreadercommon/decoders.go

@@ -5,43 +5,122 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
+type ReadResetter interface {
+	io.Reader
+	Reset(r io.Reader) error
+}
+
+type decoderPool struct {
+	pool sync.Pool
+}
+
+func (p *decoderPool) Get() ReadResetter {
+	dec, _ := p.pool.Get().(ReadResetter)
+
+	return dec
+}
+
+func (p *decoderPool) Put(dec ReadResetter) {
+	p.pool.Put(dec)
+}
+
+func newDecoderPool() *decoderPool {
+	return &decoderPool{
+		pool: sync.Pool{},
+	}
+}
+
 type DecoderMap struct {
-	m map[rawtopiccommon.Codec]PublicCreateDecoderFunc
+	m  map[rawtopiccommon.Codec]PublicCreateDecoderFunc
+	dp map[rawtopiccommon.Codec]*decoderPool
 }
 
 func NewDecoderMap() DecoderMap {
-	return DecoderMap{
-		m: map[rawtopiccommon.Codec]PublicCreateDecoderFunc{
-			rawtopiccommon.CodecRaw: func(input io.Reader) (io.Reader, error) {
-				return input, nil
-			},
-			rawtopiccommon.CodecGzip: func(input io.Reader) (io.Reader, error) {
-				return gzip.NewReader(input)
-			},
-		},
+	dm := DecoderMap{
+		m:  make(map[rawtopiccommon.Codec]PublicCreateDecoderFunc),
+		dp: make(map[rawtopiccommon.Codec]*decoderPool),
 	}
+
+	dm.AddDecoder(rawtopiccommon.CodecRaw, func(input io.Reader) (ReadResetter, error) {
+		return &nopResetter{Reader: input}, nil
+	})
+
+	dm.AddDecoder(rawtopiccommon.CodecGzip, func(input io.Reader) (ReadResetter, error) {
+		gz, err := gzip.NewReader(input)
+		if err != nil {
+			return nil, err
+		}
+
+		return gz, nil
+	})
+
+	return dm
 }
 
-func (m *DecoderMap) AddDecoder(codec rawtopiccommon.Codec, createFunc PublicCreateDecoderFunc) {
-	m.m[codec] = createFunc
+type pooledDecoder struct {
+	ReadResetter
+	pool *decoderPool
+}
+
+func (p *pooledDecoder) Close() error {
+	if closer, ok := p.ReadResetter.(io.Closer); ok {
+		closer.Close()
+	}
+	p.pool.Put(p.ReadResetter)
+
+	return nil
 }
 
 func (m *DecoderMap) Decode(codec rawtopiccommon.Codec, input io.Reader) (io.Reader, error) {
-	if f := m.m[codec]; f != nil {
-		return f(input)
+	createFunc, ok := m.m[codec]
+	if !ok {
+		return nil, xerrors.WithStackTrace(xerrors.Wrap(
+			fmt.Errorf("ydb: failed decompress message with codec %v: %w", codec, ErrPublicUnexpectedCodec),
+		))
+	}
+
+	pool := m.dp[codec]
+	decoder := pool.Get()
+	if decoder == nil {
+		var err error
+		decoder, err = createFunc(input)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		if err := decoder.Reset(input); err != nil {
+			return nil, err
+		}
 	}
 
-	return nil, xerrors.WithStackTrace(xerrors.Wrap(
-		fmt.Errorf("ydb: failed decompress message with codec %v: %w", codec, ErrPublicUnexpectedCodec),
-	))
+	return &pooledDecoder{
+		ReadResetter: decoder,
+		pool:         pool,
+	}, nil
+}
+
+func (m *DecoderMap) AddDecoder(codec rawtopiccommon.Codec, createFunc func(io.Reader) (ReadResetter, error)) {
+	m.m[codec] = createFunc
+	m.dp[codec] = newDecoderPool()
+}
+
+type nopResetter struct {
+	io.Reader
+}
+
+func (n *nopResetter) Reset(r io.Reader) error {
+	n.Reader = r
+
+	return nil
 }
 
-type PublicCreateDecoderFunc func(input io.Reader) (io.Reader, error)
+type PublicCreateDecoderFunc func(input io.Reader) (ReadResetter, error)
 
 // ErrPublicUnexpectedCodec return when try to read message content with unknown codec
 var ErrPublicUnexpectedCodec = xerrors.Wrap(errors.New("ydb: unexpected codec"))

internal/topic/topicreadercommon/decoders_test.go

@@ -0,0 +1,90 @@
+package topicreadercommon
+
+import (
+	"bytes"
+	"compress/gzip"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
+)
+
+func TestDecoderMap(t *testing.T) {
+	decoderMap := NewDecoderMap()
+
+	t.Run("DecodeRaw", func(t *testing.T) {
+		data := []byte("test data")
+		reader := bytes.NewReader(data)
+
+		decodedReader, err := decoderMap.Decode(rawtopiccommon.CodecRaw, reader)
+		require.NoError(t, err)
+
+		result, err := io.ReadAll(decodedReader)
+		require.NoError(t, err)
+		require.Equal(t, data, result)
+	})
+
+	t.Run("DecodeGzip", func(t *testing.T) {
+		data := []byte("test data")
+		var buf bytes.Buffer
+		gzipWriter := gzip.NewWriter(&buf)
+		_, err := gzipWriter.Write(data)
+		require.NoError(t, err)
+		require.NoError(t, gzipWriter.Close())
+
+		decodedReader, err := decoderMap.Decode(rawtopiccommon.CodecGzip, &buf)
+		require.NoError(t, err)
+
+		result, err := io.ReadAll(decodedReader)
+		require.NoError(t, err)
+		require.Equal(t, data, result)
+	})
+
+	t.Run("DecodeUnknownCodec", func(t *testing.T) {
+		_, err := decoderMap.Decode(rawtopiccommon.Codec(999), bytes.NewReader([]byte{}))
+		require.Error(t, err)
+		require.True(t, errors.Is(err, ErrPublicUnexpectedCodec))
+	})
+
+	t.Run("DecodeCustomCodec", func(t *testing.T) {
+		dm := NewDecoderMap()
+		customCodec := rawtopiccommon.Codec(1001)
+		dm.AddDecoder(customCodec, func(input io.Reader) (ReadResetter, error) {
+			return gzip.NewReader(input)
+		})
+		require.Len(t, dm.dp, 3)
+
+		data := []byte("custom test data")
+		var buf bytes.Buffer
+		gzipWriter := gzip.NewWriter(&buf)
+		_, err := gzipWriter.Write(data)
+		require.NoError(t, err)
+		require.NoError(t, gzipWriter.Close())
+
+		decodedReader, err := dm.Decode(customCodec, &buf)
+		require.NoError(t, err)
+		defer decodedReader.(io.Closer).Close()
+
+		result, err := io.ReadAll(decodedReader)
+		require.NoError(t, err)
+		require.Equal(t, string(data), string(result))
+
+		data2 := []byte("second test data")
+		var buf2 bytes.Buffer
+		gzipWriter2 := gzip.NewWriter(&buf2)
+		_, err = gzipWriter2.Write(data2)
+		require.NoError(t, err)
+		require.NoError(t, gzipWriter2.Close())
+
+		decodedReader2, err := dm.Decode(customCodec, &buf2)
+		require.NoError(t, err)
+		defer decodedReader2.(io.Closer).Close()
+
+		result2, err := io.ReadAll(decodedReader2)
+		require.NoError(t, err)
+		require.Equal(t, string(data2), string(result2))
+	})
+}