Skip to content

Commit 84893f0

Browse files
rekbyrekby
authored
Merge pull request #1584 Pool for encoders from DenisEvd/pool-for-encoders
2 parents 0dc77f5 + c1df891 commit 84893f0

File tree

8 files changed

+208
-49
lines changed

8 files changed

+208
-49
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
* Supported pool of encoders, which implement ResetableWriter interface
2+
13
## v3.97.0
24
* Added immutable range iterators from go1.23 into query stats to iterate over query phases and accessed tables without query stats object mutation
35

‎internal/topic/topicwriterinternal/encoders.go‎

Lines changed: 95 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package topicwriterinternal
22

33
import (
4+
"bytes"
45
"compress/gzip"
56
"fmt"
67
"io"
@@ -17,36 +18,115 @@ const (
1718
codecUnknown = rawtopiccommon.CodecUNSPECIFIED
1819
)
1920

20-
type EncoderMap struct {
21+
type PublicResetableWriter interface {
22+
io.WriteCloser
23+
Reset(wr io.Writer)
24+
}
25+
26+
type encoderPool struct {
27+
pool sync.Pool
28+
}
29+
30+
func (p *encoderPool) Get() PublicResetableWriter {
31+
enc, _ := p.pool.Get().(PublicResetableWriter)
32+
33+
return enc
34+
}
35+
36+
func (p *encoderPool) Put(wr PublicResetableWriter) {
37+
p.pool.Put(wr)
38+
}
39+
40+
func newEncoderPool() *encoderPool {
41+
return &encoderPool{
42+
pool: sync.Pool{},
43+
}
44+
}
45+
46+
type MultiEncoder struct {
2147
m map[rawtopiccommon.Codec]PublicCreateEncoderFunc
48+
49+
bp xsync.Pool[bytes.Buffer]
50+
ep map[rawtopiccommon.Codec]*encoderPool
2251
}
2352

24-
func NewEncoderMap() *EncoderMap {
25-
return &EncoderMap{
26-
m: map[rawtopiccommon.Codec]PublicCreateEncoderFunc{
27-
rawtopiccommon.CodecRaw: func(writer io.Writer) (io.WriteCloser, error) {
28-
return nopWriteCloser{writer}, nil
29-
},
30-
rawtopiccommon.CodecGzip: func(writer io.Writer) (io.WriteCloser, error) {
31-
return gzip.NewWriter(writer), nil
53+
func NewMultiEncoder() *MultiEncoder {
54+
me := &MultiEncoder{
55+
m: make(map[rawtopiccommon.Codec]PublicCreateEncoderFunc),
56+
bp: xsync.Pool[bytes.Buffer]{
57+
New: func() *bytes.Buffer {
58+
return &bytes.Buffer{}
3259
},
3360
},
61+
ep: make(map[rawtopiccommon.Codec]*encoderPool),
3462
}
63+
64+
me.AddEncoder(rawtopiccommon.CodecRaw, func(writer io.Writer) (io.WriteCloser, error) {
65+
return nopWriteCloser{writer}, nil
66+
})
67+
me.AddEncoder(rawtopiccommon.CodecGzip, func(writer io.Writer) (io.WriteCloser, error) {
68+
return gzip.NewWriter(writer), nil
69+
})
70+
71+
return me
3572
}
3673

37-
func (e *EncoderMap) AddEncoder(codec rawtopiccommon.Codec, creator PublicCreateEncoderFunc) {
74+
func (e *MultiEncoder) AddEncoder(codec rawtopiccommon.Codec, creator PublicCreateEncoderFunc) {
3875
e.m[codec] = creator
76+
e.ep[codec] = newEncoderPool()
3977
}
4078

41-
func (e *EncoderMap) CreateLazyEncodeWriter(codec rawtopiccommon.Codec, target io.Writer) (io.WriteCloser, error) {
79+
func (e *MultiEncoder) Encode(codec rawtopiccommon.Codec, target io.Writer, source io.Reader) (int, error) {
80+
buf := e.bp.GetOrNew()
81+
defer e.bp.Put(buf)
82+
83+
buf.Reset()
84+
if _, err := buf.ReadFrom(source); err != nil {
85+
return 0, err
86+
}
87+
88+
return e.EncodeBytes(codec, target, buf.Bytes())
89+
}
90+
91+
func (e *MultiEncoder) EncodeBytes(codec rawtopiccommon.Codec, target io.Writer, data []byte) (int, error) {
92+
enc, err := e.createEncodeWriter(codec, target)
93+
if err != nil {
94+
return 0, err
95+
}
96+
97+
n, err := enc.Write(data)
98+
if err == nil {
99+
err = enc.Close()
100+
}
101+
if err != nil {
102+
return 0, err
103+
}
104+
105+
if resetableEnc, ok := enc.(PublicResetableWriter); ok {
106+
e.ep[codec].Put(resetableEnc)
107+
}
108+
109+
return n, nil
110+
}
111+
112+
func (e *MultiEncoder) createEncodeWriter(codec rawtopiccommon.Codec, target io.Writer) (io.WriteCloser, error) {
113+
if ePool, ok := e.ep[codec]; ok {
114+
wr := ePool.Get()
115+
if wr != nil {
116+
wr.Reset(target)
117+
118+
return wr, nil
119+
}
120+
}
121+
42122
if encoderCreator, ok := e.m[codec]; ok {
43123
return encoderCreator(target)
44124
}
45125

46126
return nil, xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf("ydb: unexpected codec '%v' for encode message", codec)))
47127
}
48128

49-
func (e *EncoderMap) GetSupportedCodecs() rawtopiccommon.SupportedCodecs {
129+
func (e *MultiEncoder) GetSupportedCodecs() rawtopiccommon.SupportedCodecs {
50130
res := make(rawtopiccommon.SupportedCodecs, 0, len(e.m))
51131
for codec := range e.m {
52132
res = append(res, codec)
@@ -55,7 +135,7 @@ func (e *EncoderMap) GetSupportedCodecs() rawtopiccommon.SupportedCodecs {
55135
return res
56136
}
57137

58-
func (e *EncoderMap) IsSupported(codec rawtopiccommon.Codec) bool {
138+
func (e *MultiEncoder) IsSupported(codec rawtopiccommon.Codec) bool {
59139
_, ok := e.m[codec]
60140

61141
return ok
@@ -73,7 +153,7 @@ func (nopWriteCloser) Close() error {
73153

74154
// EncoderSelector not thread safe
75155
type EncoderSelector struct {
76-
m *EncoderMap
156+
m *MultiEncoder
77157

78158
tracer *trace.Topic
79159
writerReconnectorID string
@@ -87,7 +167,7 @@ type EncoderSelector struct {
87167
}
88168

89169
func NewEncoderSelector(
90-
m *EncoderMap,
170+
m *MultiEncoder,
91171
allowedCodecs rawtopiccommon.SupportedCodecs,
92172
parallelCompressors int,
93173
tracer *trace.Topic,

‎internal/topic/topicwriterinternal/encoders_test.go‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ package topicwriterinternal
22

33
import (
44
"bytes"
5+
"compress/gzip"
6+
"fmt"
7+
"io"
58
"strings"
69
"testing"
710

@@ -20,7 +23,7 @@ func TestEncoderSelector_CodecMeasure(t *testing.T) {
2023
})
2124
t.Run("One", func(t *testing.T) {
2225
s := NewEncoderSelector(
23-
NewEncoderMap(),
26+
NewMultiEncoder(),
2427
rawtopiccommon.SupportedCodecs{rawtopiccommon.CodecRaw},
2528
1,
2629
&trace.Topic{},
@@ -206,3 +209,89 @@ func TestCompressMessages(t *testing.T) {
206209
require.Error(t, cacheMessages(messages, rawtopiccommon.CodecGzip, parallelCount))
207210
})
208211
}
212+
213+
func TestMultiEncoder(t *testing.T) {
214+
decompressGzip := func(rd io.Reader) string {
215+
gzreader, err := gzip.NewReader(rd)
216+
require.NoError(t, err)
217+
218+
decompressedMsg, err := io.ReadAll(gzreader)
219+
require.NoError(t, err)
220+
221+
return string(decompressedMsg)
222+
}
223+
224+
t.Run("BuffersPool", func(t *testing.T) {
225+
testMultiEncoder := NewMultiEncoder()
226+
227+
buf := &bytes.Buffer{}
228+
for i := 0; i < 50; i++ {
229+
testMsg := []byte(fmt.Sprintf("test_data_%d", i))
230+
231+
buf.Reset()
232+
_, err := testMultiEncoder.Encode(rawtopiccommon.CodecGzip, buf, bytes.NewReader(testMsg))
233+
require.NoError(t, err)
234+
235+
require.Equal(t, string(testMsg), decompressGzip(buf))
236+
}
237+
})
238+
239+
t.Run("NotResetableWriter", func(t *testing.T) {
240+
testMultiEncoder := NewMultiEncoder()
241+
require.Len(t, testMultiEncoder.ep, 2)
242+
243+
buf := &bytes.Buffer{}
244+
_, err := testMultiEncoder.EncodeBytes(rawtopiccommon.CodecRaw, buf, []byte("test_data"))
245+
require.NoError(t, err)
246+
require.Equal(t, "test_data", buf.String())
247+
})
248+
249+
t.Run("ResetableWriterCustom", func(t *testing.T) {
250+
testMultiEncoder := NewMultiEncoder()
251+
252+
customCodecCode := rawtopiccommon.Codec(10001)
253+
testMultiEncoder.AddEncoder(customCodecCode, func(writer io.Writer) (io.WriteCloser, error) {
254+
return gzip.NewWriter(writer), nil
255+
})
256+
require.Len(t, testMultiEncoder.ep, 3)
257+
258+
buf := &bytes.Buffer{}
259+
_, err := testMultiEncoder.EncodeBytes(customCodecCode, buf, []byte("test_data_1"))
260+
require.NoError(t, err)
261+
require.Equal(t, "test_data_1", decompressGzip(buf))
262+
263+
buf.Reset()
264+
_, err = testMultiEncoder.EncodeBytes(rawtopiccommon.CodecGzip, buf, []byte("test_data_2"))
265+
require.NoError(t, err)
266+
require.Equal(t, "test_data_2", decompressGzip(buf))
267+
})
268+
269+
t.Run("ResetableWriter", func(t *testing.T) {
270+
testMultiEncoder := NewMultiEncoder()
271+
272+
buf := &bytes.Buffer{}
273+
_, err := testMultiEncoder.EncodeBytes(rawtopiccommon.CodecGzip, buf, []byte("test_data_1"))
274+
require.NoError(t, err)
275+
require.Equal(t, "test_data_1", decompressGzip(buf))
276+
277+
buf.Reset()
278+
_, err = testMultiEncoder.EncodeBytes(rawtopiccommon.CodecGzip, buf, []byte("test_data_2"))
279+
require.NoError(t, err)
280+
require.Equal(t, "test_data_2", decompressGzip(buf))
281+
})
282+
283+
t.Run("ResetableWriterManyMessages", func(t *testing.T) {
284+
testMultiEncoder := NewMultiEncoder()
285+
286+
buf := &bytes.Buffer{}
287+
for i := 0; i < 50; i++ {
288+
testMsg := []byte(fmt.Sprintf("test_data_%d", i))
289+
290+
buf.Reset()
291+
_, err := testMultiEncoder.EncodeBytes(rawtopiccommon.CodecGzip, buf, testMsg)
292+
require.NoError(t, err)
293+
294+
require.Equal(t, string(testMsg), decompressGzip(buf))
295+
}
296+
})
297+
}

‎internal/topic/topicwriterinternal/message.go‎

Lines changed: 6 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ type messageWithDataContent struct {
6767
bufCodec rawtopiccommon.Codec
6868
bufEncoded bytes.Buffer
6969
rawBuf bytes.Buffer
70-
encoders *EncoderMap
70+
encoders *MultiEncoder
7171
BufUncompressedSize int
7272
}
7373

@@ -111,18 +111,7 @@ func (m *messageWithDataContent) encodeRawContent(codec rawtopiccommon.Codec) ([
111111

112112
m.bufEncoded.Reset()
113113

114-
writer, err := m.encoders.CreateLazyEncodeWriter(codec, &m.bufEncoded)
115-
if err != nil {
116-
return nil, xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf(
117-
"ydb: failed create encoder for message, codec '%v': %w",
118-
codec,
119-
err,
120-
)))
121-
}
122-
_, err = writer.Write(m.rawBuf.Bytes())
123-
if err == nil {
124-
err = writer.Close()
125-
}
114+
_, err := m.encoders.EncodeBytes(codec, &m.bufEncoded, m.rawBuf.Bytes())
126115
if err != nil {
127116
return nil, xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf(
128117
"ydb: failed to compress message, codec '%v': %w",
@@ -157,27 +146,20 @@ func (m *messageWithDataContent) readDataToTargetCodec(codec rawtopiccommon.Code
157146
m.bufCodec = codec
158147
m.bufEncoded.Reset()
159148

160-
encoder, err := m.encoders.CreateLazyEncodeWriter(codec, &m.bufEncoded)
161-
if err != nil {
162-
return err
163-
}
164-
165149
reader := m.Data
166150
if reader == nil {
167151
reader = &bytes.Reader{}
168152
}
169-
bytesCount, err := io.Copy(encoder, reader)
170-
if err == nil {
171-
err = encoder.Close()
172-
}
153+
154+
bytesCount, err := m.encoders.Encode(codec, &m.bufEncoded, reader)
173155
if err != nil {
174156
return xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf(
175157
"ydb: failed compress message with codec '%v': %w",
176158
codec,
177159
err,
178160
)))
179161
}
180-
m.BufUncompressedSize = int(bytesCount)
162+
m.BufUncompressedSize = bytesCount
181163
m.Data = nil
182164

183165
return nil
@@ -218,7 +200,7 @@ func (m *messageWithDataContent) getEncodedBytes(codec rawtopiccommon.Codec) ([]
218200

219201
func newMessageDataWithContent(
220202
message PublicMessage, //nolint:gocritic
221-
encoders *EncoderMap,
203+
encoders *MultiEncoder,
222204
) messageWithDataContent {
223205
return messageWithDataContent{
224206
PublicMessage: message,

‎internal/topic/topicwriterinternal/writer_reconnector.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ type WriterReconnector struct {
124124
semaphore *semaphore.Weighted
125125
firstInitResponseProcessedChan empty.Chan
126126
lastSeqNo int64
127-
encodersMap *EncoderMap
127+
encodersMap *MultiEncoder
128128
initDoneCh empty.Chan
129129
initInfo InitialInfo
130130
m xsync.RWMutex
@@ -156,7 +156,7 @@ func newWriterReconnectorStopped(
156156
queue: newMessageQueue(),
157157
lastSeqNo: -1,
158158
firstInitResponseProcessedChan: make(empty.Chan),
159-
encodersMap: NewEncoderMap(),
159+
encodersMap: NewMultiEncoder(),
160160
writerInstanceID: writerInstanceID.String(),
161161
retrySettings: cfg.RetrySettings,
162162
}
@@ -760,11 +760,11 @@ func createRawMessageData(
760760
return res, err
761761
}
762762

763-
func calculateAllowedCodecs(forceCodec rawtopiccommon.Codec, encoderMap *EncoderMap,
763+
func calculateAllowedCodecs(forceCodec rawtopiccommon.Codec, multiEncoder *MultiEncoder,
764764
serverCodecs rawtopiccommon.SupportedCodecs,
765765
) rawtopiccommon.SupportedCodecs {
766766
if forceCodec != rawtopiccommon.CodecUNSPECIFIED {
767-
if serverCodecs.AllowedByCodecsList(forceCodec) && encoderMap.IsSupported(forceCodec) {
767+
if serverCodecs.AllowedByCodecsList(forceCodec) && multiEncoder.IsSupported(forceCodec) {
768768
return rawtopiccommon.SupportedCodecs{forceCodec}
769769
}
770770

@@ -779,7 +779,7 @@ func calculateAllowedCodecs(forceCodec rawtopiccommon.Codec, encoderMap *Encoder
779779

780780
res := make(rawtopiccommon.SupportedCodecs, 0, len(serverCodecs))
781781
for _, codec := range serverCodecs {
782-
if encoderMap.IsSupported(codec) {
782+
if multiEncoder.IsSupported(codec) {
783783
res = append(res, codec)
784784
}
785785
}

0 commit comments

Comments
 (0)