Implement compression

AUTHORS

@@ -20,6 +20,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>
 Bulat Gaifullin <gaifullinbf at gmail.com>
 Caine Jette <jette at alum.mit.edu>
 Carlos Nieto <jose.carlos at menteslibres.net>
@@ -55,6 +56,7 @@ Jeffrey Charles <jeffreycharles at gmail.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>

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("mysql", dsn))
+	comprStr := ""
+	if useCompression {
+		comprStr = "&compress=1"
+	}
+	db := tb.checkDB(sql.Open("mysql", dsn+comprStr))
 	for _, query := range queries {
 		if _, err := db.Exec(query); err != nil {
 			b.Fatalf("error on %q: %v", query, err)
@@ -60,10 +64,19 @@ 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")`,
@@ -222,9 +235,11 @@ func BenchmarkInterpolation(b *testing.B) {
 		},
 		maxAllowedPacket: maxPacketSize,
 		maxWriteSize:     maxPacketSize - 1,
-		buf:              newBuffer(nil),
 	}
 
+	buf := newBuffer(nil)
+	mc.reader = &buf
+
 	args := []driver.Value{
 		int64(42424242),
 		float64(math.Pi),
@@ -269,7 +284,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 +320,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 +338,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 +391,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()

buffer.go

@@ -171,6 +171,31 @@ func (b *buffer) takeCompleteBuffer() ([]byte, error) {
 	return b.buf, nil
 }
 
+// returns a buffer with the requested size.
+// If possible, a slice from the existing buffer is returned.
+// Otherwise a bigger buffer is made.
+// Only one buffer (total) can be used at a time.
+func (b *buffer) reuseBuffer(length int) ([]byte, error) {
+	if length == -1 {
+		return b.takeCompleteBuffer()
+	}
+
+	if b.length > 0 {
+		return nil, nil
+	}
+
+	// test (cheap) general case first
+	if length <= defaultBufSize || length <= cap(b.buf) {
+		return b.buf[:length], nil
+	}
+
+	if length < maxPacketSize {
+		b.buf = make([]byte, length)
+		return b.buf, nil
+	}
+	return make([]byte, length), nil
+}
+
 // store stores buf, an updated buffer, if its suitable to do so.
 func (b *buffer) store(buf []byte) error {
 	if b.length > 0 {

compress.go

@@ -0,0 +1,236 @@
+package mysql
+
+import (
+	"bytes"
+	"compress/zlib"
+	"io"
+)
+
+const (
+	minCompressLength = 50
+)
+
+type compressedReader struct {
+	buf      packetReader
+	bytesBuf []byte
+	mc       *mysqlConn
+	zr       io.ReadCloser
+}
+
+type compressedWriter struct {
+	connWriter io.Writer
+	mc         *mysqlConn
+	zw         *zlib.Writer
+}
+
+func newCompressedReader(buf packetReader, mc *mysqlConn) *compressedReader {
+	return &compressedReader{
+		buf:      buf,
+		bytesBuf: make([]byte, 0),
+		mc:       mc,
+	}
+}
+
+func newCompressedWriter(connWriter io.Writer, mc *mysqlConn) *compressedWriter {
+	return &compressedWriter{
+		connWriter: connWriter,
+		mc:         mc,
+		zw:         zlib.NewWriter(new(bytes.Buffer)),
+	}
+}
+
+func (cr *compressedReader) readNext(need int) ([]byte, error) {
+	for len(cr.bytesBuf) < need {
+		err := cr.uncompressPacket()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	data := cr.bytesBuf[:need]
+	cr.bytesBuf = cr.bytesBuf[need:]
+	return data, nil
+}
+
+func (cr *compressedReader) reuseBuffer(length int) ([]byte, error) {
+	return cr.buf.reuseBuffer(length)
+}
+
+func (cr *compressedReader) uncompressPacket() error {
+	header, err := cr.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 compressionSequence != cr.mc.compressionSequence {
+		return ErrPktSync
+	}
+
+	cr.mc.compressionSequence++
+
+	comprData, err := cr.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 {
+		cr.bytesBuf = append(cr.bytesBuf, comprData...)
+		return nil
+	}
+
+	// write comprData to a bytes.buffer, then read it using zlib into data
+	br := bytes.NewReader(comprData)
+
+	if cr.zr == nil {
+		cr.zr, err = zlib.NewReader(br)
+	} else {
+		err = cr.zr.(zlib.Resetter).Reset(br, nil)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	defer cr.zr.Close()
+
+	// use existing capacity in bytesBuf if possible
+	offset := len(cr.bytesBuf)
+	if cap(cr.bytesBuf)-offset < uncompressedLength {
+		old := cr.bytesBuf
+		cr.bytesBuf = make([]byte, offset, offset+uncompressedLength)
+		copy(cr.bytesBuf, old)
+	}
+
+	data := cr.bytesBuf[offset : offset+uncompressedLength]
+
+	lenRead := 0
+
+	// http://grokbase.com/t/gg/golang-nuts/146y9ppn6b/go-nuts-stream-compression-with-compress-flate
+	for lenRead < uncompressedLength {
+		n, err := cr.zr.Read(data[lenRead:])
+		lenRead += n
+
+		if err == io.EOF {
+			if lenRead < uncompressedLength {
+				return io.ErrUnexpectedEOF
+			}
+			break
+		}
+
+		if err != nil {
+			return err
+		}
+	}
+
+	cr.bytesBuf = append(cr.bytesBuf, data...)
+
+	return nil
+}
+
+func (cw *compressedWriter) Write(data []byte) (int, error) {
+	// when asked to write an empty packet, do nothing
+	if len(data) == 0 {
+		return 0, nil
+	}
+
+	totalBytes := len(data)
+	length := len(data) - 4
+	maxPayloadLength := maxPacketSize - 4
+	blankHeader := make([]byte, 7)
+
+	for length >= maxPayloadLength {
+		payload := data[:maxPayloadLength]
+		payloadLen := len(payload)
+
+		bytesBuf := &bytes.Buffer{}
+		bytesBuf.Write(blankHeader)
+		cw.zw.Reset(bytesBuf)
+		_, err := cw.zw.Write(payload)
+		if err != nil {
+			return 0, err
+		}
+		cw.zw.Close()
+
+		// if compression expands the payload, do not compress
+		compressedPayload := bytesBuf.Bytes()
+		if len(compressedPayload) > maxPayloadLength {
+			compressedPayload = append(blankHeader, payload...)
+			payloadLen = 0
+		}
+
+		err = cw.writeToNetwork(compressedPayload, payloadLen)
+
+		if err != nil {
+			return 0, err
+		}
+
+		length -= maxPayloadLength
+		data = data[maxPayloadLength:]
+	}
+
+	payloadLen := len(data)
+
+	// do not attempt compression if packet is too small
+	if payloadLen < minCompressLength {
+		err := cw.writeToNetwork(append(blankHeader, data...), 0)
+		if err != nil {
+			return 0, err
+		}
+		return totalBytes, nil
+	}
+
+	bytesBuf := &bytes.Buffer{}
+	bytesBuf.Write(blankHeader)
+	cw.zw.Reset(bytesBuf)
+	_, err := cw.zw.Write(data)
+	if err != nil {
+		return 0, err
+	}
+	cw.zw.Close()
+
+	compressedPayload := bytesBuf.Bytes()
+
+	if len(compressedPayload) > len(data) {
+		compressedPayload = append(blankHeader, data...)
+		payloadLen = 0
+	}
+
+	// add header and send over the wire
+	err = cw.writeToNetwork(compressedPayload, payloadLen)
+	if err != nil {
+		return 0, err
+	}
+
+	return totalBytes, nil
+}
+
+func (cw *compressedWriter) writeToNetwork(data []byte, uncomprLength int) error {
+	comprLength := len(data) - 7
+
+	// compression header
+	data[0] = byte(0xff & comprLength)
+	data[1] = byte(0xff & (comprLength >> 8))
+	data[2] = byte(0xff & (comprLength >> 16))
+
+	data[3] = cw.mc.compressionSequence
+
+	// this value is never greater than maxPayloadLength
+	data[4] = byte(0xff & uncomprLength)
+	data[5] = byte(0xff & (uncomprLength >> 8))
+	data[6] = byte(0xff & (uncomprLength >> 16))
+
+	if _, err := cw.connWriter.Write(data); err != nil {
+		return err
+	}
+
+	cw.mc.compressionSequence++
+	return nil
+}