@@ -23,6 +23,9 @@ const (
2323 writeWait = 10 * time .Second
2424 pongWait = 60 * time .Second
2525 pingPeriod = (pongWait * 9 ) / 10
26+
27+ maximumMessageSize = 2048
28+ readMessageChanSize = 20 // <- I've put 20 here because it seems like a reasonable buffer but it may to increase
2629)
2730
2831type readMessage struct {
@@ -89,38 +92,8 @@ func Websocket(ctx *context.Context) {
8992 }
9093 defer unregister ()
9194
92- readChan := make (chan readMessage , 20 )
93- go func () {
94- defer conn .Close ()
95- conn .SetReadLimit (2048 )
96- if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
97- log .Error ("unable to SetReadDeadline: %v" , err )
98- return
99- }
100- conn .SetPongHandler (func (string ) error {
101- if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
102- log .Error ("unable to SetReadDeadline: %v" , err )
103- }
104- return nil
105- })
106-
107- for {
108- messageType , message , err := conn .ReadMessage ()
109- readChan <- readMessage {
110- messageType : messageType ,
111- message : message ,
112- err : err ,
113- }
114- if err != nil {
115- close (readChan )
116- return
117- }
118- if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
119- log .Error ("unable to SetReadDeadline: %v" , err )
120- return
121- }
122- }
123- }()
95+ readMessageChan := make (chan readMessage , readMessageChanSize )
96+ go readMessagesFromConnToChan (conn , readMessageChan )
12497
12598 pingTicker := time .NewTicker (pingPeriod )
12699
@@ -139,6 +112,7 @@ func Websocket(ctx *context.Context) {
139112 return
140113 default :
141114 }
115+
142116 if err := conn .SetWriteDeadline (time .Now ().Add (writeWait )); err != nil {
143117 log .Error ("unable to SetWriteDeadline: %v" , err )
144118 return
@@ -147,10 +121,11 @@ func Websocket(ctx *context.Context) {
147121 log .Error ("unable to send PingMessage: %v" , err )
148122 return
149123 }
150- case message , ok := <- readChan :
124+ case message , ok := <- readMessageChan :
151125 if ! ok {
152126 break
153127 }
128+
154129 // ensure that we're not already cancelled
155130 select {
156131 case <- notify :
@@ -159,11 +134,14 @@ func Websocket(ctx *context.Context) {
159134 return
160135 default :
161136 }
137+
138+ // FIXME: HANDLE MESSAGES
162139 log .Info ("Got Message: %d:%s:%v" , message .messageType , message .message , message .err )
163140 case event , ok := <- eventChan :
164141 if ! ok {
165142 break
166143 }
144+
167145 // ensure that we're not already cancelled
168146 select {
169147 case <- notify :
@@ -172,6 +150,8 @@ func Websocket(ctx *context.Context) {
172150 return
173151 default :
174152 }
153+
154+ // Handle events
175155 if event .Name == "logout" {
176156 if ctx .Session .ID () == event .Data {
177157 event = & eventsource.Event {
@@ -196,14 +176,50 @@ func Websocket(ctx *context.Context) {
196176 }
197177}
198178
179+ func readMessagesFromConnToChan (conn * websocket.Conn , messageChan chan readMessage ) {
180+ defer func () {
181+ close (messageChan ) // Please note: this has to be within a wrapping anonymous func otherwise it will be evaluated when creating the defer
182+ _ = conn .Close ()
183+ }()
184+ conn .SetReadLimit (maximumMessageSize )
185+ if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
186+ log .Error ("unable to SetReadDeadline: %v" , err )
187+ return
188+ }
189+ conn .SetPongHandler (func (string ) error {
190+ if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
191+ log .Error ("unable to SetReadDeadline: %v" , err )
192+ }
193+ return nil
194+ })
195+
196+ for {
197+ messageType , message , err := conn .ReadMessage ()
198+ messageChan <- readMessage {
199+ messageType : messageType ,
200+ message : message ,
201+ err : err ,
202+ }
203+ if err != nil {
204+ // don't need to handle the error here as it is passed down the channel
205+ return
206+ }
207+ if err := conn .SetReadDeadline (time .Now ().Add (pongWait )); err != nil {
208+ log .Error ("unable to SetReadDeadline: %v" , err )
209+ return
210+ }
211+ }
212+ }
213+
199214func writeEvent (conn * websocket.Conn , event * eventsource.Event ) error {
200215 if err := conn .SetWriteDeadline (time .Now ().Add (writeWait )); err != nil {
201216 log .Error ("unable to SetWriteDeadline: %v" , err )
202217 return err
203218 }
219+
204220 w , err := conn .NextWriter (websocket .TextMessage )
205221 if err != nil {
206- log .Warn ("Unable to get writer for websocket %v" , err )
222+ log .Error ("Unable to get writer for websocket %v" , err )
207223 return err
208224 }
209225
0 commit comments