buffer WebSocket reads to reduce syscalls
Wrap the WebSocket connection in an 8KB bufio.Reader so that small reads (2-byte frame headers, length extensions) are served from buffer instead of making individual syscalls.
This commit is contained in:
parent
c57addc6c7
commit
fbaf7a07dc
2 changed files with 36 additions and 29 deletions
14
websocket.go
14
websocket.go
|
|
@ -1,6 +1,7 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"encoding/base64"
|
||||
|
|
@ -24,6 +25,7 @@ const (
|
|||
// WSConn is a minimal RFC 6455 WebSocket client.
|
||||
type WSConn struct {
|
||||
conn io.ReadWriteCloser
|
||||
reader *bufio.Reader
|
||||
mu sync.Mutex // serializes writes
|
||||
}
|
||||
|
||||
|
|
@ -90,7 +92,7 @@ func WSDial(rawURL string) (*WSConn, error) {
|
|||
return nil, fmt.Errorf("handshake failed: %s", strings.SplitN(resp, "\r\n", 2)[0])
|
||||
}
|
||||
|
||||
return &WSConn{conn: conn}, nil
|
||||
return &WSConn{conn: conn, reader: bufio.NewReaderSize(conn, 8192)}, nil
|
||||
}
|
||||
|
||||
// ReadMessage reads the next text or binary message, handling control frames internally.
|
||||
|
|
@ -127,7 +129,7 @@ func (ws *WSConn) Close() error {
|
|||
|
||||
func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
|
||||
var header [2]byte
|
||||
if _, err = io.ReadFull(ws.conn, header[:]); err != nil {
|
||||
if _, err = io.ReadFull(ws.reader, header[:]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
|
||||
|
|
@ -138,13 +140,13 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
|
|||
switch length {
|
||||
case 126:
|
||||
var ext [2]byte
|
||||
if _, err = io.ReadFull(ws.conn, ext[:]); err != nil {
|
||||
if _, err = io.ReadFull(ws.reader, ext[:]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
length = uint64(binary.BigEndian.Uint16(ext[:]))
|
||||
case 127:
|
||||
var ext [8]byte
|
||||
if _, err = io.ReadFull(ws.conn, ext[:]); err != nil {
|
||||
if _, err = io.ReadFull(ws.reader, ext[:]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
length = binary.BigEndian.Uint64(ext[:])
|
||||
|
|
@ -152,14 +154,14 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
|
|||
|
||||
var maskKey [4]byte
|
||||
if masked {
|
||||
if _, err = io.ReadFull(ws.conn, maskKey[:]); err != nil {
|
||||
if _, err = io.ReadFull(ws.reader, maskKey[:]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
}
|
||||
|
||||
payload = make([]byte, length)
|
||||
if length > 0 {
|
||||
if _, err = io.ReadFull(ws.conn, payload); err != nil {
|
||||
if _, err = io.ReadFull(ws.reader, payload); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
|
|
@ -13,7 +14,7 @@ import (
|
|||
|
||||
func TestWriteFrameMasked(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
payload := []byte("hello")
|
||||
if err := ws.writeFrame(opText, payload); err != nil {
|
||||
|
|
@ -43,7 +44,7 @@ func TestWriteFrameMasked(t *testing.T) {
|
|||
|
||||
func TestWriteFrame16BitLength(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
payload := make([]byte, 300)
|
||||
if err := ws.writeFrame(opBinary, payload); err != nil {
|
||||
|
|
@ -62,7 +63,7 @@ func TestWriteFrame16BitLength(t *testing.T) {
|
|||
|
||||
func TestWriteFrame64BitLength(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
payload := make([]byte, 70000) // > 65535, uses 8-byte extended
|
||||
if err := ws.writeFrame(opBinary, payload); err != nil {
|
||||
|
|
@ -86,7 +87,7 @@ func TestReadFrame(t *testing.T) {
|
|||
buf.WriteByte(byte(len(payload)))
|
||||
buf.Write(payload)
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
opcode, data, err := ws.readFrame()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -112,7 +113,7 @@ func TestReadFrame16BitLength(t *testing.T) {
|
|||
buf.Write(extLen[:])
|
||||
buf.Write(payload)
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
opcode, data, err := ws.readFrame()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -135,7 +136,7 @@ func TestReadFrame64BitLength(t *testing.T) {
|
|||
buf.Write(extLen[:])
|
||||
buf.Write(payload)
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
opcode, data, err := ws.readFrame()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -162,7 +163,7 @@ func TestReadFrameMasked(t *testing.T) {
|
|||
buf.Write(maskKey[:])
|
||||
buf.Write(masked)
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
opcode, data, err := ws.readFrame()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -182,7 +183,7 @@ func TestReadMessageClose(t *testing.T) {
|
|||
buf.WriteByte(0) // no payload
|
||||
|
||||
rw := &captureWriter{Reader: &buf}
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: rw}}
|
||||
ws := testWSConn(&nopCloser{readWriter: rw})
|
||||
|
||||
_, opcode, err := ws.ReadMessage()
|
||||
if err != io.EOF {
|
||||
|
|
@ -211,7 +212,7 @@ func TestReadFramePingPong(t *testing.T) {
|
|||
buf.Write(text)
|
||||
|
||||
rw := &captureWriter{Reader: &buf}
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: rw}}
|
||||
ws := testWSConn(&nopCloser{readWriter: rw})
|
||||
|
||||
data, _, err := ws.ReadMessage()
|
||||
if err != nil {
|
||||
|
|
@ -231,7 +232,7 @@ func TestReadFrame16BitLengthError(t *testing.T) {
|
|||
buf.WriteByte(0x81) // FIN + text
|
||||
buf.WriteByte(126) // 16-bit length indicator, but no length bytes follow
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
_, _, err := ws.readFrame()
|
||||
if err == nil {
|
||||
t.Error("expected error for truncated 16-bit length")
|
||||
|
|
@ -243,7 +244,7 @@ func TestReadFrame64BitLengthError(t *testing.T) {
|
|||
buf.WriteByte(0x81) // FIN + text
|
||||
buf.WriteByte(127) // 64-bit length indicator, but no length bytes follow
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
_, _, err := ws.readFrame()
|
||||
if err == nil {
|
||||
t.Error("expected error for truncated 64-bit length")
|
||||
|
|
@ -255,7 +256,7 @@ func TestReadFrameMaskKeyError(t *testing.T) {
|
|||
buf.WriteByte(0x81) // FIN + text
|
||||
buf.WriteByte(0x80 | 0x01) // masked + length 1, but no mask key or payload
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
_, _, err := ws.readFrame()
|
||||
if err == nil {
|
||||
t.Error("expected error for missing mask key")
|
||||
|
|
@ -267,7 +268,7 @@ func TestReadFramePayloadError(t *testing.T) {
|
|||
buf.WriteByte(0x81) // FIN + text
|
||||
buf.WriteByte(5) // length 5, but no payload
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
_, _, err := ws.readFrame()
|
||||
if err == nil {
|
||||
t.Error("expected error for truncated payload")
|
||||
|
|
@ -275,7 +276,7 @@ func TestReadFramePayloadError(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestWriteFrameHeaderError(t *testing.T) {
|
||||
ws := &WSConn{conn: &failOnWriteBuffer{Buffer: bytes.NewBuffer(nil)}}
|
||||
ws := testWSConn(&failOnWriteBuffer{Buffer: bytes.NewBuffer(nil)})
|
||||
err := ws.writeFrame(opText, []byte("hello"))
|
||||
if err == nil {
|
||||
t.Error("expected write error")
|
||||
|
|
@ -284,7 +285,7 @@ func TestWriteFrameHeaderError(t *testing.T) {
|
|||
|
||||
func TestWriteFramePayloadError(t *testing.T) {
|
||||
// Writer that succeeds for header but fails for payload
|
||||
ws := &WSConn{conn: &failOnNthWrite{failAfter: 1}}
|
||||
ws := testWSConn(&failOnNthWrite{failAfter: 1})
|
||||
err := ws.writeFrame(opText, []byte("hello"))
|
||||
if err == nil {
|
||||
t.Error("expected write error on payload")
|
||||
|
|
@ -310,7 +311,7 @@ func (f *failOnNthWrite) Close() error { return nil }
|
|||
func TestReadMessageError(t *testing.T) {
|
||||
// Empty buffer causes immediate EOF on readFrame
|
||||
buf := bytes.NewBuffer(nil)
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: buf})
|
||||
_, _, err := ws.ReadMessage()
|
||||
if err == nil {
|
||||
t.Error("expected error for empty buffer")
|
||||
|
|
@ -323,7 +324,7 @@ func TestReadMessagePongError(t *testing.T) {
|
|||
buf.WriteByte(0x80 | byte(opPing))
|
||||
buf.WriteByte(0)
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &failOnWriteBuffer{Buffer: &buf}}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &failOnWriteBuffer{Buffer: &buf}})
|
||||
_, _, err := ws.ReadMessage()
|
||||
if err == nil {
|
||||
t.Error("expected pong write error")
|
||||
|
|
@ -335,7 +336,7 @@ func TestReadFrameEmptyPayload(t *testing.T) {
|
|||
buf.WriteByte(0x82) // FIN + binary
|
||||
buf.WriteByte(0) // zero length
|
||||
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
opcode, data, err := ws.readFrame()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -350,7 +351,7 @@ func TestReadFrameEmptyPayload(t *testing.T) {
|
|||
|
||||
func TestWriteFrameEmptyPayload(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
if err := ws.writeFrame(opClose, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -378,7 +379,7 @@ func (f *failOnWriteBuffer) Close() error { return nil }
|
|||
|
||||
func TestWriteText(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
if err := ws.WriteText([]byte("hello")); err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -393,7 +394,7 @@ func TestWriteText(t *testing.T) {
|
|||
|
||||
func TestClose(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
||||
ws := testWSConn(&nopCloser{readWriter: &buf})
|
||||
|
||||
if err := ws.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -602,6 +603,10 @@ func TestWSDialRealHTTPServer(t *testing.T) {
|
|||
_ = ws.Close()
|
||||
}
|
||||
|
||||
func testWSConn(rw io.ReadWriteCloser) *WSConn {
|
||||
return &WSConn{conn: rw, reader: bufio.NewReader(rw)}
|
||||
}
|
||||
|
||||
// nopCloser wraps a ReadWriter with a no-op Close.
|
||||
type nopCloser struct {
|
||||
readWriter interface {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue