Add unit test

Signed-off-by: Matt Lord <[email protected]>

go/mysql/binlog_event_compression.go

@@ -26,6 +26,7 @@ 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"
@@ -325,11 +326,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 {
-			statefulDecoderPool.Put(reader)
-		}
+		statefulDecoderPool.Put(reader)
 	default:
-		reader = nil
 	}
 	tp.iterator = nil
 }
@@ -361,8 +359,10 @@ type decoderPool struct {
 
 // Get gets a pooled OR new *zstd.Decoder.
 func (dp *decoderPool) Get(reader io.Reader) (*zstd.Decoder, error) {
-	decoder := dp.pool.Get().(*zstd.Decoder)
-	if decoder == nil {
+	var decoder *zstd.Decoder
+	if pooled := dp.pool.Get(); pooled != nil {
+		decoder = pooled.(*zstd.Decoder)
+	} else {
 		d, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(zstdInMemoryDecompressorMaxSize))
 		if err != nil { // Should only happen e.g. due to ENOMEM
 			return nil, vterrors.New(vtrpcpb.Code_INTERNAL, "failed to create stateful stream decoder")
@@ -376,7 +376,10 @@ func (dp *decoderPool) Get(reader io.Reader) (*zstd.Decoder, error) {
 }
 
 func (dp *decoderPool) Put(decoder *zstd.Decoder) {
-	if err := decoder.Reset(nil); err == nil || err == io.EOF {
+	err := decoder.Reset(nil)
+	if err == nil || err == io.EOF {
 		dp.pool.Put(decoder)
+	} else {
+		log.Errorf("DEBUG: suck it: %v", err)
 	}
 }

go/mysql/binlog_event_compression_test.go

@@ -0,0 +1,63 @@
+/*
+Copyright 2023 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) {
+	type args struct {
+		r io.Reader
+	}
+	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) {
+			decoder, err := statefulDecoderPool.Get(tt.reader)
+			require.NoError(t, err)
+			require.NotNil(t, decoder)
+			require.IsType(t, &zstd.Decoder{}, decoder)
+			statefulDecoderPool.Put(decoder)
+			decoder2, err := statefulDecoderPool.Get(tt.reader)
+			require.NoError(t, err)
+			require.NotNil(t, decoder2)
+			require.IsType(t, &zstd.Decoder{}, decoder)
+			statefulDecoderPool.Put(decoder)
+			require.True(t, (decoder2 == decoder))
+			statefulDecoderPool.Put(decoder2)
+			decoder3, err := statefulDecoderPool.Get(tt.reader)
+			require.NoError(t, err)
+			require.IsType(t, &zstd.Decoder{}, decoder)
+			statefulDecoderPool.Put(decoder)
+			require.True(t, (decoder3 == decoder2))
+		})
+	}
+}