WIP cleanup

Author: nhooyr

compress_notjs.go

@@ -3,10 +3,11 @@
 package websocket
 
 import (
-	"compress/flate"
 	"io"
 	"net/http"
 	"sync"
+
+	"github.com/klauspost/compress/flate"
 )
 
 func (m CompressionMode) opts() *compressionOptions {
@@ -45,10 +46,16 @@ type trimLastFourBytesWriter struct {
 }
 
 func (tw *trimLastFourBytesWriter) reset() {
-	tw.tail = tw.tail[:0]
+	if tw != nil && tw.tail != nil {
+		tw.tail = tw.tail[:0]
+	}
 }
 
 func (tw *trimLastFourBytesWriter) Write(p []byte) (int, error) {
+	if tw.tail == nil {
+		tw.tail = make([]byte, 0, 4)
+	}
+
 	extra := len(tw.tail) + len(p) - 4
 
 	if extra <= 0 {
@@ -65,7 +72,10 @@ func (tw *trimLastFourBytesWriter) Write(p []byte) (int, error) {
 		if err != nil {
 			return 0, err
 		}
-		tw.tail = tw.tail[extra:]
+
+		// Shift remaining bytes in tail over.
+		n := copy(tw.tail, tw.tail[extra:])
+		tw.tail = tw.tail[:n]
 	}
 
 	// If p is less than or equal to 4 bytes,
@@ -118,22 +128,32 @@ type slidingWindow struct {
 	buf []byte
 }
 
-var swPoolMu sync.Mutex
+var swPoolMu sync.RWMutex
 var swPool = map[int]*sync.Pool{}
 
-func (sw *slidingWindow) init(n int) {
-	if sw.buf != nil {
-		return
+func slidingWindowPool(n int) *sync.Pool {
+	swPoolMu.RLock()
+	p, ok := swPool[n]
+	swPoolMu.RUnlock()
+	if ok {
+		return p
 	}
 
+	p = &sync.Pool{}
+
 	swPoolMu.Lock()
-	defer swPoolMu.Unlock()
+	swPool[n] = p
+	swPoolMu.Unlock()
 
-	p, ok := swPool[n]
-	if !ok {
-		p = &sync.Pool{}
-		swPool[n] = p
+	return p
+}
+
+func (sw *slidingWindow) init(n int) {
+	if sw.buf != nil {
+		return
 	}
+
+	p := slidingWindowPool(n)
 	buf, ok := p.Get().([]byte)
 	if ok {
 		sw.buf = buf[:0]

conn_notjs.go

@@ -39,16 +39,17 @@ type Conn struct {
 
 	// Read state.
 	readMu            *mu
-	readHeader        header
+	readHeaderBuf     [8]byte
 	readControlBuf    [maxControlPayload]byte
 	msgReader         *msgReader
 	readCloseFrameErr error
 
 	// Write state.
-	msgWriter    *msgWriter
-	writeFrameMu *mu
-	writeBuf     []byte
-	writeHeader  header
+	msgWriterState *msgWriterState
+	writeFrameMu   *mu
+	writeBuf       []byte
+	writeHeaderBuf [8]byte
+	writeHeader    header
 
 	closed     chan struct{}
 	closeMu    sync.Mutex
@@ -94,14 +95,14 @@ func newConn(cfg connConfig) *Conn {
 
 	c.msgReader = newMsgReader(c)
 
-	c.msgWriter = newMsgWriter(c)
+	c.msgWriterState = newMsgWriterState(c)
 	if c.client {
 		c.writeBuf = extractBufioWriterBuf(c.bw, c.rwc)
 	}
 
 	if c.flate() && c.flateThreshold == 0 {
 		c.flateThreshold = 256
-		if !c.msgWriter.flateContextTakeover() {
+		if !c.msgWriterState.flateContextTakeover() {
 			c.flateThreshold = 512
 		}
 	}
@@ -142,7 +143,7 @@ func (c *Conn) close(err error) {
 			c.writeFrameMu.Lock(context.Background())
 			putBufioWriter(c.bw)
 		}
-		c.msgWriter.close()
+		c.msgWriterState.close()
 
 		c.msgReader.close()
 		if c.client {

conn_test.go

@@ -5,14 +5,14 @@ package websocket_test
 import (
 	"bytes"
 	"context"
-	"crypto/rand"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"os/exec"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -379,60 +379,52 @@ func BenchmarkConn(b *testing.B) {
 		mode websocket.CompressionMode
 	}{
 		{
-			name: "compressionDisabled",
+			name: "disabledCompress",
 			mode: websocket.CompressionDisabled,
 		},
 		{
-			name: "compression",
+			name: "compress",
 			mode: websocket.CompressionContextTakeover,
 		},
 		{
-			name: "noContextCompression",
+			name: "compressNoContext",
 			mode: websocket.CompressionNoContextTakeover,
 		},
 	}
 	for _, bc := range benchCases {
 		b.Run(bc.name, func(b *testing.B) {
 			bb, c1, c2 := newConnTest(b, &websocket.DialOptions{
 				CompressionOptions: &websocket.CompressionOptions{Mode: bc.mode},
-			}, nil)
+			}, &websocket.AcceptOptions{
+				CompressionOptions: &websocket.CompressionOptions{Mode: bc.mode},
+			})
 			defer bb.cleanup()
 
 			bb.goEchoLoop(c2)
 
-			const n = 32768
-			writeBuf := make([]byte, n)
-			readBuf := make([]byte, n)
-			writes := make(chan websocket.MessageType)
+			msg := []byte(strings.Repeat("1234", 128))
+			readBuf := make([]byte, len(msg))
+			writes := make(chan struct{})
 			defer close(writes)
 			werrs := make(chan error)
 
 			go func() {
-				for typ := range writes {
-					werrs <- c1.Write(bb.ctx, typ, writeBuf)
+				for range writes {
+					werrs <- c1.Write(bb.ctx, websocket.MessageText, msg)
 				}
 			}()
-			b.SetBytes(n)
+			b.SetBytes(int64(len(msg)))
 			b.ReportAllocs()
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				_, err := rand.Reader.Read(writeBuf)
-				if err != nil {
-					b.Fatal(err)
-				}
-
-				expType := websocket.MessageBinary
-				if writeBuf[0]%2 == 1 {
-					expType = websocket.MessageText
-				}
-				writes <- expType
+				writes <- struct{}{}
 
 				typ, r, err := c1.Reader(bb.ctx)
 				if err != nil {
 					b.Fatal(err)
 				}
-				if expType != typ {
-					assert.Equal(b, "data type", expType, typ)
+				if websocket.MessageText != typ {
+					assert.Equal(b, "data type", websocket.MessageText, typ)
 				}
 
 				_, err = io.ReadFull(r, readBuf)
@@ -448,8 +440,8 @@ func BenchmarkConn(b *testing.B) {
 					assert.Equal(b, "n2", 0, n2)
 				}
 
-				if !bytes.Equal(writeBuf, readBuf) {
-					assert.Equal(b, "msg", writeBuf, readBuf)
+				if !bytes.Equal(msg, readBuf) {
+					assert.Equal(b, "msg", msg, readBuf)
 				}
 
 				err = <-werrs
@@ -464,3 +456,8 @@ func BenchmarkConn(b *testing.B) {
 		})
 	}
 }
+
+func TestCompression(t *testing.T) {
+	t.Parallel()
+
+}

frame.go

@@ -3,9 +3,12 @@ package websocket
 import (
 	"bufio"
 	"encoding/binary"
+	"io"
 	"math"
 	"math/bits"
 
+	"golang.org/x/xerrors"
+
 	"nhooyr.io/websocket/internal/errd"
 )
 
@@ -46,12 +49,12 @@ type header struct {
 
 // readFrameHeader reads a header from the reader.
 // See https://tools.ietf.org/html/rfc6455#section-5.2.
-func readFrameHeader(h *header, r *bufio.Reader) (err error) {
+func readFrameHeader(r *bufio.Reader, readBuf []byte) (h header, err error) {
 	defer errd.Wrap(&err, "failed to read frame header")
 
 	b, err := r.ReadByte()
 	if err != nil {
-		return err
+		return header{}, err
 	}
 
 	h.fin = b&(1<<7) != 0
@@ -63,7 +66,7 @@ func readFrameHeader(h *header, r *bufio.Reader) (err error) {
 
 	b, err = r.ReadByte()
 	if err != nil {
-		return err
+		return header{}, err
 	}
 
 	h.masked = b&(1<<7) != 0
@@ -73,24 +76,29 @@ func readFrameHeader(h *header, r *bufio.Reader) (err error) {
 	case payloadLength < 126:
 		h.payloadLength = int64(payloadLength)
 	case payloadLength == 126:
-		var pl uint16
-		err = binary.Read(r, binary.BigEndian, &pl)
-		h.payloadLength = int64(pl)
+		_, err = io.ReadFull(r, readBuf[:2])
+		h.payloadLength = int64(binary.BigEndian.Uint16(readBuf))
 	case payloadLength == 127:
-		err = binary.Read(r, binary.BigEndian, &h.payloadLength)
+		_, err = io.ReadFull(r, readBuf)
+		h.payloadLength = int64(binary.BigEndian.Uint64(readBuf))
 	}
 	if err != nil {
-		return err
+		return header{}, err
+	}
+
+	if h.payloadLength < 0 {
+		return header{}, xerrors.Errorf("received negative payload length: %v", h.payloadLength)
 	}
 
 	if h.masked {
-		err = binary.Read(r, binary.LittleEndian, &h.maskKey)
+		_, err = io.ReadFull(r, readBuf[:4])
 		if err != nil {
-			return err
+			return header{}, err
 		}
+		h.maskKey = binary.LittleEndian.Uint32(readBuf)
 	}
 
-	return nil
+	return h, nil
 }
 
 // maxControlPayload is the maximum length of a control frame payload.
@@ -99,7 +107,7 @@ const maxControlPayload = 125
 
 // writeFrameHeader writes the bytes of the header to w.
 // See https://tools.ietf.org/html/rfc6455#section-5.2
-func writeFrameHeader(h header, w *bufio.Writer) (err error) {
+func writeFrameHeader(h header, w *bufio.Writer, buf []byte) (err error) {
 	defer errd.Wrap(&err, "failed to write frame header")
 
 	var b byte
@@ -143,16 +151,19 @@ func writeFrameHeader(h header, w *bufio.Writer) (err error) {
 
 	switch {
 	case h.payloadLength > math.MaxUint16:
-		err = binary.Write(w, binary.BigEndian, h.payloadLength)
+		binary.BigEndian.PutUint64(buf, uint64(h.payloadLength))
+		_, err = w.Write(buf)
 	case h.payloadLength > 125:
-		err = binary.Write(w, binary.BigEndian, uint16(h.payloadLength))
+		binary.BigEndian.PutUint16(buf, uint16(h.payloadLength))
+		_, err = w.Write(buf[:2])
 	}
 	if err != nil {
 		return err
 	}
 
 	if h.masked {
-		err = binary.Write(w, binary.LittleEndian, h.maskKey)
+		binary.LittleEndian.PutUint32(buf, h.maskKey)
+		_, err = w.Write(buf[:4])
 		if err != nil {
 			return err
 		}

frame_test.go

@@ -80,14 +80,13 @@ func testHeader(t *testing.T, h header) {
 	w := bufio.NewWriter(b)
 	r := bufio.NewReader(b)
 
-	err := writeFrameHeader(h, w)
+	err := writeFrameHeader(h, w, make([]byte, 8))
 	assert.Success(t, err)
 
 	err = w.Flush()
 	assert.Success(t, err)
 
-	var h2 header
-	err = readFrameHeader(&h2, r)
+	h2, err := readFrameHeader(r, make([]byte, 8))
 	assert.Success(t, err)
 
 	assert.Equal(t, "read header", h, h2)

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/golang/protobuf v1.3.3
 	github.com/google/go-cmp v0.4.0
 	github.com/gorilla/websocket v1.4.1
+	github.com/klauspost/compress v1.10.0
 	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543
 )

go.sum

@@ -10,6 +10,8 @@ github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/klauspost/compress v1.10.0 h1:92XGj1AcYzA6UrVdd4qIIBrT8OroryvRvdmg/IfmC7Y=
+github.com/klauspost/compress v1.10.0/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=

internal/xsync/go.go

@@ -6,7 +6,7 @@ import (
 
 // Go allows running a function in another goroutine
 // and waiting for its error.
-func Go(fn func() error) <- chan error {
+func Go(fn func() error) <-chan error {
 	errs := make(chan error, 1)
 	go func() {
 		defer func() {