Add XOR benchmarks

Author: nhooyr

bench_test.go

@@ -228,7 +228,7 @@ func (c *Conn) handleControl(h header) {
 	}
 
 	if h.masked {
-		xor(h.maskKey, 0, b)
+		fastXOR(h.maskKey, 0, b)
 	}
 
 	switch h.opcode {
@@ -322,7 +322,7 @@ func (c *Conn) dataReadLoop(h header) (err error) {
 			left -= int64(len(b))
 
 			if h.masked {
-				maskPos = xor(h.maskKey, maskPos, b)
+				maskPos = fastXOR(h.maskKey, maskPos, b)
 			}
 
 			// Must set this before we signal the read is done.

websocket.go

@@ -701,3 +701,100 @@ func checkWSTestIndex(t *testing.T, path string) {
 		}
 	}
 }
+
+func benchConn(b *testing.B, echo, stream bool) {
+	name := "buffered"
+	if stream {
+		name = "stream"
+	}
+
+	b.Run(name, func(b *testing.B) {
+		s, closeFn := testServer(b, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			c, err := websocket.Accept(w, r, websocket.AcceptOptions{})
+			if err != nil {
+				b.Logf("server handshake failed: %+v", err)
+				return
+			}
+			if echo {
+				echoLoop(r.Context(), c)
+			} else {
+				discardLoop(r.Context(), c)
+			}
+		}))
+		defer closeFn()
+
+		wsURL := strings.Replace(s.URL, "http", "ws", 1)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
+		defer cancel()
+
+		c, _, err := websocket.Dial(ctx, wsURL, websocket.DialOptions{})
+		if err != nil {
+			b.Fatalf("failed to dial: %v", err)
+		}
+		defer c.Close(websocket.StatusInternalError, "")
+
+		sizes := []int{
+			2,
+			512,
+			4096,
+			16384,
+		}
+
+		for _, size := range sizes {
+			msg := []byte(strings.Repeat("2", size))
+			buf := make([]byte, len(msg))
+			b.Run(strconv.Itoa(size), func(b *testing.B) {
+				b.SetBytes(int64(len(msg)))
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					if stream {
+						w, err := c.Writer(ctx, websocket.MessageText)
+						if err != nil {
+							b.Fatal(err)
+						}
+
+						_, err = w.Write(msg)
+						if err != nil {
+							b.Fatal(err)
+						}
+
+						err = w.Close()
+						if err != nil {
+							b.Fatal(err)
+						}
+					} else {
+						err = c.Write(ctx, websocket.MessageText, msg)
+						if err != nil {
+							b.Fatal(err)
+						}
+					}
+
+					if echo {
+						_, r, err := c.Reader(ctx)
+						if err != nil {
+							b.Fatal(err)
+						}
+
+						_, err = io.ReadFull(r, buf)
+						if err != nil {
+							b.Fatal(err)
+						}
+					}
+				}
+			})
+		}
+
+		c.Close(websocket.StatusNormalClosure, "")
+	})
+}
+
+func BenchmarkConn(b *testing.B) {
+	b.Run("write", func(b *testing.B) {
+		benchConn(b, false, false)
+		benchConn(b, false, true)
+	})
+	b.Run("echo", func(b *testing.B) {
+		benchConn(b, true, true)
+	})
+}

websocket_test.go

@@ -12,7 +12,7 @@ import (
 // The returned value is the position of the next byte
 // to be used for masking in the key. This is so that
 // unmasking can be performed without the entire frame.
-func xor(key [4]byte, keyPos int, b []byte) int {
+func fastXOR(key [4]byte, keyPos int, b []byte) int {
 	// If the payload is greater than 16 bytes, then it's worth
 	// masking 8 bytes at a time.
 	// Optimization from https://github.com/golang/go/issues/31586#issuecomment-485530859

xor.go

@@ -1,6 +1,8 @@
 package websocket
 
 import (
+	"crypto/rand"
+	"strconv"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -12,7 +14,7 @@ func Test_xor(t *testing.T) {
 	key := [4]byte{0xa, 0xb, 0xc, 0xff}
 	p := []byte{0xa, 0xb, 0xc, 0xf2, 0xc}
 	pos := 0
-	pos = xor(key, pos, p)
+	pos = fastXOR(key, pos, p)
 
 	if exp := []byte{0, 0, 0, 0x0d, 0x6}; !cmp.Equal(exp, p) {
 		t.Fatalf("unexpected mask: %v", cmp.Diff(exp, p))
@@ -22,3 +24,58 @@ func Test_xor(t *testing.T) {
 		t.Fatalf("unexpected mask pos: %v", cmp.Diff(exp, pos))
 	}
 }
+
+func basixXOR(maskKey [4]byte, pos int, b []byte) int {
+	for i := range b {
+		b[i] ^= maskKey[pos&3]
+		pos++
+	}
+	return pos & 3
+}
+
+func BenchmarkXOR(b *testing.B) {
+	sizes := []int{
+		2,
+		32,
+		512,
+		4096,
+		16384,
+	}
+
+	fns := []struct {
+		name string
+		fn   func([4]byte, int, []byte) int
+	}{
+		{
+			"basic",
+			basixXOR,
+		},
+		{
+			"fast",
+			fastXOR,
+		},
+	}
+
+	var maskKey [4]byte
+	_, err := rand.Read(maskKey[:])
+	if err != nil {
+		b.Fatalf("failed to populate mask key: %v", err)
+	}
+
+	for _, size := range sizes {
+		data := make([]byte, size)
+
+		b.Run(strconv.Itoa(size), func(b *testing.B) {
+			for _, fn := range fns {
+				b.Run(fn.name, func(b *testing.B) {
+					b.ReportAllocs()
+					b.SetBytes(int64(size))
+
+					for i := 0; i < b.N; i++ {
+						fn.fn(maskKey, 0, data)
+					}
+				})
+			}
+		})
+	}
+}