@@ -3,158 +3,59 @@ package main
3
3
import (
4
4
"context"
5
5
"errors"
6
- "fmt"
7
- "io"
8
6
"log"
9
7
"net"
10
8
"net/http"
9
+ "os"
10
+ "os/signal"
11
11
"time"
12
-
13
- "golang.org/x/time/rate"
14
-
15
- "nhooyr.io/websocket"
16
- "nhooyr.io/websocket/wsjson"
17
12
)
18
13
19
- // This example starts a WebSocket echo server,
20
- // dials the server and then sends 5 different messages
21
- // and prints out the server's responses.
22
14
func main () {
23
- // First we listen on port 0 which means the OS will
24
- // assign us a random free port. This is the listener
25
- // the server will serve on and the client will connect to.
26
- l , err := net .Listen ("tcp" , "localhost:0" )
27
- if err != nil {
28
- log .Fatalf ("failed to listen: %v" , err )
29
- }
30
- defer l .Close ()
31
-
32
- s := & http.Server {
33
- Handler : http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
34
- err := echoServer (w , r )
35
- if err != nil {
36
- log .Printf ("echo server: %v" , err )
37
- }
38
- }),
39
- ReadTimeout : time .Second * 15 ,
40
- WriteTimeout : time .Second * 15 ,
41
- }
42
- defer s .Close ()
43
-
44
- // This starts the echo server on the listener.
45
- go func () {
46
- err := s .Serve (l )
47
- if err != http .ErrServerClosed {
48
- log .Fatalf ("failed to listen and serve: %v" , err )
49
- }
50
- }()
15
+ log .SetFlags (0 )
51
16
52
- // Now we dial the server, send the messages and echo the responses.
53
- err = client ("ws://" + l .Addr ().String ())
17
+ err := run ()
54
18
if err != nil {
55
- log .Fatalf ("client failed: %v" , err )
56
- }
57
-
58
- // Output:
59
- // received: map[i:0]
60
- // received: map[i:1]
61
- // received: map[i:2]
62
- // received: map[i:3]
63
- // received: map[i:4]
64
- }
65
-
66
- // echoServer is the WebSocket echo server implementation.
67
- // It ensures the client speaks the echo subprotocol and
68
- // only allows one message every 100ms with a 10 message burst.
69
- func echoServer (w http.ResponseWriter , r * http.Request ) error {
70
- c , err := websocket .Accept (w , r , & websocket.AcceptOptions {
71
- Subprotocols : []string {"echo" },
72
- })
73
- if err != nil {
74
- return err
75
- }
76
- defer c .Close (websocket .StatusInternalError , "the sky is falling" )
77
-
78
- if c .Subprotocol () != "echo" {
79
- c .Close (websocket .StatusPolicyViolation , "client must speak the echo subprotocol" )
80
- return errors .New ("client does not speak echo sub protocol" )
81
- }
82
-
83
- l := rate .NewLimiter (rate .Every (time .Millisecond * 100 ), 10 )
84
- for {
85
- err = echo (r .Context (), c , l )
86
- if websocket .CloseStatus (err ) == websocket .StatusNormalClosure {
87
- return nil
88
- }
89
- if err != nil {
90
- return fmt .Errorf ("failed to echo with %v: %w" , r .RemoteAddr , err )
91
- }
19
+ log .Fatal (err )
92
20
}
93
21
}
94
22
95
- // echo reads from the WebSocket connection and then writes
96
- // the received message back to it.
97
- // The entire function has 10s to complete.
98
- func echo (ctx context.Context , c * websocket.Conn , l * rate.Limiter ) error {
99
- ctx , cancel := context .WithTimeout (ctx , time .Second * 10 )
100
- defer cancel ()
101
-
102
- err := l .Wait (ctx )
103
- if err != nil {
104
- return err
23
+ // run starts a http.Server for the passed in address
24
+ // with all requests handled by echoServer.
25
+ func run () error {
26
+ if len (os .Args ) < 2 {
27
+ return errors .New ("please provide an address to listen on as the first argument" )
105
28
}
106
29
107
- typ , r , err := c . Reader ( ctx )
30
+ l , err := net . Listen ( "tcp" , os . Args [ 1 ] )
108
31
if err != nil {
109
32
return err
110
33
}
34
+ log .Printf ("listening on http://%v" , l .Addr ())
111
35
112
- w , err := c .Writer (ctx , typ )
113
- if err != nil {
114
- return err
36
+ s := & http.Server {
37
+ Handler : echoServer {
38
+ logf : log .Printf ,
39
+ },
40
+ ReadTimeout : time .Second * 10 ,
41
+ WriteTimeout : time .Second * 10 ,
115
42
}
43
+ errc := make (chan error , 1 )
44
+ go func () {
45
+ errc <- s .Serve (l )
46
+ }()
116
47
117
- _ , err = io .Copy (w , r )
118
- if err != nil {
119
- return fmt .Errorf ("failed to io.Copy: %w" , err )
48
+ sigs := make (chan os.Signal , 1 )
49
+ signal .Notify (sigs , os .Interrupt )
50
+ select {
51
+ case err := <- errc :
52
+ log .Printf ("failed to serve: %v" , err )
53
+ case sig := <- sigs :
54
+ log .Printf ("terminating: %v" , sig )
120
55
}
121
56
122
- err = w .Close ()
123
- return err
124
- }
125
-
126
- // client dials the WebSocket echo server at the given url.
127
- // It then sends it 5 different messages and echo's the server's
128
- // response to each.
129
- func client (url string ) error {
130
- ctx , cancel := context .WithTimeout (context .Background (), time .Minute )
57
+ ctx , cancel := context .WithTimeout (context .Background (), time .Second * 10 )
131
58
defer cancel ()
132
59
133
- c , _ , err := websocket .Dial (ctx , url , & websocket.DialOptions {
134
- Subprotocols : []string {"echo" },
135
- })
136
- if err != nil {
137
- return err
138
- }
139
- defer c .Close (websocket .StatusInternalError , "the sky is falling" )
140
-
141
- for i := 0 ; i < 5 ; i ++ {
142
- err = wsjson .Write (ctx , c , map [string ]int {
143
- "i" : i ,
144
- })
145
- if err != nil {
146
- return err
147
- }
148
-
149
- v := map [string ]int {}
150
- err = wsjson .Read (ctx , c , & v )
151
- if err != nil {
152
- return err
153
- }
154
-
155
- fmt .Printf ("received: %v\n " , v )
156
- }
157
-
158
- c .Close (websocket .StatusNormalClosure , "" )
159
- return nil
60
+ return s .Shutdown (ctx )
160
61
}
0 commit comments