Skip to content

Commit

Permalink
allow packet.Conn buffersize to be adjustable (#892)
Browse files Browse the repository at this point in the history
  • Loading branch information
dvilaverde authored Jun 23, 2024
1 parent 7ee1864 commit 8d0b3e3
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 5 deletions.
2 changes: 1 addition & 1 deletion client/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ func (c *Conn) writeAuthHandshake() error {
}

currentSequence := c.Sequence
c.Conn = packet.NewConn(tlsConn)
c.Conn = packet.NewBufferedConn(tlsConn, c.BufferSize)
c.Sequence = currentSequence
}

Expand Down
8 changes: 7 additions & 1 deletion client/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ import (
"github.com/go-mysql-org/go-mysql/utils"
)

const defaultBufferSize = 65536 // 64kb

type Option func(*Conn) error

type Conn struct {
Expand All @@ -33,6 +35,9 @@ type Conn struct {
ReadTimeout time.Duration
WriteTimeout time.Duration

// The buffer size to use in the packet connection
BufferSize int

serverVersion string
// server capabilities
capability uint32
Expand Down Expand Up @@ -94,6 +99,7 @@ type Dialer func(ctx context.Context, network, address string) (net.Conn, error)
func ConnectWithDialer(ctx context.Context, network, addr, user, password, dbName string, dialer Dialer, options ...Option) (*Conn, error) {
c := new(Conn)

c.BufferSize = defaultBufferSize
c.attributes = map[string]string{
"_client_name": "go-mysql",
// "_client_version": "0.1",
Expand Down Expand Up @@ -129,7 +135,7 @@ func ConnectWithDialer(ctx context.Context, network, addr, user, password, dbNam
}
}

c.Conn = packet.NewConnWithTimeout(conn, c.ReadTimeout, c.WriteTimeout)
c.Conn = packet.NewConnWithTimeout(conn, c.ReadTimeout, c.WriteTimeout, c.BufferSize)
if c.tlsConfig != nil {
seq := c.Conn.Sequence
c.Conn = packet.NewTLSConnWithTimeout(conn, c.ReadTimeout, c.WriteTimeout)
Expand Down
24 changes: 24 additions & 0 deletions driver/driver_options_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"math"
"net"
"reflect"
"strconv"
"strings"
"testing"
"time"
Expand Down Expand Up @@ -73,6 +74,29 @@ func TestDriverOptions_ConnectTimeout(t *testing.T) {
conn.Close()
}

func TestDriverOptions_BufferSize(t *testing.T) {
log.SetLevel(log.LevelDebug)
srv := CreateMockServer(t)
defer srv.Stop()

SetDSNOptions(map[string]DriverOption{
"bufferSize": func(c *client.Conn, value string) error {
var err error
c.BufferSize, err = strconv.Atoi(value)
return err
},
})

conn, err := sql.Open("mysql", "[email protected]:3307/test?bufferSize=4096")
require.NoError(t, err)

rows, err := conn.QueryContext(context.TODO(), "select * from table;")
require.NotNil(t, rows)
require.NoError(t, err)

conn.Close()
}

func TestDriverOptions_ReadTimeout(t *testing.T) {
log.SetLevel(log.LevelDebug)
srv := CreateMockServer(t)
Expand Down
10 changes: 7 additions & 3 deletions packet/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,19 +53,23 @@ type Conn struct {
}

func NewConn(conn net.Conn) *Conn {
return NewBufferedConn(conn, 65536) // 64kb
}

func NewBufferedConn(conn net.Conn, bufferSize int) *Conn {
c := new(Conn)
c.Conn = conn

c.br = bufio.NewReaderSize(c, 65536) // 64kb
c.br = bufio.NewReaderSize(c, bufferSize)
c.reader = c.br

c.copyNBuf = make([]byte, DefaultBufferSize)

return c
}

func NewConnWithTimeout(conn net.Conn, readTimeout, writeTimeout time.Duration) *Conn {
c := NewConn(conn)
func NewConnWithTimeout(conn net.Conn, readTimeout, writeTimeout time.Duration, bufferSize int) *Conn {
c := NewBufferedConn(conn, bufferSize)
c.readTimeout = readTimeout
c.writeTimeout = writeTimeout
return c
Expand Down

0 comments on commit 8d0b3e3

Please sign in to comment.