From ebec40d991c1ceb3dbb5c9957a0529c540968ba7 Mon Sep 17 00:00:00 2001
From: alkhan <alkhan@tradestation.com>
Date: Mon, 27 Apr 2020 22:50:07 -0500
Subject: [PATCH 1/2] Add buffer size option to connection

---
 conn.go         | 18 +++++++++++++++---
 conn_options.go | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/conn.go b/conn.go
index 9de9a74..11c8ab9 100644
--- a/conn.go
+++ b/conn.go
@@ -70,9 +70,6 @@ func Dial(network, addr string, opts ...func(*Conn) error) (*Conn, error) {
 // been created by the program. The opts parameter provides the
 // opportunity to specify STOMP protocol options.
 func Connect(conn io.ReadWriteCloser, opts ...func(*Conn) error) (*Conn, error) {
-	reader := frame.NewReader(conn)
-	writer := frame.NewWriter(conn)
-
 	c := &Conn{
 		conn:       conn,
 		closeMutex: &sync.Mutex{},
@@ -83,6 +80,21 @@ func Connect(conn io.ReadWriteCloser, opts ...func(*Conn) error) (*Conn, error)
 		return nil, err
 	}
 
+	var reader *frame.Reader
+	var writer *frame.Writer
+
+	if options.ReadBufferSize > 0 {
+		reader = frame.NewReaderSize(conn, options.ReadBufferSize)
+	} else {
+		reader = frame.NewReader(conn)
+	}
+
+	if options.WriteBufferSize > 0 {
+		writer = frame.NewWriterSize(conn, options.ReadBufferSize)
+	} else {
+		writer = frame.NewWriter(conn)
+	}
+
 	readChannelCapacity := 20
 	writeChannelCapacity := 20
 
diff --git a/conn_options.go b/conn_options.go
index 86020b8..dc3fbef 100644
--- a/conn_options.go
+++ b/conn_options.go
@@ -22,6 +22,7 @@ type connOptions struct {
 	AcceptVersions                            []string
 	Header                                    *frame.Header
 	ReadChannelCapacity, WriteChannelCapacity int
+	ReadBufferSize, WriteBufferSize           int
 }
 
 func newConnOptions(conn *Conn, opts []func(*Conn) error) (*connOptions, error) {
@@ -156,6 +157,16 @@ var ConnOpt struct {
 	// same time. A high number may affect memory usage while a too low number may lock the
 	// system up. Default is set to 20.
 	WriteChannelCapacity func(capacity int) func(*Conn) error
+
+	// ReadBufferSize specifies number of bytes that can be used to read the message
+	// A high number may affect memory usage while a too low number may lock the
+	// system up. Default is set to 4096.
+	ReadBufferSize func(size int) func(*Conn) error
+
+	// WriteBufferSize specifies number of bytes that can be used to write the message
+	// A high number may affect memory usage while a too low number may lock the
+	// system up. Default is set to 4096.
+	WriteBufferSize func(size int) func(*Conn) error
 }
 
 func init() {
@@ -244,4 +255,18 @@ func init() {
 			return nil
 		}
 	}
+
+	ConnOpt.ReadBufferSize = func(size int) func(*Conn) error {
+		return func(c *Conn) error {
+			c.options.ReadBufferSize = size
+			return nil
+		}
+	}
+
+	ConnOpt.WriteBufferSize = func(size int) func(*Conn) error {
+		return func(c *Conn) error {
+			c.options.WriteBufferSize = size
+			return nil
+		}
+	}
 }

From cd5b9c834f111e9e8b9a6c49c295ca2c5a887865 Mon Sep 17 00:00:00 2001
From: alkhan <alkhan@tradestation.com>
Date: Tue, 28 Apr 2020 09:58:38 -0500
Subject: [PATCH 2/2] Make buffer size option consistent with other options

---
 conn.go | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/conn.go b/conn.go
index 11c8ab9..494cdcc 100644
--- a/conn.go
+++ b/conn.go
@@ -70,6 +70,9 @@ func Dial(network, addr string, opts ...func(*Conn) error) (*Conn, error) {
 // been created by the program. The opts parameter provides the
 // opportunity to specify STOMP protocol options.
 func Connect(conn io.ReadWriteCloser, opts ...func(*Conn) error) (*Conn, error) {
+	reader := frame.NewReader(conn)
+	writer := frame.NewWriter(conn)
+
 	c := &Conn{
 		conn:       conn,
 		closeMutex: &sync.Mutex{},
@@ -80,19 +83,12 @@ func Connect(conn io.ReadWriteCloser, opts ...func(*Conn) error) (*Conn, error)
 		return nil, err
 	}
 
-	var reader *frame.Reader
-	var writer *frame.Writer
-
 	if options.ReadBufferSize > 0 {
 		reader = frame.NewReaderSize(conn, options.ReadBufferSize)
-	} else {
-		reader = frame.NewReader(conn)
 	}
 
 	if options.WriteBufferSize > 0 {
 		writer = frame.NewWriterSize(conn, options.ReadBufferSize)
-	} else {
-		writer = frame.NewWriter(conn)
 	}
 
 	readChannelCapacity := 20