refactoring

Author: methane

benchmark_test.go

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

buffer.go

@@ -10,31 +10,30 @@ package mysql
 
 import (
 	"io"
-	"net"
-	"time"
 )
 
 const defaultBufSize = 4096
 const maxCachedBufSize = 256 * 1024
 
+// readwriteFunc is a function that compatible with io.Reader and io.Writer.
+// We use this function type instead of io.ReadWriter because we want to
+// just pass mc.readWithTimeout or mc.writeWithTimeout functions.
+type readwriteFunc func([]byte) (int, error)
+
 // A buffer which is used for both reading and writing.
 // This is possible since communication on each connection is synchronous.
 // In other words, we can't write and read simultaneously on the same connection.
 // The buffer is similar to bufio.Reader / Writer but zero-copy-ish
 // Also highly optimized for this particular use case.
 type buffer struct {
-	buf          []byte // read buffer.
-	cachedBuf    []byte // buffer that will be reused. len(cachedBuf) <= maxCachedBufSize.
-	nc           net.Conn
-	readTimeout  time.Duration
-	writeTimeout time.Duration
+	buf       []byte // read buffer.
+	cachedBuf []byte // buffer that will be reused. len(cachedBuf) <= maxCachedBufSize.
 }
 
 // newBuffer allocates and returns a new buffer.
