packets: implemented compression protocol CR changes

Author: Brigitte Lamarche

benchmark_go18_test.go

@@ -42,7 +42,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")`,
@@ -78,7 +78,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")`,

benchmark_test.go

@@ -43,9 +43,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 {
 			if w, ok := err.(MySQLWarnings); ok {
@@ -61,10 +65,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")`,

compress.go

@@ -18,6 +18,8 @@ type compressedReader struct {
 	buf      packetReader
 	bytesBuf []byte
 	mc       *mysqlConn
+	br       *bytes.Reader
+	zr       io.ReadCloser
 }
 
 type compressedWriter struct {
@@ -48,12 +50,8 @@ func (cr *compressedReader) readNext(need int) ([]byte, error) {
 		}
 	}
 
-	data := make([]byte, need)
-
-	copy(data, cr.bytesBuf[:len(data)])
-
-	cr.bytesBuf = cr.bytesBuf[len(data):]
-
+	data := cr.bytesBuf[:need]
+	cr.bytesBuf = cr.bytesBuf[need:]
 	return data, nil
 }
 
@@ -88,27 +86,43 @@ func (cr *compressedReader) uncompressPacket() error {
 	}
 
 	// write comprData to a bytes.buffer, then read it using zlib into data
-	var b bytes.Buffer
-	b.Write(comprData)
-	r, err := zlib.NewReader(&b)
+	if cr.br == nil {
+		cr.br = bytes.NewReader(comprData)
+	} else {
+		cr.br.Reset(comprData)
+	}
+
+	resetter, ok := cr.zr.(zlib.Resetter)
 
-	if r != nil {
-		defer r.Close()
+	if ok {
+		err := resetter.Reset(cr.br, []byte{})
+		if err != nil {
+			return err
+		}
+	} else {
+		cr.zr, err = zlib.NewReader(cr.br)
+		if err != nil {
+			return err
+		}
 	}
 
-	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 := make([]byte, uncompressedLength)
+	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 {
-
-		tmp := data[lenRead:]
-
-		n, err := r.Read(tmp)
+		n, err := cr.zr.Read(data[lenRead:])
 		lenRead += n
 
 		if err == io.EOF {
@@ -152,7 +166,15 @@ func (cw *compressedWriter) Write(data []byte) (int, error) {
 			return 0, err
 		}
 
-		err = cw.writeComprPacketToNetwork(b.Bytes(), lenSmall)
+		// if compression expands the payload, do not compress
+		useData := b.Bytes()
+
+		if len(useData) > len(dataSmall) {
+			useData = dataSmall
+			lenSmall = 0
+		}
+
+		err = cw.writeComprPacketToNetwork(useData, lenSmall)
 		if err != nil {
 			return 0, err
 		}
@@ -163,7 +185,7 @@ func (cw *compressedWriter) Write(data []byte) (int, error) {
 
 	lenSmall := len(data)
 
-	// do not compress if packet is too small
+	// do not attempt compression if packet is too small
 	if lenSmall < minCompressLength {
 		err := cw.writeComprPacketToNetwork(data, 0)
 		if err != nil {
@@ -183,7 +205,15 @@ func (cw *compressedWriter) Write(data []byte) (int, error) {
 		return 0, err
 	}
 
-	err = cw.writeComprPacketToNetwork(b.Bytes(), lenSmall)
+	// if compression expands the payload, do not compress
+	useData := b.Bytes()
+
+	if len(useData) > len(data) {
+		useData = data
+		lenSmall = 0
+	}
+
+	err = cw.writeComprPacketToNetwork(useData, lenSmall)
 
 	if err != nil {
 		return 0, err

packets.go

@@ -881,6 +881,7 @@ func (stmt *mysqlStmt) writeCommandLongData(paramID int, arg []byte) error {
 		}
 
 		stmt.mc.sequence = 0
+		stmt.mc.compressionSequence = 0
 		// Add command byte [1 byte]
 		data[4] = comStmtSendLongData
 
@@ -906,6 +907,7 @@ func (stmt *mysqlStmt) writeCommandLongData(paramID int, arg []byte) error {
 
 	// Reset Packet Sequence
 	stmt.mc.sequence = 0
+	stmt.mc.compressionSequence = 0
 	return nil
 }
 
@@ -925,6 +927,7 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 
 	// Reset packet-sequence
 	mc.sequence = 0
+	mc.compressionSequence = 0
 
 	var data []byte