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