graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

buffer WebSocket reads to reduce syscalls

Browse source 
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:
Graham McIntire 2026-02-12 09:53:16 -06:00
parent c57addc6c7
commit fbaf7a07dc
No known key found for this signature in database
2 changed files with 36 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
websocket.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"bufio"
"crypto/rand" "crypto/rand"
"crypto/tls" "crypto/tls"
"encoding/base64" "encoding/base64"
@ -23,8 +24,9 @@ const (
// WSConn is a minimal RFC 6455 WebSocket client. // WSConn is a minimal RFC 6455 WebSocket client.
type WSConn struct { type WSConn struct {
conn io.ReadWriteCloser conn io.ReadWriteCloser
mu sync.Mutex // serializes writes reader *bufio.Reader
mu sync.Mutex // serializes writes
} }
var randRead = rand.Read var randRead = rand.Read
@ -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 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. // 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) { func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
var header [2]byte 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 return 0, nil, err
} }
@ -138,13 +140,13 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
switch length { switch length {
case 126: case 126:
var ext [2]byte 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 return 0, nil, err
} }
length = uint64(binary.BigEndian.Uint16(ext[:])) length = uint64(binary.BigEndian.Uint16(ext[:]))
case 127: case 127:
var ext [8]byte 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 return 0, nil, err
} }
length = binary.BigEndian.Uint64(ext[:]) length = binary.BigEndian.Uint64(ext[:])
@ -152,14 +154,14 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
var maskKey [4]byte var maskKey [4]byte
if masked { if masked {
if _, err = io.ReadFull(ws.conn, maskKey[:]); err != nil { if _, err = io.ReadFull(ws.reader, maskKey[:]); err != nil {
return 0, nil, err return 0, nil, err
} }
} }
payload = make([]byte, length) payload = make([]byte, length)
if length > 0 { 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 return 0, nil, err
} }
} }

