reduce memory allocations when zlib compression is enabled (#880)

* simple metrics to capture packet data size both compressed & uncompressed

* reusing zlib readers and writers via a sync.Pool

* cache zlib writer/readers on a connection to avoid excessive Writer creation, which reduces memory allocations

* fixing issue when reading large packet responses

* revert to standard go zlib implemenation

* Revert "revert to standard go zlib implemenation"

This reverts commit a7893fa.

* Revert "fixing issue when reading large packet responses"

This reverts commit 4f512d8.

* Revert "cache zlib writer/readers on a connection to avoid excessive Writer creation, which reduces memory allocations"

This reverts commit 25e0cee.

* refactoring to reduce memory allocations

* cleanup changes in prep for PR into upstream

* make zlib poll return ReadCloser instead of Reader

* fixing broken tests and linting issues

* removed duplicate assignment

* update link to doxygen docs

* addressing PR review feedback

---------

Co-authored-by: dvilaverde <[email protected]>

Author: dvilaverde

client/req.go

@@ -39,19 +39,23 @@ func (c *Conn) writeCommandStr(command byte, arg string) error {
 func (c *Conn) writeCommandUint32(command byte, arg uint32) error {
 	c.ResetSequence()
 
-	return c.WritePacket([]byte{
-		0x05, //5 bytes long
-		0x00,
-		0x00,
-		0x00, //sequence
+	buf := utils.ByteSliceGet(9)
 
-		command,
+	buf.B[0] = 0x05 //5 bytes long
+	buf.B[1] = 0x00
+	buf.B[2] = 0x00
+	buf.B[3] = 0x00 //sequence
 
-		byte(arg),
-		byte(arg >> 8),
-		byte(arg >> 16),
-		byte(arg >> 24),
-	})
+	buf.B[4] = command
+
+	buf.B[5] = byte(arg)
+	buf.B[6] = byte(arg >> 8)
+	buf.B[7] = byte(arg >> 16)
+	buf.B[8] = byte(arg >> 24)
+
+	err := c.WritePacket(buf.B)
+	utils.ByteSlicePut(buf)
+	return err
 }
 
 func (c *Conn) writeCommandStrStr(command byte, arg1 string, arg2 string) error {

client/stmt.go

@@ -7,6 +7,7 @@ import (
 	"math"
 
 	. "github.com/go-mysql-org/go-mysql/mysql"
+	"github.com/go-mysql-org/go-mysql/utils"
 	"github.com/pingcap/errors"
 )
 
@@ -146,37 +147,44 @@ func (s *Stmt) write(args ...interface{}) error {
 		length += len(paramValues[i])
 	}
 
-	data := make([]byte, 4, 4+length)
+	data := utils.BytesBufferGet()
+	defer func() {
+		utils.BytesBufferPut(data)
+	}()
+	if data.Len() < length+4 {
+		data.Grow(4 + length)
+	}
 
-	data = append(data, COM_STMT_EXECUTE)
-	data = append(data, byte(s.id), byte(s.id>>8), byte(s.id>>16), byte(s.id>>24))
+	data.Write([]byte{0, 0, 0, 0})
+	data.WriteByte(COM_STMT_EXECUTE)
+	data.Write([]byte{byte(s.id), byte(s.id >> 8), byte(s.id >> 16), byte(s.id >> 24)})
 
 	//flag: CURSOR_TYPE_NO_CURSOR
-	data = append(data, 0x00)
+	data.WriteByte(0x00)
 
 	//iteration-count, always 1
-	data = append(data, 1, 0, 0, 0)
+	data.Write([]byte{1, 0, 0, 0})
 
 	if s.params > 0 {
-		data = append(data, nullBitmap...)
+		data.Write(nullBitmap)
 
 		//new-params-bound-flag
-		data = append(data, newParamBoundFlag)
+		data.WriteByte(newParamBoundFlag)
 
 		if newParamBoundFlag == 1 {
 			//type of each parameter, length: num-params * 2
-			data = append(data, paramTypes...)
+			data.Write(paramTypes)
 
 			//value of each parameter
 			for _, v := range paramValues {
-				data = append(data, v...)
+				data.Write(v)
 			}
 		}
 	}
 
 	s.conn.ResetSequence()
 
-	return s.conn.WritePacket(data)
+	return s.conn.WritePacket(data.Bytes())
 }
 
 func (c *Conn) Prepare(query string) (*Stmt, error) {

compress/zlib.go

@@ -0,0 +1,96 @@
+package compress
+
+import (
+	"bytes"
+	"io"
+	"sync"
+
+	"github.com/klauspost/compress/zlib"
+)
+
+const DefaultCompressionLevel = 6
+
+var (
+	zlibReaderPool *sync.Pool
+	zlibWriterPool sync.Pool
+)
+
+func init() {
+	zlibReaderPool = &sync.Pool{
+		New: func() interface{} {
+			return nil
+		},
+	}
+
+	zlibWriterPool = sync.Pool{
+		New: func() interface{} {
+			w, err := zlib.NewWriterLevel(new(bytes.Buffer), DefaultCompressionLevel)
+			if err != nil {
+				panic(err)
+			}
+			return w
+		},
+	}
+}
+
+var _ io.WriteCloser = zlibWriter{}
+var _ io.ReadCloser = zlibReader{}
+
+type zlibWriter struct {
+	w *zlib.Writer
+}
+
+type zlibReader struct {
+	r io.ReadCloser
+}
+
+func GetPooledZlibWriter(target io.Writer) (io.WriteCloser, error) {
+	w := zlibWriterPool.Get().(*zlib.Writer)
+	w.Reset(target)
+
+	return zlibWriter{
+		w: w,
+	}, nil
+}
+
+func GetPooledZlibReader(src io.Reader) (io.ReadCloser, error) {
+	var (
+		rc  io.ReadCloser
+		err error
+	)
+
+	if r := zlibReaderPool.Get(); r != nil {
+		rc = r.(io.ReadCloser)
+		if rc.(zlib.Resetter).Reset(src, nil) != nil {
+			return nil, err
+		}
+	} else {
+		if rc, err = zlib.NewReader(src); err != nil {
+			return nil, err
+		}
+	}
+
+	return zlibReader{
+		r: rc,
+	}, nil
+}
+
+func (c zlibWriter) Write(data []byte) (n int, err error) {
+	return c.w.Write(data)
+}
+
+func (c zlibWriter) Close() error {
+	err := c.w.Close()
+	zlibWriterPool.Put(c.w)
+	return err
+}
+
+func (d zlibReader) Read(buf []byte) (n int, err error) {
+	return d.r.Read(buf)
+}
+
+func (d zlibReader) Close() error {
+	err := d.r.Close()
+	zlibReaderPool.Put(d.r)
+	return err
+}

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/goccy/go-json v0.10.2
 	github.com/google/uuid v1.3.0
 	github.com/jmoiron/sqlx v1.3.3
-	github.com/klauspost/compress v1.17.1
+	github.com/klauspost/compress v1.17.8
 	github.com/pingcap/errors v0.11.5-0.20221009092201-b66cddb77c32
 	github.com/pingcap/tidb/pkg/parser v0.0.0-20231103042308-035ad5ccbe67
 	github.com/shopspring/decimal v1.2.0

go.sum

@@ -19,8 +19,8 @@ github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/jmoiron/sqlx v1.3.3 h1:j82X0bf7oQ27XeqxicSZsTU5suPwKElg3oyxNn43iTk=
 github.com/jmoiron/sqlx v1.3.3/go.mod h1:2BljVx/86SuTyjE+aPYlHCTNvZrnJXghYGpNiXLBMCQ=
-github.com/klauspost/compress v1.17.1 h1:NE3C767s2ak2bweCZo3+rdP4U/HoyVXLv/X9f2gPS5g=
-github.com/klauspost/compress v1.17.1/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
+github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
+github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=