Implement zlib compression

.github/workflows/test.yml

@@ -83,7 +83,7 @@ jobs:
           my-cnf: |
             innodb_log_file_size=256MB
             innodb_buffer_pool_size=512MB
-            max_allowed_packet=16MB
+            max_allowed_packet=48MB
             ; TestConcurrent fails if max_connections is too large
             max_connections=50
             local_infile=1

AUTHORS

@@ -21,6 +21,7 @@ Animesh Ray <mail.rayanimesh at gmail.com>
 Arne Hormann <arnehormann at gmail.com>
 Ariel Mashraki <ariel at mashraki.co.il>
 Asta Xie <xiemengjun at gmail.com>
+B Lamarche <blam413 at gmail.com>
 Brian Hendriks <brian at dolthub.com>
 Bulat Gaifullin <gaifullinbf at gmail.com>
 Caine Jette <jette at alum.mit.edu>
@@ -60,6 +61,7 @@ Jennifer Purevsuren <jennifer at dolthub.com>
 Jerome Meyer <jxmeyer at gmail.com>
 Jiajia Zhong <zhong2plus at gmail.com>
 Jian Zhen <zhenjl at gmail.com>
+Joe Mann <contact at joemann.co.uk>
 Joshua Prunier <joshua.prunier at gmail.com>
 Julien Lefevre <julien.lefevr at gmail.com>
 Julien Schmidt <go-sql-driver at julienschmidt.com>

README.md

@@ -38,6 +38,7 @@ A MySQL-Driver for Go's [database/sql](https://golang.org/pkg/database/sql/) pac
   * Secure `LOAD DATA LOCAL INFILE` support with file allowlisting and `io.Reader` support
   * Optional `time.Time` parsing
   * Optional placeholder interpolation
+  * Supports zlib compression.
 
 ## Requirements
 
@@ -267,6 +268,16 @@ SELECT u.id FROM users as u
 
 will return `u.id` instead of just `id` if `columnsWithAlias=true`.
 
+##### `compress`
+
+```
+Type:           bool
+Valid Values:   true, false
+Default:        false
+```
+
+Toggles zlib compression. false by default.
+
 ##### `interpolateParams`
 
 ```
@@ -310,6 +321,15 @@ Default:       64*1024*1024
 
 Max packet size allowed in bytes. The default value is 64 MiB and should be adjusted to match the server settings. `maxAllowedPacket=0` can be used to automatically fetch the `max_allowed_packet` variable from server *on every connection*.
 
+##### `minCompressLength`
+
+```
+Type:          decimal number
+Default:       50
+```
+
+Min packet size in bytes to compress, when compression is enabled (see the `compress` parameter). Packets smaller than this will be sent uncompressed.
+
 ##### `multiStatements`
 
 ```

benchmark_test.go

@@ -46,9 +46,13 @@ func (tb *TB) checkStmt(stmt *sql.Stmt, err error) *sql.Stmt {
 	return stmt
 }
 
