Skip to content

Commit

Permalink
Use proper zstd decoder pool for binlog event compression handling (#…
Browse files Browse the repository at this point in the history
…17042)

Signed-off-by: Matt Lord <[email protected]>
  • Loading branch information
mattlord authored Oct 22, 2024
1 parent 03b93a1 commit 2203cfb
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 31 deletions.
77 changes: 46 additions & 31 deletions go/mysql/binlog_event_compression.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@ import (
"github.com/klauspost/compress/zstd"

"vitess.io/vitess/go/stats"
"vitess.io/vitess/go/vt/log"
"vitess.io/vitess/go/vt/vterrors"

vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
Expand Down Expand Up @@ -89,26 +88,9 @@ var (
// allocations and GC overhead so this pool allows us to handle
// concurrent cases better while still scaling to 0 when there's no
// usage.
statefulDecoderPool sync.Pool
statefulDecoderPool = &decoderPool{}
)

func init() {
var err error
statelessDecoder, err = zstd.NewReader(nil, zstd.WithDecoderConcurrency(0))
if err != nil { // Should only happen e.g. due to ENOMEM
log.Errorf("Error creating stateless decoder: %v", err)
}
statefulDecoderPool = sync.Pool{
New: func() any {
d, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(zstdInMemoryDecompressorMaxSize))
if err != nil { // Should only happen e.g. due to ENOMEM
log.Errorf("Error creating stateful decoder: %v", err)
}
return d
},
}
}

type TransactionPayload struct {
size uint64
compressionType uint64
Expand Down Expand Up @@ -304,21 +286,22 @@ func (tp *TransactionPayload) decompress() error {
// larger payloads.
if tp.uncompressedSize > zstdInMemoryDecompressorMaxSize {
in := bytes.NewReader(tp.payload)
streamDecoder := statefulDecoderPool.Get().(*zstd.Decoder)
if streamDecoder == nil {
return vterrors.New(vtrpcpb.Code_INTERNAL, "failed to create stateful stream decoder")
}
if err := streamDecoder.Reset(in); err != nil {
return vterrors.Wrap(err, "error resetting stateful stream decoder")
streamDecoder, err := statefulDecoderPool.Get(in)
if err != nil {
return err
}
compressedTrxPayloadsUsingStream.Add(1)
tp.reader = streamDecoder
return nil
}

// Process smaller payloads using only in-memory buffers.
if statelessDecoder == nil { // Should never happen
return vterrors.New(vtrpcpb.Code_INTERNAL, "failed to create stateless decoder")
if statelessDecoder == nil { // Should only need to be done once
var err error
statelessDecoder, err = zstd.NewReader(nil, zstd.WithDecoderConcurrency(0))
if err != nil { // Should only happen e.g. due to ENOMEM
return vterrors.Wrap(err, "failed to create stateless decoder")
}
}
decompressedBytes := make([]byte, 0, tp.uncompressedSize) // Perform a single pre-allocation
decompressedBytes, err := statelessDecoder.DecodeAll(tp.payload, decompressedBytes[:0])
Expand All @@ -340,11 +323,8 @@ func (tp *TransactionPayload) decompress() error {
func (tp *TransactionPayload) Close() {
switch reader := tp.reader.(type) {
case *zstd.Decoder:
if err := reader.Reset(nil); err == nil || err == io.EOF {
readersPool.Put(reader)
}
statefulDecoderPool.Put(reader)
default:
reader = nil
}
tp.iterator = nil
}
Expand All @@ -368,3 +348,38 @@ func (tp *TransactionPayload) GetNextEvent() (BinlogEvent, error) {
//func (tp *TransactionPayload) Events() iter.Seq[BinlogEvent] {
// return tp.iterator
//}

// decoderPool manages a sync.Pool of *zstd.Decoders.
type decoderPool struct {
pool sync.Pool
}

// Get gets a pooled OR new *zstd.Decoder.
func (dp *decoderPool) Get(reader io.Reader) (*zstd.Decoder, error) {
var (
decoder *zstd.Decoder
ok bool
)
if pooled := dp.pool.Get(); pooled != nil {
decoder, ok = pooled.(*zstd.Decoder)
if !ok {
return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] expected *zstd.Decoder but got %T", pooled)
}
} else {
d, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(zstdInMemoryDecompressorMaxSize))
if err != nil { // Should only happen e.g. due to ENOMEM
return nil, vterrors.Wrap(err, "failed to create stateful stream decoder")
}
decoder = d
}
if err := decoder.Reset(reader); err != nil {
return nil, vterrors.Wrap(err, "error resetting stateful stream decoder")
}
return decoder, nil
}

func (dp *decoderPool) Put(decoder *zstd.Decoder) {
if err := decoder.Reset(nil); err == nil || err == io.EOF {
dp.pool.Put(decoder)
}
}
75 changes: 75 additions & 0 deletions go/mysql/binlog_event_compression_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*
Copyright 2024 The Vitess 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 mysql

import (
"bytes"
"io"
"testing"

"github.com/klauspost/compress/zstd"
"github.com/stretchr/testify/require"
)

func TestDecoderPool(t *testing.T) {
validateDecoder := func(t *testing.T, err error, decoder *zstd.Decoder) {
require.NoError(t, err)
require.NotNil(t, decoder)
require.IsType(t, &zstd.Decoder{}, decoder)
}
tests := []struct {
name string
reader io.Reader
wantErr bool
}{
{
name: "happy path",
reader: bytes.NewReader([]byte{0x68, 0x61, 0x70, 0x70, 0x79}),
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// It's not guaranteed that we get the same decoder back from the pool
// that we just put in, so we use a loop and ensure that it worked at
// least one of the times. Without doing this the test would be flaky.
poolingUsed := false

for i := 0; i < 20; i++ {
decoder, err := statefulDecoderPool.Get(tt.reader)
validateDecoder(t, err, decoder)
statefulDecoderPool.Put(decoder)

decoder2, err := statefulDecoderPool.Get(tt.reader)
validateDecoder(t, err, decoder2)
if decoder2 == decoder {
poolingUsed = true
}
statefulDecoderPool.Put(decoder2)

decoder3, err := statefulDecoderPool.Get(tt.reader)
validateDecoder(t, err, decoder3)
if decoder3 == decoder || decoder3 == decoder2 {
poolingUsed = true
}
statefulDecoderPool.Put(decoder3)
}

require.True(t, poolingUsed)
})
}
}

0 comments on commit 2203cfb

Please sign in to comment.