47
websocket_test.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"bufio"
"bytes" "bytes"
"encoding/binary" "encoding/binary"
"fmt" "fmt"
@ -13,7 +14,7 @@ import (
func TestWriteFrameMasked(t *testing.T) { func TestWriteFrameMasked(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
payload := []byte("hello") payload := []byte("hello")
if err := ws.writeFrame(opText, payload); err != nil { if err := ws.writeFrame(opText, payload); err != nil {
@ -43,7 +44,7 @@ func TestWriteFrameMasked(t *testing.T) {
func TestWriteFrame16BitLength(t *testing.T) { func TestWriteFrame16BitLength(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
payload := make([]byte, 300) payload := make([]byte, 300)
if err := ws.writeFrame(opBinary, payload); err != nil { if err := ws.writeFrame(opBinary, payload); err != nil {
@ -62,7 +63,7 @@ func TestWriteFrame16BitLength(t *testing.T) {
func TestWriteFrame64BitLength(t *testing.T) { func TestWriteFrame64BitLength(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
payload := make([]byte, 70000) // > 65535, uses 8-byte extended payload := make([]byte, 70000) // > 65535, uses 8-byte extended
if err := ws.writeFrame(opBinary, payload); err != nil { if err := ws.writeFrame(opBinary, payload); err != nil {
@ -86,7 +87,7 @@ func TestReadFrame(t *testing.T) {
buf.WriteByte(byte(len(payload))) buf.WriteByte(byte(len(payload)))
buf.Write(payload) buf.Write(payload)
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
opcode, data, err := ws.readFrame() opcode, data, err := ws.readFrame()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -112,7 +113,7 @@ func TestReadFrame16BitLength(t *testing.T) {
buf.Write(extLen[:]) buf.Write(extLen[:])
buf.Write(payload) buf.Write(payload)
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
opcode, data, err := ws.readFrame() opcode, data, err := ws.readFrame()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -135,7 +136,7 @@ func TestReadFrame64BitLength(t *testing.T) {
buf.Write(extLen[:]) buf.Write(extLen[:])
buf.Write(payload) buf.Write(payload)
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
opcode, data, err := ws.readFrame() opcode, data, err := ws.readFrame()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -162,7 +163,7 @@ func TestReadFrameMasked(t *testing.T) {
buf.Write(maskKey[:]) buf.Write(maskKey[:])
buf.Write(masked) buf.Write(masked)
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
opcode, data, err := ws.readFrame() opcode, data, err := ws.readFrame()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -182,7 +183,7 @@ func TestReadMessageClose(t *testing.T) {
buf.WriteByte(0) // no payload buf.WriteByte(0) // no payload
rw := &captureWriter{Reader: &buf} rw := &captureWriter{Reader: &buf}
ws := &WSConn{conn: &nopCloser{readWriter: rw}} ws := testWSConn(&nopCloser{readWriter: rw})
_, opcode, err := ws.ReadMessage() _, opcode, err := ws.ReadMessage()
if err != io.EOF { if err != io.EOF {
@ -211,7 +212,7 @@ func TestReadFramePingPong(t *testing.T) {
buf.Write(text) buf.Write(text)
rw := &captureWriter{Reader: &buf} rw := &captureWriter{Reader: &buf}
ws := &WSConn{conn: &nopCloser{readWriter: rw}} ws := testWSConn(&nopCloser{readWriter: rw})
data, _, err := ws.ReadMessage() data, _, err := ws.ReadMessage()
if err != nil { if err != nil {
@ -231,7 +232,7 @@ func TestReadFrame16BitLengthError(t *testing.T) {
buf.WriteByte(0x81) // FIN + text buf.WriteByte(0x81) // FIN + text
buf.WriteByte(126) // 16-bit length indicator, but no length bytes follow 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() _, _, err := ws.readFrame()
if err == nil { if err == nil {
t.Error("expected error for truncated 16-bit length") 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(0x81) // FIN + text
buf.WriteByte(127) // 64-bit length indicator, but no length bytes follow 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() _, _, err := ws.readFrame()
if err == nil { if err == nil {
t.Error("expected error for truncated 64-bit length") 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(0x81) // FIN + text
buf.WriteByte(0x80 | 0x01) // masked + length 1, but no mask key or payload 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() _, _, err := ws.readFrame()
if err == nil { if err == nil {
t.Error("expected error for missing mask key") t.Error("expected error for missing mask key")
@ -267,7 +268,7 @@ func TestReadFramePayloadError(t *testing.T) {
buf.WriteByte(0x81) // FIN + text buf.WriteByte(0x81) // FIN + text
buf.WriteByte(5) // length 5, but no payload buf.WriteByte(5) // length 5, but no payload
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
_, _, err := ws.readFrame() _, _, err := ws.readFrame()
if err == nil { if err == nil {
t.Error("expected error for truncated payload") t.Error("expected error for truncated payload")
@ -275,7 +276,7 @@ func TestReadFramePayloadError(t *testing.T) {
} }
func TestWriteFrameHeaderError(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")) err := ws.writeFrame(opText, []byte("hello"))
if err == nil { if err == nil {
t.Error("expected write error") t.Error("expected write error")
@ -284,7 +285,7 @@ func TestWriteFrameHeaderError(t *testing.T) {
func TestWriteFramePayloadError(t *testing.T) { func TestWriteFramePayloadError(t *testing.T) {
// Writer that succeeds for header but fails for payload // 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")) err := ws.writeFrame(opText, []byte("hello"))
if err == nil { if err == nil {
t.Error("expected write error on payload") t.Error("expected write error on payload")
@ -310,7 +311,7 @@ func (f *failOnNthWrite) Close() error { return nil }
func TestReadMessageError(t *testing.T) { func TestReadMessageError(t *testing.T) {
// Empty buffer causes immediate EOF on readFrame // Empty buffer causes immediate EOF on readFrame
buf := bytes.NewBuffer(nil) buf := bytes.NewBuffer(nil)
ws := &WSConn{conn: &nopCloser{readWriter: buf}} ws := testWSConn(&nopCloser{readWriter: buf})
_, _, err := ws.ReadMessage() _, _, err := ws.ReadMessage()
if err == nil { if err == nil {
t.Error("expected error for empty buffer") t.Error("expected error for empty buffer")
@ -323,7 +324,7 @@ func TestReadMessagePongError(t *testing.T) {
buf.WriteByte(0x80 | byte(opPing)) buf.WriteByte(0x80 | byte(opPing))
buf.WriteByte(0) buf.WriteByte(0)
ws := &WSConn{conn: &nopCloser{readWriter: &failOnWriteBuffer{Buffer: &buf}}} ws := testWSConn(&nopCloser{readWriter: &failOnWriteBuffer{Buffer: &buf}})
_, _, err := ws.ReadMessage() _, _, err := ws.ReadMessage()
if err == nil { if err == nil {
t.Error("expected pong write error") t.Error("expected pong write error")
@ -335,7 +336,7 @@ func TestReadFrameEmptyPayload(t *testing.T) {
buf.WriteByte(0x82) // FIN + binary buf.WriteByte(0x82) // FIN + binary
buf.WriteByte(0) // zero length buf.WriteByte(0) // zero length
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
opcode, data, err := ws.readFrame() opcode, data, err := ws.readFrame()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -350,7 +351,7 @@ func TestReadFrameEmptyPayload(t *testing.T) {
func TestWriteFrameEmptyPayload(t *testing.T) { func TestWriteFrameEmptyPayload(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
if err := ws.writeFrame(opClose, nil); err != nil { if err := ws.writeFrame(opClose, nil); err != nil {
t.Fatal(err) t.Fatal(err)
@ -378,7 +379,7 @@ func (f *failOnWriteBuffer) Close() error { return nil }
func TestWriteText(t *testing.T) { func TestWriteText(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
if err := ws.WriteText([]byte("hello")); err != nil { if err := ws.WriteText([]byte("hello")); err != nil {
t.Fatal(err) t.Fatal(err)
@ -393,7 +394,7 @@ func TestWriteText(t *testing.T) {
func TestClose(t *testing.T) { func TestClose(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
ws := &WSConn{conn: &nopCloser{readWriter: &buf}} ws := testWSConn(&nopCloser{readWriter: &buf})
if err := ws.Close(); err != nil { if err := ws.Close(); err != nil {
t.Fatal(err) t.Fatal(err)
@ -602,6 +603,10 @@ func TestWSDialRealHTTPServer(t *testing.T) {
_ = ws.Close() _ = 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. // nopCloser wraps a ReadWriter with a no-op Close.
type nopCloser struct { type nopCloser struct {
readWriter interface { readWriter interface {