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.
This commit is contained in:
Gary Burd 2017-03-02 14:46:13 -08:00
parent 4873052237
commit b258b4fadb
4 changed files with 69 additions and 24 deletions

View File

@ -66,8 +66,9 @@ type Dialer struct {
// HandshakeTimeout specifies the duration for the handshake to complete. // HandshakeTimeout specifies the duration for the handshake to complete.
HandshakeTimeout time.Duration HandshakeTimeout time.Duration
// Input and output buffer sizes. If the buffer size is zero, then a // ReadBufferSize and WriteBufferSize specify I/O buffer sizes. If a buffer
// default value of 4096 is used. // 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 ReadBufferSize, WriteBufferSize int
// Subprotocols specifies the client's requested subprotocols. // Subprotocols specifies the client's requested subprotocols.

51
conn.go
View File

@ -268,34 +268,59 @@ func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int)
return newConnBRW(conn, isServer, readBufferSize, writeBufferSize, nil) 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 { func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, brw *bufio.ReadWriter) *Conn {
mu := make(chan bool, 1) mu := make(chan bool, 1)
mu <- true mu <- true
var br *bufio.Reader
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) >= 256 {
br = brw.Reader
}
}
if br == nil {
if readBufferSize == 0 { if readBufferSize == 0 {
readBufferSize = defaultReadBufferSize readBufferSize = defaultReadBufferSize
} }
if readBufferSize < maxControlFramePayloadSize { if readBufferSize < maxControlFramePayloadSize {
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
// bufio.Reader.buf[:0].
brw.Reader.Reset(conn)
if p, err := brw.Reader.Peek(0); err == nil && cap(p) == readBufferSize {
br = brw.Reader
}
}
if br == nil {
br = bufio.NewReaderSize(conn, readBufferSize) br = bufio.NewReaderSize(conn, readBufferSize)
} }
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 { if writeBufferSize == 0 {
writeBufferSize = defaultWriteBufferSize writeBufferSize = defaultWriteBufferSize
} }
writeBuf = make([]byte, writeBufferSize+maxFrameHeaderSize)
}
c := &Conn{ c := &Conn{
isServer: isServer, isServer: isServer,
@ -303,7 +328,7 @@ func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize in
conn: conn, conn: conn,
mu: mu, mu: mu,
readFinal: true, readFinal: true,
writeBuf: make([]byte, writeBufferSize+maxFrameHeaderSize), writeBuf: writeBuf,
enableWriteCompression: true, enableWriteCompression: true,
compressionLevel: defaultCompressionLevel, compressionLevel: defaultCompressionLevel,
} }

View File

@ -464,16 +464,34 @@ func TestFailedConnectionReadPanic(t *testing.T) {
t.Fatal("should not get here") t.Fatal("should not get here")
} }
func TestBufioReaderReuse(t *testing.T) { func TestBufioReuse(t *testing.T) {
brw := bufio.NewReadWriter(bufio.NewReader(nil), nil) brw := bufio.NewReadWriter(bufio.NewReader(nil), bufio.NewWriter(nil))
c := newConnBRW(nil, false, 0, 0, brw) c := newConnBRW(nil, false, 0, 0, brw)
if c.br != brw.Reader { if c.br != brw.Reader {
t.Error("connection did not reuse bufio.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
c = newConnBRW(nil, false, 0, 0, brw) brw.Writer.Reset(&wh)
if c.br == brw.Reader { brw.WriteByte(0)
t.Error("connection reuse bufio.Reader with wrong size") 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")
}
} }

View File

@ -28,8 +28,9 @@ type Upgrader struct {
HandshakeTimeout time.Duration HandshakeTimeout time.Duration
// ReadBufferSize and WriteBufferSize specify I/O buffer sizes. If a buffer // 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 // size is zero, then buffers allocated by the HTTP server are used. The
// do not limit the size of the messages that can be sent or received. // I/O buffer sizes do not limit the size of the messages that can be sent
// or received.
ReadBufferSize, WriteBufferSize int ReadBufferSize, WriteBufferSize int
// Subprotocols specifies the server's supported protocols in order of // Subprotocols specifies the server's supported protocols in order of