diff --git a/conn.go b/conn.go index 9de9a74..494cdcc 100644 --- a/conn.go +++ b/conn.go @@ -83,6 +83,14 @@ func Connect(conn io.ReadWriteCloser, opts ...func(*Conn) error) (*Conn, error) return nil, err } + if options.ReadBufferSize > 0 { + reader = frame.NewReaderSize(conn, options.ReadBufferSize) + } + + if options.WriteBufferSize > 0 { + writer = frame.NewWriterSize(conn, options.ReadBufferSize) + } + 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 + } + } }