55package protocol
66
77import (
8+ "context"
9+ "encoding/json"
810 "errors"
911 "net/http"
12+ "sync"
1013 "time"
1114
1215 "github.com/gorilla/websocket"
@@ -23,11 +26,14 @@ type ReconnectingWebsocket struct {
2326 maxReconnectionDelay time.Duration
2427 reconnectionDelayGrowFactor float64
2528
29+ once sync.Once
2630 closedCh chan struct {}
27- connCh chan chan * websocket. Conn
31+ connCh chan chan * WebsocketConnection
2832 errCh chan error
2933
3034 log * logrus.Entry
35+
36+ ReconnectionHandler func ()
3137}
3238
3339// NewReconnectingWebsocket creates a new instance of ReconnectingWebsocket
@@ -39,7 +45,7 @@ func NewReconnectingWebsocket(url string, reqHeader http.Header, log *logrus.Ent
3945 maxReconnectionDelay : 30 * time .Second ,
4046 reconnectionDelayGrowFactor : 1.5 ,
4147 handshakeTimeout : 2 * time .Second ,
42- connCh : make (chan chan * websocket. Conn ),
48+ connCh : make (chan chan * WebsocketConnection ),
4349 closedCh : make (chan struct {}),
4450 errCh : make (chan error ),
4551 log : log ,
@@ -48,26 +54,27 @@ func NewReconnectingWebsocket(url string, reqHeader http.Header, log *logrus.Ent
4854
4955// Close closes the underlying webscoket connection.
5056func (rc * ReconnectingWebsocket ) Close () error {
51- close (rc .closedCh )
57+ rc .once .Do (func () {
58+ close (rc .closedCh )
59+ })
5260 return nil
5361}
5462
55- // WriteObject writes the JSON encoding of v as a message.
56- // See the documentation for encoding/json Marshal for details about the conversion of Go values to JSON.
57- func (rc * ReconnectingWebsocket ) WriteObject (v interface {}) error {
63+ // EnsureConnection ensures ws connections
64+ // Returns only if connection is permanently failed
65+ // If the passed handler returns false as closed then err is returned to the client,
66+ // otherwise err is treated as a connection error, and new conneciton is provided.
67+ func (rc * ReconnectingWebsocket ) EnsureConnection (handler func (conn * WebsocketConnection ) (closed bool , err error )) error {
5868 for {
59- connCh := make (chan * websocket. Conn , 1 )
69+ connCh := make (chan * WebsocketConnection , 1 )
6070 select {
6171 case <- rc .closedCh :
6272 return errors .New ("closed" )
6373 case rc .connCh <- connCh :
6474 }
6575 conn := <- connCh
66- err := conn .WriteJSON (v )
67- if err == nil {
68- return nil
69- }
70- if ! websocket .IsUnexpectedCloseError (err ) {
76+ closed , err := handler (conn )
77+ if ! closed {
7178 return err
7279 }
7380 select {
@@ -78,35 +85,62 @@ func (rc *ReconnectingWebsocket) WriteObject(v interface{}) error {
7885 }
7986}
8087
88+ func isJSONError (err error ) bool {
89+ _ , isJsonErr := err .(* json.InvalidUTF8Error )
90+ if isJsonErr {
91+ return true
92+ }
93+ _ , isJsonErr = err .(* json.InvalidUnmarshalError )
94+ if isJsonErr {
95+ return true
96+ }
97+ _ , isJsonErr = err .(* json.MarshalerError )
98+ if isJsonErr {
99+ return true
100+ }
101+ _ , isJsonErr = err .(* json.SyntaxError )
102+ if isJsonErr {
103+ return true
104+ }
105+ _ , isJsonErr = err .(* json.UnmarshalFieldError )
106+ if isJsonErr {
107+ return true
108+ }
109+ _ , isJsonErr = err .(* json.UnmarshalTypeError )
110+ if isJsonErr {
111+ return true
112+ }
113+ _ , isJsonErr = err .(* json.UnsupportedTypeError )
114+ if isJsonErr {
115+ return true
116+ }
117+ _ , isJsonErr = err .(* json.UnsupportedValueError )
118+ return isJsonErr
119+ }
120+
121+ // WriteObject writes the JSON encoding of v as a message.
122+ // See the documentation for encoding/json Marshal for details about the conversion of Go values to JSON.
123+ func (rc * ReconnectingWebsocket ) WriteObject (v interface {}) error {
124+ return rc .EnsureConnection (func (conn * WebsocketConnection ) (bool , error ) {
125+ err := conn .WriteJSON (v )
126+ closed := err != nil && ! isJSONError (err )
127+ return closed , err
128+ })
129+ }
130+
81131// ReadObject reads the next JSON-encoded message from the connection and stores it in the value pointed to by v.
82132// See the documentation for the encoding/json Unmarshal function for details about the conversion of JSON to a Go value.
83133func (rc * ReconnectingWebsocket ) ReadObject (v interface {}) error {
84- for {
85- connCh := make (chan * websocket.Conn , 1 )
86- select {
87- case <- rc .closedCh :
88- return errors .New ("closed" )
89- case rc .connCh <- connCh :
90- }
91- conn := <- connCh
134+ return rc .EnsureConnection (func (conn * WebsocketConnection ) (bool , error ) {
92135 err := conn .ReadJSON (v )
93- if err == nil {
94- return nil
95- }
96- if ! websocket .IsUnexpectedCloseError (err ) {
97- return err
98- }
99- select {
100- case <- rc .closedCh :
101- return errors .New ("closed" )
102- case rc .errCh <- err :
103- }
104- }
136+ closed := err != nil && ! isJSONError (err )
137+ return closed , err
138+ })
105139}
106140
107141// Dial creates a new client connection.
108142func (rc * ReconnectingWebsocket ) Dial () {
109- var conn * websocket. Conn
143+ var conn * WebsocketConnection
110144 defer func () {
111145 if conn == nil {
112146 return
@@ -129,19 +163,26 @@ func (rc *ReconnectingWebsocket) Dial() {
129163
130164 time .Sleep (1 * time .Second )
131165 conn = rc .connect ()
166+ if conn != nil && rc .ReconnectionHandler != nil {
167+ go rc .ReconnectionHandler ()
168+ }
132169 }
133170 }
134171}
135172
136- func (rc * ReconnectingWebsocket ) connect () * websocket. Conn {
173+ func (rc * ReconnectingWebsocket ) connect () * WebsocketConnection {
137174 delay := rc .minReconnectionDelay
138175 for {
139176 dialer := websocket.Dialer {HandshakeTimeout : rc .handshakeTimeout }
140177 conn , _ , err := dialer .Dial (rc .url , rc .reqHeader )
141178 if err == nil {
142179 rc .log .WithField ("url" , rc .url ).Info ("connection was successfully established" )
143-
144- return conn
180+ ws , err := NewWebsocketConnection (context .Background (), conn , func (staleErr error ) {
181+ rc .errCh <- staleErr
182+ })
183+ if err == nil {
184+ return ws
185+ }
145186 }
146187
147188 rc .log .WithError (err ).WithField ("url" , rc .url ).Errorf ("failed to connect, trying again in %d seconds..." , uint32 (delay .Seconds ()))
0 commit comments