Skip to content

Commit a681f25

Browse files
nhooyrnhooyr
committed
Move benchmarks into separate file
1 parent a819b57 commit a681f25

File tree

2 files changed

+155
-129
lines changed

2 files changed

+155
-129
lines changed

‎websocket_bench_test.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package websocket_test
2+
3+
import (
4+
"context"
5+
"io"
6+
"io/ioutil"
7+
"net/http"
8+
"nhooyr.io/websocket"
9+
"strconv"
10+
"strings"
11+
"testing"
12+
"time"
13+
)
14+
15+
16+
func BenchmarkConn(b *testing.B) {
17+
sizes := []int{
18+
2,
19+
16,
20+
32,
21+
512,
22+
4096,
23+
16384,
24+
}
25+
26+
b.Run("write", func(b *testing.B) {
27+
for _, size := range sizes {
28+
b.Run(strconv.Itoa(size), func(b *testing.B) {
29+
b.Run("stream", func(b *testing.B) {
30+
benchConn(b, false, true, size)
31+
})
32+
b.Run("buffer", func(b *testing.B) {
33+
benchConn(b, false, false, size)
34+
})
35+
})
36+
}
37+
})
38+
39+
b.Run("echo", func(b *testing.B) {
40+
for _, size := range sizes {
41+
b.Run(strconv.Itoa(size), func(b *testing.B) {
42+
benchConn(b, true, true, size)
43+
})
44+
}
45+
})
46+
}
47+
48+
func benchConn(b *testing.B, echo, stream bool, size int) {
49+
s, closeFn := testServer(b, func(w http.ResponseWriter, r *http.Request) error {
50+
c, err := websocket.Accept(w, r, nil)
51+
if err != nil {
52+
return err
53+
}
54+
if echo {
55+
echoLoop(r.Context(), c)
56+
} else {
57+
discardLoop(r.Context(), c)
58+
}
59+
return nil
60+
}, false)
61+
defer closeFn()
62+
63+
wsURL := strings.Replace(s.URL, "http", "ws", 1)
64+
65+
ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
66+
defer cancel()
67+
68+
c, _, err := websocket.Dial(ctx, wsURL, nil)
69+
if err != nil {
70+
b.Fatal(err)
71+
}
72+
defer c.Close(websocket.StatusInternalError, "")
73+
74+
msg := []byte(strings.Repeat("2", size))
75+
readBuf := make([]byte, len(msg))
76+
b.SetBytes(int64(len(msg)))
77+
b.ReportAllocs()
78+
b.ResetTimer()
79+
for i := 0; i < b.N; i++ {
80+
if stream {
81+
w, err := c.Writer(ctx, websocket.MessageText)
82+
if err != nil {
83+
b.Fatal(err)
84+
}
85+
86+
_, err = w.Write(msg)
87+
if err != nil {
88+
b.Fatal(err)
89+
}
90+
91+
err = w.Close()
92+
if err != nil {
93+
b.Fatal(err)
94+
}
95+
} else {
96+
err = c.Write(ctx, websocket.MessageText, msg)
97+
if err != nil {
98+
b.Fatal(err)
99+
}
100+
}
101+
102+
if echo {
103+
_, r, err := c.Reader(ctx)
104+
if err != nil {
105+
b.Fatal(err)
106+
}
107+
108+
_, err = io.ReadFull(r, readBuf)
109+
if err != nil {
110+
b.Fatal(err)
111+
}
112+
}
113+
}
114+
b.StopTimer()
115+
116+
c.Close(websocket.StatusNormalClosure, "")
117+
}
118+
119+
120+
func discardLoop(ctx context.Context, c *websocket.Conn) {
121+
defer c.Close(websocket.StatusInternalError, "")
122+
123+
ctx, cancel := context.WithTimeout(ctx, time.Minute)
124+
defer cancel()
125+
126+
b := make([]byte, 32768)
127+
echo := func() error {
128+
_, r, err := c.Reader(ctx)
129+
if err != nil {
130+
return err
131+
}
132+
133+
_, err = io.CopyBuffer(ioutil.Discard, r, b)
134+
if err != nil {
135+
return err
136+
}
137+
return nil
138+
}
139+
140+
for {
141+
err := echo()
142+
if err != nil {
143+
return
144+
}
145+
}
146+
}

‎websocket_test.go‎

