1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
|
package websocketconn
import (
"bytes"
"fmt"
"io"
"io/ioutil"
"net"
"net/http"
"net/url"
"sync"
"testing"
"time"
"github.com/gorilla/websocket"
)
// Returns a (server, client) pair of websocketconn.Conns.
func connPair() (*Conn, *Conn, error) {
// Will be assigned inside server.Handler.
var serverConn *Conn
// Start up a web server to receive the request.
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return nil, nil, err
}
defer ln.Close()
errCh := make(chan error)
server := http.Server{
Handler: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
upgrader := websocket.Upgrader{
CheckOrigin: func(*http.Request) bool { return true },
}
ws, err := upgrader.Upgrade(rw, req, nil)
if err != nil {
errCh <- err
return
}
serverConn = New(ws)
close(errCh)
}),
}
defer server.Close()
go func() {
err := server.Serve(ln)
if err != nil && err != http.ErrServerClosed {
errCh <- err
}
}()
// Make a request to the web server.
urlStr := (&url.URL{Scheme: "ws", Host: ln.Addr().String()}).String()
ws, _, err := (&websocket.Dialer{}).Dial(urlStr, nil)
if err != nil {
return nil, nil, err
}
clientConn := New(ws)
// The server is finished when errCh is written to or closed.
err = <-errCh
if err != nil {
return nil, nil, err
}
return serverConn, clientConn, nil
}
// Test that you can write in chunks and read the result concatenated.
func TestWrite(t *testing.T) {
tests := [][][]byte{
{},
{[]byte("foo")},
{[]byte("foo"), []byte("bar")},
{{}, []byte("foo"), {}, {}, []byte("bar")},
}
for _, test := range tests {
s, c, err := connPair()
if err != nil {
t.Fatal(err)
}
// This is a little awkward because we need to read to and write
// from both ends of the Conn, and we need to do it in separate
// goroutines because otherwise a Write may block waiting for
// someone to Read it. Here we set up a loop in a separate
// goroutine, reading from the Conn s and writing to the dataCh
// and errCh channels, whose ultimate effect in the select loop
// below is like
// data, err := ioutil.ReadAll(s)
dataCh := make(chan []byte)
errCh := make(chan error)
go func() {
for {
var buf [1024]byte
n, err := s.Read(buf[:])
if err != nil {
errCh <- err
return
}
p := make([]byte, n)
copy(p, buf[:])
dataCh <- p
}
}()
// Write the data to the client side of the Conn, one chunk at a
// time.
for i, chunk := range test {
n, err := c.Write(chunk)
if err != nil || n != len(chunk) {
t.Fatalf("%+q Write chunk %d: got (%d, %v), expected (%d, %v)",
test, i, n, err, len(chunk), nil)
}
}
// We cannot immediately c.Close here, because that closes the
// connection right away, without waiting for buffered data to
// be sent.
// Pull data and err from the server goroutine above.
var data []byte
err = nil
loop:
for {
select {
case p := <-dataCh:
data = append(data, p...)
case err = <-errCh:
break loop
case <-time.After(100 * time.Millisecond):
break loop
}
}
s.Close()
c.Close()
// Now data and err contain the result of reading everything
// from s.
expected := bytes.Join(test, []byte{})
if err != nil || !bytes.Equal(data, expected) {
t.Fatalf("%+q ReadAll: got (%+q, %v), expected (%+q, %v)",
test, data, err, expected, nil)
}
}
}
// Test that multiple goroutines may call Read on a Conn simultaneously. Run
// this with
// go test -race
func TestConcurrentRead(t *testing.T) {
s, c, err := connPair()
if err != nil {
t.Fatal(err)
}
defer s.Close()
// Set up multiple threads reading from the same conn.
errCh := make(chan error, 2)
var wg sync.WaitGroup
wg.Add(2)
for i := 0; i < 2; i++ {
go func() {
defer wg.Done()
_, err := io.Copy(ioutil.Discard, s)
if err != nil {
errCh <- err
}
}()
}
// Write a bunch of data to the other end.
for i := 0; i < 2000; i++ {
_, err := c.Write([]byte(fmt.Sprintf("%d", i)))
if err != nil {
c.Close()
t.Fatalf("Write: %v", err)
}
}
c.Close()
wg.Wait()
close(errCh)
err = <-errCh
if err != nil {
t.Fatalf("Read: %v", err)
}
}
// Test that multiple goroutines may call Write on a Conn simultaneously. Run
// this with
// go test -race
func TestConcurrentWrite(t *testing.T) {
s, c, err := connPair()
if err != nil {
t.Fatal(err)
}
// Set up multiple threads writing to the same conn.
errCh := make(chan error, 3)
var wg sync.WaitGroup
wg.Add(2)
for i := 0; i < 2; i++ {
go func() {
defer wg.Done()
for j := 0; j < 1000; j++ {
_, err := fmt.Fprintf(s, "%d", j)
if err != nil {
errCh <- err
break
}
}
}()
}
go func() {
wg.Wait()
err := s.Close()
if err != nil {
errCh <- err
}
close(errCh)
}()
// Read from the other end.
_, err = io.Copy(ioutil.Discard, c)
c.Close()
if err != nil {
t.Fatalf("Read: %v", err)
}
err = <-errCh
if err != nil {
t.Fatalf("Write: %v", err)
}
}
|
