Add conn benchmark

Author: nhooyr

.gitignore

@@ -0,0 +1 @@
+websocket.test

autobahn_test.go

@@ -59,8 +59,6 @@ func TestAutobahn(t *testing.T) {
 		for i := 1; i <= cases; i++ {
 			i := i
 			t.Run("", func(t *testing.T) {
-				t.Parallel()
-
 				ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
 				defer cancel()
 

compress_notjs.go

@@ -118,13 +118,17 @@ type slidingWindow struct {
 	buf []byte
 }
 
+var swPoolMu sync.Mutex
 var swPool = map[int]*sync.Pool{}
 
 func (sw *slidingWindow) init(n int) {
 	if sw.buf != nil {
 		return
 	}
 
+	swPoolMu.Lock()
+	defer swPoolMu.Unlock()
+
 	p, ok := swPool[n]
 	if !ok {
 		p = &sync.Pool{}
@@ -143,6 +147,9 @@ func (sw *slidingWindow) close() {
 		return
 	}
 
+	swPoolMu.Lock()
+	defer swPoolMu.Unlock()
+
 	swPool[cap(sw.buf)].Put(sw.buf)
 	sw.buf = nil
 }

conn_notjs.go

@@ -39,6 +39,7 @@ type Conn struct {
 
 	// Read state.
 	readMu            *mu
+	readHeader        header
 	readControlBuf    [maxControlPayload]byte
 	msgReader         *msgReader
 	readCloseFrameErr error

conn_test.go

@@ -3,7 +3,9 @@
 package websocket_test
 
 import (
+	"bytes"
 	"context"
+	"crypto/rand"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -48,7 +50,7 @@ func TestConn(t *testing.T) {
 				}, &websocket.AcceptOptions{
 					CompressionOptions: copts(),
 				})
-				defer tt.done()
+				defer tt.cleanup()
 
 				tt.goEchoLoop(c2)
 
@@ -67,15 +69,15 @@ func TestConn(t *testing.T) {
 
 	t.Run("badClose", func(t *testing.T) {
 		tt, c1, _ := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		err := c1.Close(-1, "")
 		assert.Contains(t, err, "failed to marshal close frame: status code StatusCode(-1) cannot be set")
 	})
 
 	t.Run("ping", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		c1.CloseRead(tt.ctx)
 		c2.CloseRead(tt.ctx)
@@ -91,7 +93,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("badPing", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		c2.CloseRead(tt.ctx)
 
@@ -104,7 +106,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("concurrentWrite", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		tt.goDiscardLoop(c2)
 
@@ -129,7 +131,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("concurrentWriteError", func(t *testing.T) {
 		tt, c1, _ := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		_, err := c1.Writer(tt.ctx, websocket.MessageText)
 		assert.Success(t, err)
@@ -143,7 +145,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("netConn", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		n1 := websocket.NetConn(tt.ctx, c1, websocket.MessageBinary)
 		n2 := websocket.NetConn(tt.ctx, c2, websocket.MessageBinary)
@@ -179,7 +181,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("netConn/BadMsg", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		n1 := websocket.NetConn(tt.ctx, c1, websocket.MessageBinary)
 		n2 := websocket.NetConn(tt.ctx, c2, websocket.MessageText)
@@ -201,7 +203,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("wsjson", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		tt.goEchoLoop(c2)
 
@@ -227,7 +229,7 @@ func TestConn(t *testing.T) {
 
 	t.Run("wspb", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
-		defer tt.done()
+		defer tt.cleanup()
 
 		tt.goEchoLoop(c2)
 
@@ -297,14 +299,16 @@ func assertCloseStatus(exp websocket.StatusCode, err error) error {
 }
 
 type connTest struct {
-	t   *testing.T
+	t   testing.TB
 	ctx context.Context
 
 	doneFuncs []func()
 }
 
-func newConnTest(t *testing.T, dialOpts *websocket.DialOptions, acceptOpts *websocket.AcceptOptions) (tt *connTest, c1, c2 *websocket.Conn) {
-	t.Parallel()
+func newConnTest(t testing.TB, dialOpts *websocket.DialOptions, acceptOpts *websocket.AcceptOptions) (tt *connTest, c1, c2 *websocket.Conn) {
+	if t, ok := t.(*testing.T); ok {
+		t.Parallel()
+	}
 	t.Helper()
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
@@ -325,7 +329,7 @@ func (tt *connTest) appendDone(f func()) {
 	tt.doneFuncs = append(tt.doneFuncs, f)
 }
 
-func (tt *connTest) done() {
+func (tt *connTest) cleanup() {
 	for i := len(tt.doneFuncs) - 1; i >= 0; i-- {
 		tt.doneFuncs[i]()
 	}
@@ -368,3 +372,95 @@ func (tt *connTest) goDiscardLoop(c *websocket.Conn) {
 		}
 	})
 }
+
+func BenchmarkConn(b *testing.B) {
+	var benchCases = []struct {
+		name string
+		mode websocket.CompressionMode
+	}{
+		{
+			name: "compressionDisabled",
+			mode: websocket.CompressionDisabled,
+		},
+		{
+			name: "compression",
+			mode: websocket.CompressionContextTakeover,
+		},
+		{
+			name: "noContextCompression",
+			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)
+			defer bb.cleanup()
+
+			bb.goEchoLoop(c2)
+
+			const n = 32768
+			writeBuf := make([]byte, n)
+			readBuf := make([]byte, n)
+			writes := make(chan websocket.MessageType)
+			defer close(writes)
+			werrs := make(chan error)
+
+			go func() {
+				for typ := range writes {
+					werrs <- c1.Write(bb.ctx, typ, writeBuf)
+				}
+			}()
+			b.SetBytes(n)
+			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
+
+				typ, r, err := c1.Reader(bb.ctx)
+				if err != nil {
+					b.Fatal(err)
+				}
+				if expType != typ {
+					assert.Equal(b, "data type", expType, typ)
+				}
+
+				_, err = io.ReadFull(r, readBuf)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				n2, err := r.Read(readBuf)
+				if err != io.EOF {
+					assert.Equal(b, "read err", io.EOF, err)
+				}
+				if n2 != 0 {
+					assert.Equal(b, "n2", 0, n2)
+				}
+
+				if !bytes.Equal(writeBuf, readBuf) {
+					assert.Equal(b, "msg", writeBuf, readBuf)
+				}
+
+				err = <-werrs
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+			b.StopTimer()
+
+			err := c1.Close(websocket.StatusNormalClosure, "")
+			assert.Success(b, err)
+		})
+	}
+}

frame.go

@@ -46,15 +46,14 @@ type header struct {
 
 // readFrameHeader reads a header from the reader.
 // See https://tools.ietf.org/html/rfc6455#section-5.2.
-func readFrameHeader(r *bufio.Reader) (_ header, err error) {
+func readFrameHeader(h *header, r *bufio.Reader) (err error) {
 	defer errd.Wrap(&err, "failed to read frame header")
 
 	b, err := r.ReadByte()
 	if err != nil {
-		return header{}, err
+		return err
 	}
 
-	var h header
 	h.fin = b&(1<<7) != 0
 	h.rsv1 = b&(1<<6) != 0
 	h.rsv2 = b&(1<<5) != 0
@@ -64,7 +63,7 @@ func readFrameHeader(r *bufio.Reader) (_ header, err error) {
 
 	b, err = r.ReadByte()
 	if err != nil {
-		return header{}, err
+		return err
 	}
 
 	h.masked = b&(1<<7) != 0
@@ -81,17 +80,17 @@ func readFrameHeader(r *bufio.Reader) (_ header, err error) {
 		err = binary.Read(r, binary.BigEndian, &h.payloadLength)
 	}
 	if err != nil {
-		return header{}, err
+		return err
 	}
 
 	if h.masked {
 		err = binary.Read(r, binary.LittleEndian, &h.maskKey)
 		if err != nil {
-			return header{}, err
+			return err
 		}
 	}
 
-	return h, nil
+	return nil
 }
 
 // maxControlPayload is the maximum length of a control frame payload.

frame_test.go

@@ -86,7 +86,8 @@ func testHeader(t *testing.T, h header) {
 	err = w.Flush()
 	assert.Success(t, err)
 
-	h2, err := readFrameHeader(r)
+	var h2 header
+	err = readFrameHeader(&h2, r)
 	assert.Success(t, err)
 
 	assert.Equal(t, "read header", h, h2)

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() {

read.go

@@ -173,7 +173,7 @@ func (c *Conn) readFrameHeader(ctx context.Context) (header, error) {
 	case c.readTimeout <- ctx:
 	}
 
-	h, err := readFrameHeader(c.br)
+	err := readFrameHeader(&c.readHeader, c.br)
 	if err != nil {
 		select {
 		case <-c.closed:
@@ -192,7 +192,7 @@ func (c *Conn) readFrameHeader(ctx context.Context) (header, error) {
 	case c.readTimeout <- context.Background():
 	}
 
-	return h, nil
+	return c.readHeader, nil
 }
 
 func (c *Conn) readFramePayload(ctx context.Context, p []byte) (int, error) {
@@ -390,6 +390,8 @@ func (mr *msgReader) read(p []byte) (int, error) {
 			return 0, err
 		}
 		mr.setFrame(h)
+
+		return mr.read(p)
 	}
 
 	if int64(len(p)) > mr.payloadLength {