Merge pull request #1202 from uptrace/add-config-with-conn-buffer-size

feat(pg): allow user config buffer size of pg's connect

Author: j2gg0s

driver/pgdriver/config.go

@@ -48,6 +48,9 @@ type Config struct {
 
 	// Enable tracing
 	EnableTracing bool
+
+	// size of reader buffer, default is 4096
+	BufferSize int
 }
 
 func newDefaultConfig() *Config {
@@ -67,6 +70,8 @@ func newDefaultConfig() *Config {
 		WriteTimeout: 5 * time.Second,
 
 		EnableTracing: true,
+
+		BufferSize: 4096,
 	}
 
 	conf.Dialer = func(ctx context.Context, network, addr string) (net.Conn, error) {
@@ -90,6 +95,12 @@ func WithOptions(opts ...Option) Option {
 	}
 }
 
+func WithConfig(after *Config) Option {
+	return func(before *Config) {
+		*before = *after
+	}
+}
+
 func WithNetwork(network string) Option {
 	if network == "" {
 		panic("network is empty")
@@ -186,6 +197,12 @@ func WithWriteTimeout(writeTimeout time.Duration) Option {
 	}
 }
 
+func WithBufferSize(size int) Option {
+	return func(conf *Config) {
+		conf.BufferSize = size
+	}
+}
+
 // WithResetSessionFunc configures a function that is called prior to executing
 // a query on a connection that has been used before.
 // If the func returns driver.ErrBadConn, the connection is discarded.

driver/pgdriver/config_test.go

@@ -28,6 +28,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -42,6 +43,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   2 * time.Second,
 				WriteTimeout:  3 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -59,6 +61,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -73,6 +76,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -87,6 +91,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -101,6 +106,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 		{
@@ -115,6 +121,7 @@ func TestParseDSN(t *testing.T) {
 				ReadTimeout:   10 * time.Second,
 				WriteTimeout:  5 * time.Second,
 				EnableTracing: true,
+				BufferSize:    4096,
 			},
 		},
 	}

driver/pgdriver/driver.go

@@ -123,7 +123,7 @@ func newConn(ctx context.Context, conf *Config) (*Conn, error) {
 	cn := &Conn{
 		conf:    conf,
 		netConn: netConn,
-		rd:      newReader(netConn),
+		rd:      newReader(netConn, conf.BufferSize),
 	}
 
 	if conf.TLSConfig != nil {

driver/pgdriver/listener.go

@@ -36,11 +36,19 @@ type Listener struct {
 }
 
 func NewListener(db *bun.DB) *Listener {
-	return &Listener{
+	ln := &Listener{
 		db:     db,
 		driver: db.Driver().(Driver).connector,
 		exit:   make(chan struct{}),
 	}
+	if conf := ln.driver.Config(); conf.BufferSize < 8000 {
+		// https://github.com/uptrace/bun/issues/1201
+		// listener's payloads can be up to 8000 bytes
+		newConf := *conf
+		newConf.BufferSize = 8192
+		ln.driver = NewConnector(WithConfig(&newConf))
+	}
+	return ln
 }
 
 // Close closes the listener, releasing any open resources.

driver/pgdriver/proto.go

@@ -84,9 +84,9 @@ type reader struct {
 	buf []byte
 }
 
-func newReader(r io.Reader) *reader {
+func newReader(r io.Reader, size int) *reader {
 	return &reader{
-		Reader: bufio.NewReader(r),
+		Reader: bufio.NewReaderSize(r, size),
 		buf:    make([]byte, 128),
 	}
 }