-func initDB(b *testing.B, queries ...string) *sql.DB {
+func initDB(b *testing.B, useCompression bool, queries ...string) *sql.DB {
 	tb := (*TB)(b)
-	db := tb.checkDB(sql.Open(driverNameTest, dsn))
+	comprStr := ""
+	if useCompression {
+		comprStr = "&compress=1"
+	}
+	db := tb.checkDB(sql.Open(driverNameTest, dsn+comprStr))
 	for _, query := range queries {
 		if _, err := db.Exec(query); err != nil {
 			b.Fatalf("error on %q: %v", query, err)
@@ -60,10 +64,18 @@ func initDB(b *testing.B, queries ...string) *sql.DB {
 const concurrencyLevel = 10
 
 func BenchmarkQuery(b *testing.B) {
+	benchmarkQueryHelper(b, false)
+}
+
+func BenchmarkQueryCompression(b *testing.B) {
+	benchmarkQueryHelper(b, true)
+}
+
+func benchmarkQueryHelper(b *testing.B, compr bool) {
 	tb := (*TB)(b)
 	b.StopTimer()
 	b.ReportAllocs()
-	db := initDB(b,
+	db := initDB(b, compr,
 		"DROP TABLE IF EXISTS foo",
 		"CREATE TABLE foo (id INT PRIMARY KEY, val CHAR(50))",
 		`INSERT INTO foo VALUES (1, "one")`,
@@ -224,6 +236,7 @@ func BenchmarkInterpolation(b *testing.B) {
 		maxWriteSize:     maxPacketSize - 1,
 		buf:              newBuffer(nil),
 	}
+	mc.packetReader = &mc.buf
 
 	args := []driver.Value{
 		int64(42424242),
@@ -269,7 +282,7 @@ func benchmarkQueryContext(b *testing.B, db *sql.DB, p int) {
 }
 
 func BenchmarkQueryContext(b *testing.B) {
-	db := initDB(b,
+	db := initDB(b, false,
 		"DROP TABLE IF EXISTS foo",
 		"CREATE TABLE foo (id INT PRIMARY KEY, val CHAR(50))",
 		`INSERT INTO foo VALUES (1, "one")`,
@@ -305,7 +318,7 @@ func benchmarkExecContext(b *testing.B, db *sql.DB, p int) {
 }
 
 func BenchmarkExecContext(b *testing.B) {
-	db := initDB(b,
+	db := initDB(b, false,
 		"DROP TABLE IF EXISTS foo",
 		"CREATE TABLE foo (id INT PRIMARY KEY, val CHAR(50))",
 		`INSERT INTO foo VALUES (1, "one")`,
@@ -323,7 +336,7 @@ func BenchmarkExecContext(b *testing.B) {
 // "size=" means size of each blobs.
 func BenchmarkQueryRawBytes(b *testing.B) {
 	var sizes []int = []int{100, 1000, 2000, 4000, 8000, 12000, 16000, 32000, 64000, 256000}
-	db := initDB(b,
+	db := initDB(b, false,
 		"DROP TABLE IF EXISTS bench_rawbytes",
 		"CREATE TABLE bench_rawbytes (id INT PRIMARY KEY, val LONGBLOB)",
 	)
@@ -376,7 +389,7 @@ func BenchmarkQueryRawBytes(b *testing.B) {
 // BenchmarkReceiveMassiveRows measures performance of receiving large number of rows.
 func BenchmarkReceiveMassiveRows(b *testing.B) {
 	// Setup -- prepare 10000 rows.
-	db := initDB(b,
+	db := initDB(b, false,
 		"DROP TABLE IF EXISTS foo",
 		"CREATE TABLE foo (id INT PRIMARY KEY, val TEXT)")
 	defer db.Close()

compress.go

@@ -0,0 +1,246 @@
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2024 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"bytes"
+	"compress/zlib"
+	"fmt"
+	"io"
+	"sync"
+)
+
+var (
+	zrPool *sync.Pool // Do not use directly. Use zDecompress() instead.
+	zwPool *sync.Pool // Do not use directly. Use zCompress() instead.
+)
+
+func init() {
+	zrPool = &sync.Pool{
+		New: func() any { return nil },
+	}
+	zwPool = &sync.Pool{
+		New: func() any {
+			zw, err := zlib.NewWriterLevel(new(bytes.Buffer), 2)
+			if err != nil {
+				panic(err) // compress/zlib return non-nil error only if level is invalid
+			}
+			return zw
+		},
+	}
+}
+
+func zDecompress(src, dst []byte) (int, error) {
+	br := bytes.NewReader(src)
+	var zr io.ReadCloser
+	var err error
+
+	if a := zrPool.Get(); a == nil {
+		if zr, err = zlib.NewReader(br); err != nil {
+			return 0, err
+		}
+	} else {
+		zr = a.(io.ReadCloser)
+		if zr.(zlib.Resetter).Reset(br, nil); err != nil {
+			return 0, err
+		}
+	}
+	defer func() {
+		zr.Close()
+		zrPool.Put(zr)
+	}()
+
+	lenRead := 0
+	size := len(dst)
+
+	for lenRead < size {
+		n, err := zr.Read(dst[lenRead:])
+		lenRead += n
+
+		if err == io.EOF {
+			if lenRead < size {
+				return lenRead, io.ErrUnexpectedEOF
+			}
+		} else if err != nil {
+			return lenRead, err
+		}
+	}
+	return lenRead, nil
+}
+
+func zCompress(src []byte, dst io.Writer) error {
+	zw := zwPool.Get().(*zlib.Writer)
+	zw.Reset(dst)
+	if _, err := zw.Write(src); err != nil {
+		return err
+	}
+	zw.Close()
+	zwPool.Put(zw)
+	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
+}
+
+func newDecompressor(mc *mysqlConn) *decompressor {
+	return &decompressor{
+		mc: mc,
+	}
+}
+
+func (c *decompressor) readNext(need int) ([]byte, error) {
+	for len(c.bytesBuf) < 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
+}
+
+func (c *decompressor) uncompressPacket() error {
+	header, err := c.mc.buf.readNext(7) // size of compressed header
+	if err != nil {
+		return err
+	}
+
+	// compressed header structure
+	comprLength := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
+	uncompressedLength := int(uint32(header[4]) | uint32(header[5])<<8 | uint32(header[6])<<16)
+	compressionSequence := uint8(header[3])
+	if debugTrace {
+		traceLogger.Printf("uncompress cmplen=%v uncomplen=%v pkt_cmp_seq=%v expected_cmp_seq=%v\n",
+			comprLength, uncompressedLength, compressionSequence, c.mc.sequence)
+	}
+	if compressionSequence != c.mc.sequence {
+		// return ErrPktSync
+		// server may return error packet (e.g. 1153 Got a packet bigger than 'max_allowed_packet' bytes)
+		// before receiving all packets from client. In this case, seqnr is younger than expected.
+		if debugTrace {
+			traceLogger.Printf("WARN: unexpected cmpress seq nr: expected %v, got %v",
+				c.mc.sequence, compressionSequence)
+		}
+		// TODO(methane): report error when the packet is not an error packet.
+	}
+	c.mc.sequence = compressionSequence + 1
+	c.mc.compressSequence = c.mc.sequence
+
+	comprData, err := c.mc.buf.readNext(comprLength)
+	if err != nil {
+		return err
+	}
+
+	// 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...)
+		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])
+	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]
+	return nil
+}
+
+const maxPayloadLen = maxPacketSize - 4
+
+// writeCompressed 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) {
+	totalBytes := len(packets)
+	dataLen := len(packets)
+	blankHeader := make([]byte, 7)
+	var buf bytes.Buffer
+
+	for dataLen > 0 {
+		payloadLen := dataLen
+		if payloadLen > maxPayloadLen {
+			payloadLen = maxPayloadLen
+		}
+		payload := packets[:payloadLen]
+		uncompressedLen := payloadLen
+
+		if _, err := buf.Write(blankHeader); err != nil {
+			return 0, err
+		}
+
+		// If payload is less than minCompressLength, don't compress.
+		if uncompressedLen < minCompressLength {
+			if _, err := buf.Write(payload); err != nil {
+				return 0, err
+			}
+			uncompressedLen = 0
+		} else {
+			zCompress(payload, &buf)
+		}
+
+		if err := mc.writeCompressedPacket(buf.Bytes(), uncompressedLen); err != nil {
+			return 0, err
+		}
+		dataLen -= payloadLen
+		packets = packets[payloadLen:]
+		buf.Reset()
+	}
+
+	return totalBytes, nil
+}
+
+// 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 {
+	comprLength := len(data) - 7
+	if debugTrace {
+		traceLogger.Printf(
+			"writeCompressedPacket: comprLength=%v, uncompressedLen=%v, seq=%v",
+			comprLength, uncompressedLen, mc.compressSequence)
+	}
+
+	// compression header
+	data[0] = byte(0xff & comprLength)
+	data[1] = byte(0xff & (comprLength >> 8))
+	data[2] = byte(0xff & (comprLength >> 16))
+
+	data[3] = mc.compressSequence
+
+	// this value is never greater than maxPayloadLength
+	data[4] = byte(0xff & uncompressedLen)
+	data[5] = byte(0xff & (uncompressedLen >> 8))
+	data[6] = byte(0xff & (uncompressedLen >> 16))
+
+	if _, err := mc.netConn.Write(data); err != nil {
+		mc.log("writing compressed packet:", err)
+		return err
+	}
+
+	mc.compressSequence++
+	return nil
+}