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.

Author: garyburd

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.

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,
 }

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
+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 reuse bufio.Reader with wrong size")
+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")
+}
+
 }

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