compress: better buffer reuse

Author: methane

benchmark_test.go

@@ -236,7 +236,7 @@ func BenchmarkInterpolation(b *testing.B) {
 		maxWriteSize:     maxPacketSize - 1,
 		buf:              newBuffer(nil),
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	args := []driver.Value{
 		int64(42424242),

buffer.go

@@ -155,3 +155,9 @@ func (b *buffer) store(buf []byte) {
 		b.cachedBuf = buf[:cap(buf)]
 	}
 }
+
+// writePackets is a proxy function to nc.Write.
+// This is used to make the buffer type compatible with compressed I/O.
+func (b *buffer) writePackets(packets []byte) (int, error) {
+	return b.nc.Write(packets)
+}

compress.go

@@ -85,34 +85,28 @@ func zCompress(src []byte, dst io.Writer) error {
 	return nil
 }
 
-type decompressor struct {
-	mc *mysqlConn
-	// read buffer (FIFO).
-	// We can not reuse already-read buffer until dropping Go 1.20 support.
-	// It is because of database/mysql's weired behavior.
-	// See https://github.com/go-sql-driver/mysql/issues/1435
-	bytesBuf []byte
+type compIO struct {
+	mc   *mysqlConn
+	buff bytes.Buffer
 }
 
-func newDecompressor(mc *mysqlConn) *decompressor {
-	return &decompressor{
+func newCompIO(mc *mysqlConn) *compIO {
+	return &compIO{
 		mc: mc,
 	}
 }
 
-func (c *decompressor) readNext(need int) ([]byte, error) {
-	for len(c.bytesBuf) < need {
+func (c *compIO) readNext(need int) ([]byte, error) {
+	for c.buff.Len() < need {
 		if err := c.uncompressPacket(); err != nil {
 			return nil, err
 		}
 	}
-
-	data := c.bytesBuf[:need:need] // prevent caller writes into r.bytesBuf
-	c.bytesBuf = c.bytesBuf[need:]
-	return data, nil
+	data := c.buff.Next(need)
+	return data[:need:need], nil // prevent caller writes into c.buff
 }
 
-func (c *decompressor) uncompressPacket() error {
+func (c *compIO) uncompressPacket() error {
 	header, err := c.mc.buf.readNext(7) // size of compressed header
 	if err != nil {
 		return err
@@ -147,41 +141,37 @@ func (c *decompressor) uncompressPacket() error {
 	// if payload is uncompressed, its length will be specified as zero, and its
 	// true length is contained in comprLength
 	if uncompressedLength == 0 {
-		c.bytesBuf = append(c.bytesBuf, comprData...)
+		c.buff.Write(comprData)
 		return nil
 	}
 
 	// use existing capacity in bytesBuf if possible
-	offset := len(c.bytesBuf)
-	if cap(c.bytesBuf)-offset < uncompressedLength {
-		old := c.bytesBuf
-		c.bytesBuf = make([]byte, offset, offset+uncompressedLength)
-		copy(c.bytesBuf, old)
-	}
-
-	lenRead, err := zDecompress(comprData, c.bytesBuf[offset:offset+uncompressedLength])
+	c.buff.Grow(uncompressedLength)
+	dec := c.buff.AvailableBuffer()[:uncompressedLength]
+	lenRead, err := zDecompress(comprData, dec)
 	if err != nil {
 		return err
 	}
 	if lenRead != uncompressedLength {
 		return fmt.Errorf("invalid compressed packet: uncompressed length in header is %d, actual %d",
 			uncompressedLength, lenRead)
 	}
-	c.bytesBuf = c.bytesBuf[:offset+uncompressedLength]
+	c.buff.Write(dec) // fast copy. See bytes.Buffer.AvailableBuffer() doc.
 	return nil
 }
 
 const maxPayloadLen = maxPacketSize - 4
 
-// writeCompressed sends one or some packets with compression.
+// writePackets sends one or some packets with compression.
 // Use this instead of mc.netConn.Write() when mc.compress is true.
-func (mc *mysqlConn) writeCompressed(packets []byte) (int, error) {
+func (c *compIO) writePackets(packets []byte) (int, error) {
 	totalBytes := len(packets)
 	dataLen := len(packets)
 	blankHeader := make([]byte, 7)
-	var buf bytes.Buffer
+	buf := &c.buff
 
 	for dataLen > 0 {
+		buf.Reset()
 		payloadLen := dataLen
 		if payloadLen > maxPayloadLen {
 			payloadLen = maxPayloadLen
@@ -200,10 +190,10 @@ func (mc *mysqlConn) writeCompressed(packets []byte) (int, error) {
 			}
 			uncompressedLen = 0
 		} else {
-			zCompress(payload, &buf)
+			zCompress(payload, buf)
 		}
 
-		if err := mc.writeCompressedPacket(buf.Bytes(), uncompressedLen); err != nil {
+		if err := c.writeCompressedPacket(buf.Bytes(), uncompressedLen); err != nil {
 			return 0, err
 		}
 		dataLen -= payloadLen
@@ -216,7 +206,8 @@ func (mc *mysqlConn) writeCompressed(packets []byte) (int, error) {
 
 // writeCompressedPacket writes a compressed packet with header.
 // data should start with 7 size space for header followed by payload.
-func (mc *mysqlConn) writeCompressedPacket(data []byte, uncompressedLen int) error {
+func (c *compIO) writeCompressedPacket(data []byte, uncompressedLen int) error {
+	mc := c.mc
 	comprLength := len(data) - 7
 	if debugTrace {
 		fmt.Printf(

compress_test.go

@@ -26,8 +26,9 @@ func makeRandByteSlice(size int) []byte {
 func compressHelper(t *testing.T, mc *mysqlConn, uncompressedPacket []byte) []byte {
 	conn := new(mockConn)
 	mc.netConn = conn
+	comp := newCompIO(mc)
 
-	n, err := mc.writeCompressed(uncompressedPacket)
+	n, err := comp.writePackets(uncompressedPacket)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -43,7 +44,7 @@ func uncompressHelper(t *testing.T, mc *mysqlConn, compressedPacket []byte, expS
 	conn := new(mockConn)
 	conn.data = compressedPacket
 	mc.buf.nc = conn
-	cr := newDecompressor(mc)
+	cr := newCompIO(mc)
 
 	uncompressedPacket, err := cr.readNext(expSize)
 	if err != nil {

connection.go

@@ -28,7 +28,7 @@ type mysqlConn struct {
 	netConn          net.Conn
 	rawConn          net.Conn    // underlying connection when netConn is TLS connection.
 	result           mysqlResult // managed by clearResult() and handleOkPacket().
-	packetReader     packetReader
+	packetRW         packetIO
 	cfg              *Config
 	connector        *connector
 	maxAllowedPacket int
@@ -65,8 +65,9 @@ func (mc *mysqlConn) log(v ...any) {
 	mc.cfg.Logger.Print(v...)
 }
 
-type packetReader interface {
+type packetIO interface {
 	readNext(need int) ([]byte, error)
+	writePackets(data []byte) (int, error)
 }
 
 func (mc *mysqlConn) resetSequenceNr() {

connection_test.go

@@ -25,7 +25,7 @@ func TestInterpolateParams(t *testing.T) {
 			InterpolateParams: true,
 		},
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	q, err := mc.interpolateParams("SELECT ?+?", []driver.Value{int64(42), "gopher"})
 	if err != nil {
@@ -73,7 +73,7 @@ func TestInterpolateParamsTooManyPlaceholders(t *testing.T) {
 			InterpolateParams: true,
 		},
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	q, err := mc.interpolateParams("SELECT ?+?", []driver.Value{int64(42)})
 	if err != driver.ErrSkip {
@@ -92,7 +92,7 @@ func TestInterpolateParamsPlaceholderInString(t *testing.T) {
 		},
 	}
 
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	q, err := mc.interpolateParams("SELECT 'abc?xyz',?", []driver.Value{int64(42)})
 	// When InterpolateParams support string literal, this should return `"SELECT 'abc?xyz', 42`
@@ -168,7 +168,7 @@ func TestPingMarkBadConnection(t *testing.T) {
 	mc := &mysqlConn{
 		netConn:          nc,
 		buf:              buf,
-		packetReader:     &buf,
+		packetRW:         &buf,
 		maxAllowedPacket: defaultMaxAllowedPacket,
 		closech:          make(chan struct{}),
 		cfg:              NewConfig(),
@@ -188,7 +188,7 @@ func TestPingErrInvalidConn(t *testing.T) {
 	mc := &mysqlConn{
 		netConn:          nc,
 		buf:              buf,
-		packetReader:     &buf,
+		packetRW:         &buf,
 		maxAllowedPacket: defaultMaxAllowedPacket,
 		closech:          make(chan struct{}),
 		cfg:              NewConfig(),

connector.go

@@ -129,7 +129,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	mc.buf = newBuffer(mc.netConn)
 	// packet reader and writer in handshake are never compressed
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	// Set I/O timeouts
 	mc.buf.timeout = mc.cfg.ReadTimeout
@@ -174,7 +174,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	if mc.cfg.compress && mc.flags&clientCompress == clientCompress {
 		mc.compress = true
-		mc.packetReader = newDecompressor(mc)
+		mc.packetRW = newCompIO(mc)
 	}
 	if mc.cfg.MaxAllowedPacket > 0 {
 		mc.maxAllowedPacket = mc.cfg.MaxAllowedPacket

packets.go

@@ -32,7 +32,7 @@ func (mc *mysqlConn) readPacket() ([]byte, error) {
 
 	for {
 		// read packet header
-		data, err := mc.packetReader.readNext(4)
+		data, err := mc.packetRW.readNext(4)
 		if err != nil {
 			mc.close()
 			if cerr := mc.canceled.Value(); cerr != nil {
@@ -80,7 +80,7 @@ func (mc *mysqlConn) readPacket() ([]byte, error) {
 		}
 
 		// read packet body [pktLen bytes]
-		data, err = mc.packetReader.readNext(pktLen)
+		data, err = mc.packetRW.readNext(pktLen)
 		if err != nil {
 			mc.close()
 			if cerr := mc.canceled.Value(); cerr != nil {
@@ -143,15 +143,7 @@ func (mc *mysqlConn) writePacket(data []byte) error {
 			}
 		}
 
-		var (
-			n   int
-			err error
-		)
-		if mc.compress {
-			n, err = mc.writeCompressed(data[:4+size])
-		} else {
-			n, err = mc.netConn.Write(data[:4+size])
-		}
+		n, err := mc.packetRW.writePackets(data[:4+size])
 		if err != nil {
 			mc.cleanup()
 			if cerr := mc.canceled.Value(); cerr != nil {

packets_test.go

@@ -100,7 +100,7 @@ func newRWMockConn(sequence uint8) (*mockConn, *mysqlConn) {
 	buf := newBuffer(conn)
 	mc := &mysqlConn{
 		buf:              buf,
-		packetReader:     &buf,
+		packetRW:         &buf,
 		cfg:              connector.cfg,
 		connector:        connector,
 		netConn:          conn,
@@ -116,7 +116,7 @@ func TestReadPacketSingleByte(t *testing.T) {
 	mc := &mysqlConn{
 		buf: newBuffer(conn),
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	conn.data = []byte{0x01, 0x00, 0x00, 0x00, 0xff}
 	conn.maxReads = 1
@@ -169,7 +169,7 @@ func TestReadPacketSplit(t *testing.T) {
 	mc := &mysqlConn{
 		buf: newBuffer(conn),
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	data := make([]byte, maxPacketSize*2+4*3)
 	const pkt2ofs = maxPacketSize + 4
@@ -277,7 +277,7 @@ func TestReadPacketFail(t *testing.T) {
 		closech: make(chan struct{}),
 		cfg:     NewConfig(),
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	// illegal empty (stand-alone) packet
 	conn.data = []byte{0x00, 0x00, 0x00, 0x00}
@@ -323,7 +323,7 @@ func TestRegression801(t *testing.T) {
 		sequence: 42,
 		closech:  make(chan struct{}),
 	}
-	mc.packetReader = &mc.buf
+	mc.packetRW = &mc.buf
 
 	conn.data = []byte{72, 0, 0, 42, 10, 53, 46, 53, 46, 56, 0, 165, 0, 0, 0,
 		60, 70, 63, 58, 68, 104, 34, 97, 0, 223, 247, 33, 2, 0, 15, 128, 21, 0,