@@ -3,158 +3,57 @@ package main
33import (
44 "context"
55 "errors"
6- "fmt"
7- "io"
86 "log"
97 "net"
108 "net/http"
9+ "os"
10+ "os/signal"
1111 "time"
12-
13- "golang.org/x/time/rate"
14-
15- "nhooyr.io/websocket"
16- "nhooyr.io/websocket/wsjson"
1712)
1813
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.
2214func 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 )
5116
52- // Now we dial the server, send the messages and echo the responses.
53- err = client ("ws://" + l .Addr ().String ())
17+ err := run ()
5418 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 )
9220 }
9321}
9422
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 initializes the chatServer and routes and then
24+ // starts a http.Server for the passed in address.
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" )
10528 }
10629
107- typ , r , err := c . Reader ( ctx )
30+ l , err := net . Listen ( "tcp" , os . Args [ 1 ] )
10831 if err != nil {
10932 return err
11033 }
34+ log .Printf ("listening on http://%v" , l .Addr ())
11135
112- w , err := c .Writer (ctx , typ )
113- if err != nil {
114- return err
36+ s := & http.Server {
37+ Handler : http .HandlerFunc (echoServer ),
38+ ReadTimeout : time .Second * 10 ,
39+ WriteTimeout : time .Second * 10 ,
11540 }
41+ errc := make (chan error , 1 )
42+ go func () {
43+ errc <- s .Serve (l )
44+ }()
11645
117- _ , err = io .Copy (w , r )
118- if err != nil {
119- return fmt .Errorf ("failed to io.Copy: %w" , err )
46+ sigs := make (chan os.Signal , 1 )
47+ signal .Notify (sigs , os .Interrupt )
48+ select {
49+ case err := <- errc :
50+ log .Printf ("failed to serve: %v" , err )
51+ case sig := <- sigs :
52+ log .Printf ("terminating: %v" , sig )
12053 }
12154
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 )
55+ ctx , cancel := context .WithTimeout (context .Background (), time .Second * 10 )
13156 defer cancel ()
13257
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
58+ return s .Shutdown (ctx )
16059}
0 commit comments