From b258b4fadb573ac412f187b9f31974ea99d32f50 Mon Sep 17 00:00:00 2001 From: Gary Burd Date: Thu, 2 Mar 2017 14:46:13 -0800 Subject: [PATCH] Use bufio.Writer returned from hijack in upgrade Reuse the buffer backing the bufio.Writer returned from hijack if that buffer is large enough to be generally useful and Upgrader.WriteBufferSize == 0. Update the logic for reusing bufio.Reader returned from hijack to match the logic for bufio.Reader: The buffer backing the reader must be sufficiently large to be generally useful and Upgrader.ReadBufferSize == 0. Improve the documentation for ReadBufferSize and WriterBufferSize in Dialer and Upgrader. --- client.go | 5 +++-- conn.go | 53 ++++++++++++++++++++++++++++++++++++++-------------- conn_test.go | 30 +++++++++++++++++++++++------ server.go | 5 +++-- 4 files changed, 69 insertions(+), 24 deletions(-) diff --git a/client.go b/client.go index 1b0e69a..eb20640 100644 --- a/client.go +++ b/client.go @@ -66,8 +66,9 @@ type Dialer struct { // HandshakeTimeout specifies the duration for the handshake to complete. HandshakeTimeout time.Duration - // Input and output buffer sizes. If the buffer size is zero, then a - // default value of 4096 is used. + // ReadBufferSize and WriteBufferSize specify I/O buffer sizes. If a buffer + // size is zero, then a useful default size is used. The I/O buffer sizes + // do not limit the size of the messages that can be sent or received. ReadBufferSize, WriteBufferSize int // Subprotocols specifies the client's requested subprotocols. diff --git a/conn.go b/conn.go index 7acc58a..97e1dba 100644 --- a/conn.go +++ b/conn.go @@ -268,33 +268,58 @@ func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) return newConnBRW(conn, isServer, readBufferSize, writeBufferSize, nil) } +type writeHook struct { + p []byte +} + +func (wh *writeHook) Write(p []byte) (int, error) { + wh.p = p + return len(p), nil +} + func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, brw *bufio.ReadWriter) *Conn { mu := make(chan bool, 1) mu <- true - if readBufferSize == 0 { - readBufferSize = defaultReadBufferSize - } - if readBufferSize < maxControlFramePayloadSize { - readBufferSize = maxControlFramePayloadSize - } - - // Reuse the supplied brw.Reader if brw.Reader's buf is the requested size. var br *bufio.Reader - if brw != nil && brw.Reader != nil { - // This code assumes that peek on a reset reader returns + if readBufferSize == 0 && brw != nil && brw.Reader != nil { + // Reuse the supplied bufio.Reader if the buffer has a useful size. + // This code assumes that peek on a reader returns // bufio.Reader.buf[:0]. brw.Reader.Reset(conn) - if p, err := brw.Reader.Peek(0); err == nil && cap(p) == readBufferSize { + if p, err := brw.Reader.Peek(0); err == nil && cap(p) >= 256 { br = brw.Reader } } if br == nil { + if readBufferSize == 0 { + readBufferSize = defaultReadBufferSize + } + if readBufferSize < maxControlFramePayloadSize { + readBufferSize = maxControlFramePayloadSize + } br = bufio.NewReaderSize(conn, readBufferSize) } - if writeBufferSize == 0 { - writeBufferSize = defaultWriteBufferSize + var writeBuf []byte + if writeBufferSize == 0 && brw != nil && brw.Writer != nil { + // Use the bufio.Writer's buffer if the buffer has a useful size. This + // code assumes that bufio.Writer.buf[:1] is passed to the + // bufio.Writer's underlying writer. + var wh writeHook + brw.Writer.Reset(&wh) + brw.Writer.WriteByte(0) + brw.Flush() + if cap(wh.p) >= maxFrameHeaderSize+256 { + writeBuf = wh.p[:cap(wh.p)] + } + } + + if writeBuf == nil { + if writeBufferSize == 0 { + writeBufferSize = defaultWriteBufferSize + } + writeBuf = make([]byte, writeBufferSize+maxFrameHeaderSize) } c := &Conn{ @@ -303,7 +328,7 @@ func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize in conn: conn, mu: mu, readFinal: true, - writeBuf: make([]byte, writeBufferSize+maxFrameHeaderSize), + writeBuf: writeBuf, enableWriteCompression: true, compressionLevel: defaultCompressionLevel, } diff --git a/conn_test.go b/conn_test.go index ba347f8..06e9bc3 100644 --- a/conn_test.go +++ b/conn_test.go @@ -464,16 +464,34 @@ func TestFailedConnectionReadPanic(t *testing.T) { t.Fatal("should not get here") } -func TestBufioReaderReuse(t *testing.T) { - brw := bufio.NewReadWriter(bufio.NewReader(nil), nil) +func TestBufioReuse(t *testing.T) { + brw := bufio.NewReadWriter(bufio.NewReader(nil), bufio.NewWriter(nil)) c := newConnBRW(nil, false, 0, 0, brw) + if c.br != brw.Reader { t.Error("connection did not reuse bufio.Reader") } - brw = bufio.NewReadWriter(bufio.NewReaderSize(nil, 1234), nil) // size must not equal bufio.defaultBufSize - c = newConnBRW(nil, false, 0, 0, brw) - if c.br == brw.Reader { - t.Error("connection reuse bufio.Reader with wrong size") + var wh writeHook + brw.Writer.Reset(&wh) + brw.WriteByte(0) + brw.Flush() + if &c.writeBuf[0] != &wh.p[0] { + t.Error("connection did not reuse bufio.Writer") } + + brw = bufio.NewReadWriter(bufio.NewReaderSize(nil, 0), bufio.NewWriterSize(nil, 0)) + c = newConnBRW(nil, false, 0, 0, brw) + + if c.br == brw.Reader { + t.Error("connection used bufio.Reader with small size") + } + + brw.Writer.Reset(&wh) + brw.WriteByte(0) + brw.Flush() + if &c.writeBuf[0] != &wh.p[0] { + t.Error("connection used bufio.Writer with small size") + } + } diff --git a/server.go b/server.go index 95c1656..3495e0f 100644 --- a/server.go +++ b/server.go @@ -28,8 +28,9 @@ type Upgrader struct { HandshakeTimeout time.Duration // ReadBufferSize and WriteBufferSize specify I/O buffer sizes. If a buffer - // size is zero, then a default value of 4096 is used. The I/O buffer sizes - // do not limit the size of the messages that can be sent or received. + // size is zero, then buffers allocated by the HTTP server are used. The + // I/O buffer sizes do not limit the size of the messages that can be sent + // or received. ReadBufferSize, WriteBufferSize int // Subprotocols specifies the server's supported protocols in order of