Lines changed: 9 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -1092,6 +1092,8 @@ func TestAutobahn(t *testing.T) {
10921092

10931093
// Section 2.
10941094
t.Run("pingPong", func(t *testing.T) {
1095+
t.Parallel()
1096+
10951097
run(t, "emptyPayload", func(ctx context.Context, c *websocket.Conn) error {
10961098
ctx = c.CloseRead(ctx)
10971099
return c.PingWithPayload(ctx, "")
@@ -1197,40 +1199,17 @@ func TestAutobahn(t *testing.T) {
11971199
})
11981200
})
11991201

1200-
}
1201-
1202-
func BenchmarkConn(b *testing.B) {
1203-
sizes := []int{
1204-
2,
1205-
16,
1206-
32,
1207-
512,
1208-
4096,
1209-
16384,
1210-
}
1211-
1212-
b.Run("write", func(b *testing.B) {
1213-
for _, size := range sizes {
1214-
b.Run(strconv.Itoa(size), func(b *testing.B) {
1215-
b.Run("stream", func(b *testing.B) {
1216-
benchConn(b, false, true, size)
1217-
})
1218-
b.Run("buffer", func(b *testing.B) {
1219-
benchConn(b, false, false, size)
1220-
})
1221-
})
1222-
}
1223-
})
1202+
// Section 3.
1203+
t.Run("reserved", func(t *testing.T) {
1204+
t.Parallel()
12241205

1225-
b.Run("echo", func(b *testing.B) {
1226-
for _, size := range sizes {
1227-
b.Run(strconv.Itoa(size), func(b *testing.B) {
1228-
benchConn(b, true, true, size)
1229-
})
1230-
}
1206+
run(t, "rsv1", func(ctx context.Context, c *websocket.Conn) error {
1207+
c.WriteFrame()
1208+
})
12311209
})
12321210
}
12331211

1212+
12341213
func echoLoop(ctx context.Context, c *websocket.Conn) {
12351214
defer c.Close(websocket.StatusInternalError, "")
12361215

@@ -1272,105 +1251,6 @@ func echoLoop(ctx context.Context, c *websocket.Conn) {
12721251
}
12731252
}
12741253

1275-
func discardLoop(ctx context.Context, c *websocket.Conn) {
1276-
defer c.Close(websocket.StatusInternalError, "")
1277-
1278-
ctx, cancel := context.WithTimeout(ctx, time.Minute)
1279-
defer cancel()
1280-
1281-
b := make([]byte, 32768)
1282-
echo := func() error {
1283-
_, r, err := c.Reader(ctx)
1284-
if err != nil {
1285-
return err
1286-
}
1287-
1288-
_, err = io.CopyBuffer(ioutil.Discard, r, b)
1289-
if err != nil {
1290-
return err
1291-
}
1292-
return nil
1293-
}
1294-
1295-
for {
1296-
err := echo()
1297-
if err != nil {
1298-
return
1299-
}
1300-
}
1301-
}
1302-
1303-
func benchConn(b *testing.B, echo, stream bool, size int) {
1304-
s, closeFn := testServer(b, func(w http.ResponseWriter, r *http.Request) error {
1305-
c, err := websocket.Accept(w, r, nil)
1306-
if err != nil {
1307-
return err
1308-
}
1309-
if echo {
1310-
echoLoop(r.Context(), c)
1311-
} else {
1312-
discardLoop(r.Context(), c)
1313-
}
1314-
return nil
1315-
}, false)
1316-
defer closeFn()
1317-
1318-
wsURL := strings.Replace(s.URL, "http", "ws", 1)
1319-
1320-
ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
1321-
defer cancel()
1322-
1323-
c, _, err := websocket.Dial(ctx, wsURL, nil)
1324-
if err != nil {
1325-
b.Fatal(err)
1326-
}
1327-
defer c.Close(websocket.StatusInternalError, "")
1328-
1329-
msg := []byte(strings.Repeat("2", size))
1330-
readBuf := make([]byte, len(msg))
1331-
b.SetBytes(int64(len(msg)))
1332-
b.ReportAllocs()
1333-
b.ResetTimer()
1334-
for i := 0; i < b.N; i++ {
1335-
if stream {
1336-
w, err := c.Writer(ctx, websocket.MessageText)
1337-
if err != nil {
1338-
b.Fatal(err)
1339-
}
1340-
1341-
_, err = w.Write(msg)
1342-
if err != nil {
1343-
b.Fatal(err)
1344-
}
1345-
1346-
err = w.Close()
1347-
if err != nil {
1348-
b.Fatal(err)
1349-
}
1350-
} else {
1351-
err = c.Write(ctx, websocket.MessageText, msg)
1352-
if err != nil {
1353-
b.Fatal(err)
1354-
}
1355-
}
1356-
1357-
if echo {
1358-
_, r, err := c.Reader(ctx)
1359-
if err != nil {
1360-
b.Fatal(err)
1361-
}
1362-
1363-
_, err = io.ReadFull(r, readBuf)
1364-
if err != nil {
1365-
b.Fatal(err)
1366-
}
1367-
}
1368-
}
1369-
b.StopTimer()
1370-
1371-
c.Close(websocket.StatusNormalClosure, "")
1372-
}
1373-
13741254
func assertCloseStatus(err error, code websocket.StatusCode) error {
13751255
var cerr websocket.CloseError
13761256
if !xerrors.As(err, &cerr) {

0 commit comments

Comments
 (0)