-func newBuffer(nc net.Conn) buffer {
+func newBuffer() buffer {
 	return buffer{
 		cachedBuf: make([]byte, defaultBufSize),
-		nc:        nc,
 	}
 }
 
@@ -44,7 +43,7 @@ func (b *buffer) busy() bool {
 }
 
 // fill reads into the read buffer until at least _need_ bytes are in it.
-func (b *buffer) fill(need int) error {
+func (b *buffer) fill(need int, r readwriteFunc) error {
 	// we'll move the contents of the current buffer to dest before filling it.
 	dest := b.cachedBuf
 
@@ -65,13 +64,7 @@ func (b *buffer) fill(need int) error {
 	copy(dest[:n], b.buf)
 
 	for {
-		if b.readTimeout > 0 {
-			if err := b.nc.SetReadDeadline(time.Now().Add(b.readTimeout)); err != nil {
-				return err
-			}
-		}
-
-		nn, err := b.nc.Read(dest[n:])
+		nn, err := r(dest[n:])
 		n += nn
 
 		if err == nil && n < need {
@@ -93,10 +86,10 @@ func (b *buffer) fill(need int) error {
 
 // returns next N bytes from buffer.
 // The returned slice is only guaranteed to be valid until the next read
-func (b *buffer) readNext(need int) ([]byte, error) {
+func (b *buffer) readNext(need int, r readwriteFunc) ([]byte, error) {
 	if len(b.buf) < need {
 		// refill
-		if err := b.fill(need); err != nil {
+		if err := b.fill(need, r); err != nil {
 			return nil, err
 		}
 	}
@@ -156,14 +149,3 @@ 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) {
-	if b.writeTimeout > 0 {
-		if err := b.nc.SetWriteDeadline(time.Now().Add(b.writeTimeout)); err != nil {
-			return 0, err
-		}
-	}
-	return b.nc.Write(packets)
-}

compress.go

@@ -47,7 +47,7 @@ func zDecompress(src, dst []byte) (int, error) {
 		}
 	} else {
 		zr = a.(io.ReadCloser)
-		if zr.(zlib.Resetter).Reset(br, nil); err != nil {
+		if err := zr.(zlib.Resetter).Reset(br, nil); err != nil {
 			return 0, err
 		}
 	}
@@ -96,18 +96,18 @@ func newCompIO(mc *mysqlConn) *compIO {
 	}
 }
 
-func (c *compIO) readNext(need int) ([]byte, error) {
+func (c *compIO) readNext(need int, r readwriteFunc) ([]byte, error) {
 	for c.buff.Len() < need {
-		if err := c.readCompressedPacket(); err != nil {
+		if err := c.readCompressedPacket(r); err != nil {
 			return nil, err
 		}
 	}
 	data := c.buff.Next(need)
 	return data[:need:need], nil // prevent caller writes into c.buff
 }
 
-func (c *compIO) readCompressedPacket() error {
-	header, err := c.mc.buf.readNext(7) // size of compressed header
+func (c *compIO) readCompressedPacket(r readwriteFunc) error {
+	header, err := c.mc.buf.readNext(7, r) // size of compressed header
 	if err != nil {
 		return err
 	}
@@ -134,7 +134,7 @@ func (c *compIO) readCompressedPacket() error {
 	c.mc.sequence = compressionSequence + 1
 	c.mc.compressSequence = c.mc.sequence
 
-	comprData, err := c.mc.buf.readNext(comprLength)
+	comprData, err := c.mc.buf.readNext(comprLength, r)
 	if err != nil {
 		return err
 	}
@@ -221,7 +221,7 @@ func (c *compIO) writeCompressedPacket(data []byte, uncompressedLen int) error {
 	data[3] = mc.compressSequence
 	putUint24(data[4:7], uncompressedLen)
 
-	if _, err := mc.buf.writePackets(data); err != nil {
+	if _, err := mc.writeWithTimeout(data); err != nil {
 		mc.log("writing compressed packet:", err)
 		return err
 	}

compress_test.go

@@ -11,7 +11,6 @@ package mysql
 import (
 	"bytes"
 	"crypto/rand"
-	"fmt"
 	"io"
 	"testing"
 )
@@ -26,15 +25,12 @@ 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 := comp.writePackets(uncompressedPacket)
+	err := mc.writePacket(append(make([]byte, 4), uncompressedPacket...))
 	if err != nil {
 		t.Fatal(err)
 	}
-	if n != len(uncompressedPacket) {
-		t.Fatalf("expected to write %d bytes, wrote %d bytes", len(uncompressedPacket), n)
-	}
+
 	return conn.written
 }
 
@@ -43,10 +39,9 @@ func uncompressHelper(t *testing.T, mc *mysqlConn, compressedPacket []byte, expS
 	// mocking out buf variable
 	conn := new(mockConn)
 	conn.data = compressedPacket
-	mc.buf.nc = conn
-	cr := newCompIO(mc)
+	mc.netConn = conn
 
-	uncompressedPacket, err := cr.readNext(expSize)
+	uncompressedPacket, err := mc.readPacket()
 	if err != nil {
 		if err != io.EOF {
 			t.Fatalf("non-nil/non-EOF error when reading contents: %s", err.Error())
@@ -72,8 +67,6 @@ func TestRoundtrip(t *testing.T) {
 	}{
 		{uncompressed: []byte("a"),
 			desc: "a"},
-		{uncompressed: []byte{0},
-			desc: "0 byte"},
 		{uncompressed: []byte("hello world"),
 			desc: "hello world"},
 		{uncompressed: make([]byte, 100),
@@ -82,8 +75,6 @@ func TestRoundtrip(t *testing.T) {
 			desc: "32768 bytes"},
 		{uncompressed: make([]byte, 330000),
 			desc: "33000 bytes"},
-		{uncompressed: make([]byte, 0),
-			desc: "nothing"},
 		{uncompressed: makeRandByteSlice(10),
 			desc: "10 rand bytes",
 		},
@@ -100,26 +91,29 @@ func TestRoundtrip(t *testing.T) {
 
 	_, cSend := newRWMockConn(0)
 	cSend.compress = true
+	cSend.compIO = newCompIO(cSend)
 	_, cReceive := newRWMockConn(0)
 	cReceive.compress = true
+	cReceive.compIO = newCompIO(cReceive)
 
 	for _, test := range tests {
-		s := fmt.Sprintf("Test roundtrip with %s", test.desc)
-		cSend.resetSequenceNr()
-		cReceive.resetSequenceNr()
+		t.Run(test.desc, func(t *testing.T) {
+			cSend.resetSequenceNr()
+			cReceive.resetSequenceNr()
 
-		uncompressed := roundtripHelper(t, cSend, cReceive, test.uncompressed)
-		if !bytes.Equal(uncompressed, test.uncompressed) {
-			t.Fatalf("%s: roundtrip failed", s)
-		}
+			uncompressed := roundtripHelper(t, cSend, cReceive, test.uncompressed)
+			if !bytes.Equal(uncompressed, test.uncompressed) {
+				t.Fatalf("roundtrip failed")
+			}
 
-		if cSend.sequence != cReceive.sequence {
-			t.Errorf("inconsistent sequence number: send=%v recv=%v",
-				cSend.sequence, cReceive.sequence)
-		}
-		if cSend.compressSequence != cReceive.compressSequence {
-			t.Errorf("inconsistent compress sequence number: send=%v recv=%v",
-				cSend.compressSequence, cReceive.compressSequence)
-		}
+			if cSend.sequence != cReceive.sequence {
+				t.Errorf("inconsistent sequence number: send=%v recv=%v",
+					cSend.sequence, cReceive.sequence)
+			}
+			if cSend.compressSequence != cReceive.compressSequence {
+				t.Errorf("inconsistent compress sequence number: send=%v recv=%v",
+					cSend.compressSequence, cReceive.compressSequence)
+			}
+		})
 	}
 }

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().
-	packetRW         packetIO
+	compIO           *compIO
 	cfg              *Config
 	connector        *connector
 	maxAllowedPacket int
@@ -64,9 +64,24 @@ func (mc *mysqlConn) log(v ...any) {
 	mc.cfg.Logger.Print(v...)
 }
 
-type packetIO interface {
-	readNext(need int) ([]byte, error)
-	writePackets(data []byte) (int, error)
+func (mc *mysqlConn) readWithTimeout(b []byte) (int, error) {
+	to := mc.cfg.ReadTimeout
+	if to > 0 {
+		if err := mc.netConn.SetReadDeadline(time.Now().Add(to)); err != nil {
+			return 0, err
+		}
+	}
+	return mc.netConn.Read(b)
+}
+
+func (mc *mysqlConn) writeWithTimeout(b []byte) (int, error) {
+	to := mc.cfg.WriteTimeout
+	if to > 0 {
+		if err := mc.netConn.SetWriteDeadline(time.Now().Add(to)); err != nil {
+			return 0, err
+		}
+	}
+	return mc.netConn.Write(b)
 }
 
 func (mc *mysqlConn) resetSequenceNr() {

connection_test.go

@@ -19,13 +19,12 @@ import (
 
 func TestInterpolateParams(t *testing.T) {
 	mc := &mysqlConn{
-		buf:              newBuffer(nil),
+		buf:              newBuffer(),
 		maxAllowedPacket: maxPacketSize,
 		cfg: &Config{
 			InterpolateParams: true,
 		},
 	}
-	mc.packetRW = &mc.buf
 
 	q, err := mc.interpolateParams("SELECT ?+?", []driver.Value{int64(42), "gopher"})
 	if err != nil {
@@ -40,7 +39,7 @@ func TestInterpolateParams(t *testing.T) {
 
 func TestInterpolateParamsJSONRawMessage(t *testing.T) {
 	mc := &mysqlConn{
-		buf:              newBuffer(nil),
+		buf:              newBuffer(),
 		maxAllowedPacket: maxPacketSize,
 		cfg: &Config{
 			InterpolateParams: true,
@@ -67,13 +66,12 @@ func TestInterpolateParamsJSONRawMessage(t *testing.T) {
 
 func TestInterpolateParamsTooManyPlaceholders(t *testing.T) {
 	mc := &mysqlConn{
-		buf:              newBuffer(nil),
+		buf:              newBuffer(),
 		maxAllowedPacket: maxPacketSize,
 		cfg: &Config{
 			InterpolateParams: true,
 		},
 	}
-	mc.packetRW = &mc.buf
 
 	q, err := mc.interpolateParams("SELECT ?+?", []driver.Value{int64(42)})
 	if err != driver.ErrSkip {
@@ -85,15 +83,13 @@ func TestInterpolateParamsTooManyPlaceholders(t *testing.T) {
 // https://github.com/go-sql-driver/mysql/pull/490
 func TestInterpolateParamsPlaceholderInString(t *testing.T) {
 	mc := &mysqlConn{
-		buf:              newBuffer(nil),
+		buf:              newBuffer(),
 		maxAllowedPacket: maxPacketSize,
 		cfg: &Config{
 			InterpolateParams: true,
 		},
 	}
 
-	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`
 	if err != driver.ErrSkip {
@@ -103,7 +99,7 @@ func TestInterpolateParamsPlaceholderInString(t *testing.T) {
 
 func TestInterpolateParamsUint64(t *testing.T) {
 	mc := &mysqlConn{
-		buf:              newBuffer(nil),
+		buf:              newBuffer(),
 		maxAllowedPacket: maxPacketSize,
 		cfg: &Config{
 			InterpolateParams: true,
@@ -164,11 +160,10 @@ func TestCleanCancel(t *testing.T) {
 func TestPingMarkBadConnection(t *testing.T) {
 	nc := badConnection{err: errors.New("boom")}
 
-	buf := newBuffer(nc)
+	buf := newBuffer()
 	mc := &mysqlConn{
 		netConn:          nc,
 		buf:              buf,
-		packetRW:         &buf,
 		maxAllowedPacket: defaultMaxAllowedPacket,
 		closech:          make(chan struct{}),
 		cfg:              NewConfig(),
@@ -184,11 +179,9 @@ func TestPingMarkBadConnection(t *testing.T) {
 func TestPingErrInvalidConn(t *testing.T) {
 	nc := badConnection{err: errors.New("failed to write"), n: 10}
 
-	buf := newBuffer(nc)
 	mc := &mysqlConn{
 		netConn:          nc,
-		buf:              buf,
-		packetRW:         &buf,
+		buf:              newBuffer(),
 		maxAllowedPacket: defaultMaxAllowedPacket,
 		closech:          make(chan struct{}),
 		cfg:              NewConfig(),