Browse Source
Merge pull request #604 from tendermint/bugfix/ws-io-timeout
Biff up RPC WSClientpull/623/head
Ethan Buchman
7 years ago
committed by
GitHub
GitHub
10 changed files with 816 additions and 237 deletions
Split View
Diff Options
-
+10 -22benchmarks/simu/counter.go
-
+4 -3glide.lock
-
+11 -9glide.yaml
-
+2 -2node/node.go
-
+4 -3rpc/client/httpclient.go
-
+386 -90rpc/lib/client/ws_client.go
-
+193 -0rpc/lib/client/ws_client_test.go
-
+53 -14rpc/lib/rpc_test.go
-
+140 -94rpc/lib/server/handlers.go
-
+13 -0rpc/lib/types/types.go
+ 10
- 22
benchmarks/simu/counter.go
View File
+ 4
- 3
glide.lock
View File
+ 11
- 9
glide.yaml
View File
+ 2
- 2
node/node.go
View File
+ 4
- 3
rpc/client/httpclient.go
View File
+ 386
- 90
rpc/lib/client/ws_client.go
View File
| @ -1,160 +1,456 @@ | |||
| package rpcclient | |||
| import ( | |||
| "context" | |||
| "encoding/json" | |||
| "fmt" | |||
| "math" | |||
| "net" | |||
| "net/http" | |||
| "sync" | |||
| "time" | |||
| "github.com/gorilla/websocket" | |||
| "github.com/pkg/errors" | |||
| metrics "github.com/rcrowley/go-metrics" | |||
| types "github.com/tendermint/tendermint/rpc/lib/types" | |||
| cmn "github.com/tendermint/tmlibs/common" | |||
| ) | |||
| const ( | |||
| wsResultsChannelCapacity = 10 | |||
| wsErrorsChannelCapacity = 1 | |||
| wsWriteTimeoutSeconds = 10 | |||
| defaultMaxReconnectAttempts = 25 | |||
| defaultWriteWait = 0 | |||
| defaultReadWait = 0 | |||
| defaultPingPeriod = 0 | |||
| ) | |||
| // WSClient is a WebSocket client. The methods of WSClient are safe for use by | |||
| // multiple goroutines. | |||
| type WSClient struct { | |||
| cmn.BaseService | |||
| conn *websocket.Conn | |||
| Address string // IP:PORT or /path/to/socket | |||
| Endpoint string // /websocket/url/endpoint | |||
| Dialer func(string, string) (net.Conn, error) | |||
| *websocket.Conn | |||
| ResultsCh chan json.RawMessage // closes upon WSClient.Stop() | |||
| ErrorsCh chan error // closes upon WSClient.Stop() | |||
| // Time between sending a ping and receiving a pong. See | |||
| // https://godoc.org/github.com/rcrowley/go-metrics#Timer. | |||
| PingPongLatencyTimer metrics.Timer | |||
| // user facing channels, closed only when the client is being stopped. | |||
| ResultsCh chan json.RawMessage | |||
| ErrorsCh chan error | |||
| // internal channels | |||
| send chan types.RPCRequest // user requests | |||
| backlog chan types.RPCRequest // stores a single user request received during a conn failure | |||
| reconnectAfter chan error // reconnect requests | |||
| readRoutineQuit chan struct{} // a way for readRoutine to close writeRoutine | |||
| wg sync.WaitGroup | |||
| mtx sync.RWMutex | |||
| sentLastPingAt time.Time | |||
| reconnecting bool | |||
| // Maximum reconnect attempts (0 or greater; default: 25). | |||
| maxReconnectAttempts int | |||
| // Time allowed to write a message to the server. 0 means block until operation succeeds. | |||
| writeWait time.Duration | |||
| // Time allowed to read the next message from the server. 0 means block until operation succeeds. | |||
| readWait time.Duration | |||
| // Send pings to server with this period. Must be less than readWait. If 0, no pings will be sent. | |||
| pingPeriod time.Duration | |||
| } | |||
| // create a new connection | |||
| func NewWSClient(remoteAddr, endpoint string) *WSClient { | |||
| // NewWSClient returns a new client. See the commentary on the func(*WSClient) | |||
| // functions for a detailed description of how to configure ping period and | |||
| // pong wait time. The endpoint argument must begin with a `/`. | |||
| func NewWSClient(remoteAddr, endpoint string, options ...func(*WSClient)) *WSClient { | |||
| addr, dialer := makeHTTPDialer(remoteAddr) | |||
| wsClient := &WSClient{ | |||
| Address: addr, | |||
| Dialer: dialer, | |||
| Endpoint: endpoint, | |||
| Conn: nil, | |||
| c := &WSClient{ | |||
| Address: addr, | |||
| Dialer: dialer, | |||
| Endpoint: endpoint, | |||
| PingPongLatencyTimer: metrics.NewTimer(), | |||
| maxReconnectAttempts: defaultMaxReconnectAttempts, | |||
| readWait: defaultReadWait, | |||
| writeWait: defaultWriteWait, | |||
| pingPeriod: defaultPingPeriod, | |||
| } | |||
| c.BaseService = *cmn.NewBaseService(nil, "WSClient", c) | |||
| for _, option := range options { | |||
| option(c) | |||
| } | |||
| wsClient.BaseService = *cmn.NewBaseService(nil, "WSClient", wsClient) | |||
| return wsClient | |||
| return c | |||
| } | |||
| func (wsc *WSClient) String() string { | |||
| return wsc.Address + ", " + wsc.Endpoint | |||
| // MaxReconnectAttempts sets the maximum number of reconnect attempts before returning an error. | |||
| // It should only be used in the constructor and is not Goroutine-safe. | |||
| func MaxReconnectAttempts(max int) func(*WSClient) { | |||
| return func(c *WSClient) { | |||
| c.maxReconnectAttempts = max | |||
| } | |||
| } | |||
| // ReadWait sets the amount of time to wait before a websocket read times out. | |||
| // It should only be used in the constructor and is not Goroutine-safe. | |||
| func ReadWait(readWait time.Duration) func(*WSClient) { | |||
| return func(c *WSClient) { | |||
| c.readWait = readWait | |||
| } | |||
| } | |||
| // WriteWait sets the amount of time to wait before a websocket write times out. | |||
| // It should only be used in the constructor and is not Goroutine-safe. | |||
| func WriteWait(writeWait time.Duration) func(*WSClient) { | |||
| return func(c *WSClient) { | |||
| c.writeWait = writeWait | |||
| } | |||
| } | |||
| // OnStart implements cmn.BaseService interface | |||
| func (wsc *WSClient) OnStart() error { | |||
| wsc.BaseService.OnStart() | |||
| err := wsc.dial() | |||
| // PingPeriod sets the duration for sending websocket pings. | |||
| // It should only be used in the constructor - not Goroutine-safe. | |||
| func PingPeriod(pingPeriod time.Duration) func(*WSClient) { | |||
| return func(c *WSClient) { | |||
| c.pingPeriod = pingPeriod | |||
| } | |||
| } | |||
| // String returns WS client full address. | |||
| func (c *WSClient) String() string { | |||
| return fmt.Sprintf("%s (%s)", c.Address, c.Endpoint) | |||
| } | |||
| // OnStart implements cmn.Service by dialing a server and creating read and | |||
| // write routines. | |||
| func (c *WSClient) OnStart() error { | |||
| err := c.dial() | |||
| if err != nil { | |||
| return err | |||
| } | |||
| wsc.ResultsCh = make(chan json.RawMessage, wsResultsChannelCapacity) | |||
| wsc.ErrorsCh = make(chan error, wsErrorsChannelCapacity) | |||
| go wsc.receiveEventsRoutine() | |||
| c.ResultsCh = make(chan json.RawMessage) | |||
| c.ErrorsCh = make(chan error) | |||
| c.send = make(chan types.RPCRequest) | |||
| // 1 additional error may come from the read/write | |||
| // goroutine depending on which failed first. | |||
| c.reconnectAfter = make(chan error, 1) | |||
| // capacity for 1 request. a user won't be able to send more because the send | |||
| // channel is unbuffered. | |||
| c.backlog = make(chan types.RPCRequest, 1) | |||
| c.startReadWriteRoutines() | |||
| go c.reconnectRoutine() | |||
| return nil | |||
| } | |||
| // OnReset implements cmn.BaseService interface | |||
| func (wsc *WSClient) OnReset() error { | |||
| return nil | |||
| // OnStop implements cmn.Service. | |||
| func (c *WSClient) OnStop() {} | |||
| // Stop overrides cmn.Service#Stop. There is no other way to wait until Quit | |||
| // channel is closed. | |||
| func (c *WSClient) Stop() bool { | |||
| success := c.BaseService.Stop() | |||
| // only close user-facing channels when we can't write to them | |||
| c.wg.Wait() | |||
| close(c.ResultsCh) | |||
| close(c.ErrorsCh) | |||
| return success | |||
| } | |||
| // IsReconnecting returns true if the client is reconnecting right now. | |||
| func (c *WSClient) IsReconnecting() bool { | |||
| c.mtx.RLock() | |||
| defer c.mtx.RUnlock() | |||
| return c.reconnecting | |||
| } | |||
| // IsActive returns true if the client is running and not reconnecting. | |||
| func (c *WSClient) IsActive() bool { | |||
| return c.IsRunning() && !c.IsReconnecting() | |||
| } | |||
| // Send the given RPC request to the server. Results will be available on | |||
| // ResultsCh, errors, if any, on ErrorsCh. Will block until send succeeds or | |||
| // ctx.Done is closed. | |||
| func (c *WSClient) Send(ctx context.Context, request types.RPCRequest) error { | |||
| select { | |||
| case c.send <- request: | |||
| c.Logger.Info("sent a request", "req", request) | |||
| return nil | |||
| case <-ctx.Done(): | |||
| return ctx.Err() | |||
| } | |||
| } | |||
| // Call the given method. See Send description. | |||
| func (c *WSClient) Call(ctx context.Context, method string, params map[string]interface{}) error { | |||
| request, err := types.MapToRequest("", method, params) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| return c.Send(ctx, request) | |||
| } | |||
| // CallWithArrayParams the given method with params in a form of array. See | |||
| // Send description. | |||
| func (c *WSClient) CallWithArrayParams(ctx context.Context, method string, params []interface{}) error { | |||
| request, err := types.ArrayToRequest("", method, params) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| return c.Send(ctx, request) | |||
| } | |||
| func (wsc *WSClient) dial() error { | |||
| /////////////////////////////////////////////////////////////////////////////// | |||
| // Private methods | |||
| // Dial | |||
| func (c *WSClient) dial() error { | |||
| dialer := &websocket.Dialer{ | |||
| NetDial: wsc.Dialer, | |||
| NetDial: c.Dialer, | |||
| Proxy: http.ProxyFromEnvironment, | |||
| } | |||
| rHeader := http.Header{} | |||
| con, _, err := dialer.Dial("ws://"+wsc.Address+wsc.Endpoint, rHeader) | |||
| conn, _, err := dialer.Dial("ws://"+c.Address+c.Endpoint, rHeader) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| // Set the ping/pong handlers | |||
| con.SetPingHandler(func(m string) error { | |||
| // NOTE: https://github.com/gorilla/websocket/issues/97 | |||
| go con.WriteControl(websocket.PongMessage, []byte(m), time.Now().Add(time.Second*wsWriteTimeoutSeconds)) | |||
| return nil | |||
| }) | |||
| con.SetPongHandler(func(m string) error { | |||
| // NOTE: https://github.com/gorilla/websocket/issues/97 | |||
| return nil | |||
| }) | |||
| wsc.Conn = con | |||
| c.conn = conn | |||
| return nil | |||
| } | |||
| // OnStop implements cmn.BaseService interface | |||
| func (wsc *WSClient) OnStop() { | |||
| wsc.BaseService.OnStop() | |||
| wsc.Conn.Close() | |||
| // ResultsCh/ErrorsCh is closed in receiveEventsRoutine. | |||
| } | |||
| // reconnect tries to redial up to maxReconnectAttempts with exponential | |||
| // backoff. | |||
| func (c *WSClient) reconnect() error { | |||
| attempt := 0 | |||
| c.mtx.Lock() | |||
| c.reconnecting = true | |||
| c.mtx.Unlock() | |||
| defer func() { | |||
| c.mtx.Lock() | |||
| c.reconnecting = false | |||
| c.mtx.Unlock() | |||
| }() | |||
| func (wsc *WSClient) receiveEventsRoutine() { | |||
| for { | |||
| _, data, err := wsc.ReadMessage() | |||
| c.Logger.Info("reconnecting", "attempt", attempt+1) | |||
| d := time.Duration(math.Exp2(float64(attempt))) | |||
| time.Sleep(d * time.Second) | |||
| err := c.dial() | |||
| if err != nil { | |||
| wsc.Logger.Info("WSClient failed to read message", "err", err, "data", string(data)) | |||
| wsc.Stop() | |||
| break | |||
| c.Logger.Error("failed to redial", "err", err) | |||
| } else { | |||
| var response types.RPCResponse | |||
| err := json.Unmarshal(data, &response) | |||
| c.Logger.Info("reconnected") | |||
| return nil | |||
| } | |||
| attempt++ | |||
| if attempt > c.maxReconnectAttempts { | |||
| return errors.Wrap(err, "reached maximum reconnect attempts") | |||
| } | |||
| } | |||
| } | |||
| func (c *WSClient) startReadWriteRoutines() { | |||
| c.wg.Add(2) | |||
| c.readRoutineQuit = make(chan struct{}) | |||
| go c.readRoutine() | |||
| go c.writeRoutine() | |||
| } | |||
| func (c *WSClient) processBacklog() error { | |||
| select { | |||
| case request := <-c.backlog: | |||
| if c.writeWait > 0 { | |||
| c.conn.SetWriteDeadline(time.Now().Add(c.writeWait)) | |||
| } | |||
| err := c.conn.WriteJSON(request) | |||
| if err != nil { | |||
| c.Logger.Error("failed to resend request", "err", err) | |||
| c.reconnectAfter <- err | |||
| // requeue request | |||
| c.backlog <- request | |||
| return err | |||
| } | |||
| c.Logger.Info("resend a request", "req", request) | |||
| default: | |||
| } | |||
| return nil | |||
| } | |||
| func (c *WSClient) reconnectRoutine() { | |||
| for { | |||
| select { | |||
| case originalError := <-c.reconnectAfter: | |||
| // wait until writeRoutine and readRoutine finish | |||
| c.wg.Wait() | |||
| err := c.reconnect() | |||
| if err != nil { | |||
| c.Logger.Error("failed to reconnect", "err", err, "original_err", originalError) | |||
| c.Stop() | |||
| return | |||
| } else { | |||
| // drain reconnectAfter | |||
| LOOP: | |||
| for { | |||
| select { | |||
| case <-c.reconnectAfter: | |||
| default: | |||
| break LOOP | |||
| } | |||
| } | |||
| err = c.processBacklog() | |||
| if err == nil { | |||
| c.startReadWriteRoutines() | |||
| } | |||
| } | |||
| case <-c.Quit: | |||
| return | |||
| } | |||
| } | |||
| } | |||
| // The client ensures that there is at most one writer to a connection by | |||
| // executing all writes from this goroutine. | |||
| func (c *WSClient) writeRoutine() { | |||
| var ticker *time.Ticker | |||
| if c.pingPeriod > 0 { | |||
| // ticker with a predefined period | |||
| ticker = time.NewTicker(c.pingPeriod) | |||
| } else { | |||
| // ticker that never fires | |||
| ticker = &time.Ticker{C: make(<-chan time.Time)} | |||
| } | |||
| defer func() { | |||
| ticker.Stop() | |||
| c.conn.Close() | |||
| c.wg.Done() | |||
| }() | |||
| for { | |||
| select { | |||
| case request := <-c.send: | |||
| if c.writeWait > 0 { | |||
| c.conn.SetWriteDeadline(time.Now().Add(c.writeWait)) | |||
| } | |||
| err := c.conn.WriteJSON(request) | |||
| if err != nil { | |||
| wsc.Logger.Info("WSClient failed to parse message", "err", err, "data", string(data)) | |||
| wsc.ErrorsCh <- err | |||
| continue | |||
| c.Logger.Error("failed to send request", "err", err) | |||
| c.reconnectAfter <- err | |||
| // add request to the backlog, so we don't lose it | |||
| c.backlog <- request | |||
| return | |||
| } | |||
| if response.Error != "" { | |||
| wsc.ErrorsCh <- errors.Errorf(response.Error) | |||
| continue | |||
| case <-ticker.C: | |||
| if c.writeWait > 0 { | |||
| c.conn.SetWriteDeadline(time.Now().Add(c.writeWait)) | |||
| } | |||
| wsc.ResultsCh <- *response.Result | |||
| err := c.conn.WriteMessage(websocket.PingMessage, []byte{}) | |||
| if err != nil { | |||
| c.Logger.Error("failed to write ping", "err", err) | |||
| c.reconnectAfter <- err | |||
| return | |||
| } | |||
| c.mtx.Lock() | |||
| c.sentLastPingAt = time.Now() | |||
| c.mtx.Unlock() | |||
| c.Logger.Debug("sent ping") | |||
| case <-c.readRoutineQuit: | |||
| return | |||
| case <-c.Quit: | |||
| c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) | |||
| return | |||
| } | |||
| } | |||
| // this must be modified in the same go-routine that reads from the | |||
| // connection to avoid race conditions | |||
| wsc.Conn = nil | |||
| } | |||
| // The client ensures that there is at most one reader to a connection by | |||
| // executing all reads from this goroutine. | |||
| func (c *WSClient) readRoutine() { | |||
| defer func() { | |||
| c.conn.Close() | |||
| c.wg.Done() | |||
| }() | |||
| c.conn.SetPongHandler(func(string) error { | |||
| // gather latency stats | |||
| c.mtx.RLock() | |||
| t := c.sentLastPingAt | |||
| c.mtx.RUnlock() | |||
| c.PingPongLatencyTimer.UpdateSince(t) | |||
| c.Logger.Debug("got pong") | |||
| return nil | |||
| }) | |||
| for { | |||
| // reset deadline for every message type (control or data) | |||
| if c.readWait > 0 { | |||
| c.conn.SetReadDeadline(time.Now().Add(c.readWait)) | |||
| } | |||
| _, data, err := c.conn.ReadMessage() | |||
| if err != nil { | |||
| if !websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) { | |||
| return | |||
| } | |||
| c.Logger.Error("failed to read response", "err", err) | |||
| close(c.readRoutineQuit) | |||
| c.reconnectAfter <- err | |||
| return | |||
| } | |||
| // Cleanup | |||
| close(wsc.ResultsCh) | |||
| close(wsc.ErrorsCh) | |||
| var response types.RPCResponse | |||
| err = json.Unmarshal(data, &response) | |||
| if err != nil { | |||
| c.Logger.Error("failed to parse response", "err", err, "data", string(data)) | |||
| c.ErrorsCh <- err | |||
| continue | |||
| } | |||
| if response.Error != "" { | |||
| c.ErrorsCh <- errors.Errorf(response.Error) | |||
| continue | |||
| } | |||
| c.Logger.Info("got response", "resp", response.Result) | |||
| c.ResultsCh <- *response.Result | |||
| } | |||
| } | |||
| /////////////////////////////////////////////////////////////////////////////// | |||
| // Predefined methods | |||
| // Subscribe to an event. Note the server must have a "subscribe" route | |||
| // defined. | |||
| func (wsc *WSClient) Subscribe(eventid string) error { | |||
| params := map[string]interface{}{"event": eventid} | |||
| request, err := types.MapToRequest("", "subscribe", params) | |||
| if err == nil { | |||
| err = wsc.WriteJSON(request) | |||
| } | |||
| return err | |||
| func (c *WSClient) Subscribe(ctx context.Context, eventType string) error { | |||
| params := map[string]interface{}{"event": eventType} | |||
| return c.Call(ctx, "subscribe", params) | |||
| } | |||
| // Unsubscribe from an event. Note the server must have a "unsubscribe" route | |||
| // defined. | |||
| func (wsc *WSClient) Unsubscribe(eventid string) error { | |||
| params := map[string]interface{}{"event": eventid} | |||
| request, err := types.MapToRequest("", "unsubscribe", params) | |||
| if err == nil { | |||
| err = wsc.WriteJSON(request) | |||
| } | |||
| return err | |||
| func (c *WSClient) Unsubscribe(ctx context.Context, eventType string) error { | |||
| params := map[string]interface{}{"event": eventType} | |||
| return c.Call(ctx, "unsubscribe", params) | |||
| } | |||
| // Call asynchronously calls a given method by sending an RPCRequest to the | |||
| // server. Results will be available on ResultsCh, errors, if any, on ErrorsCh. | |||
| func (wsc *WSClient) Call(method string, params map[string]interface{}) error { | |||
| request, err := types.MapToRequest("", method, params) | |||
| if err == nil { | |||
| err = wsc.WriteJSON(request) | |||
| } | |||
| return err | |||
| // UnsubscribeAll from all. Note the server must have a "unsubscribe_all" route | |||
| // defined. | |||
| func (c *WSClient) UnsubscribeAll(ctx context.Context) error { | |||
| params := map[string]interface{}{} | |||
| return c.Call(ctx, "unsubscribe_all", params) | |||
| } | |||
+ 193
- 0
rpc/lib/client/ws_client_test.go
View File
| @ -0,0 +1,193 @@ | |||
| package rpcclient | |||
| import ( | |||
| "context" | |||
| "encoding/json" | |||
| "net" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "sync" | |||
| "testing" | |||
| "time" | |||
| "github.com/gorilla/websocket" | |||
| "github.com/stretchr/testify/require" | |||
| "github.com/tendermint/tmlibs/log" | |||
| types "github.com/tendermint/tendermint/rpc/lib/types" | |||
| ) | |||
| type myHandler struct { | |||
| closeConnAfterRead bool | |||
| mtx sync.RWMutex | |||
| } | |||
| var upgrader = websocket.Upgrader{ | |||
| ReadBufferSize: 1024, | |||
| WriteBufferSize: 1024, | |||
| } | |||
| func (h *myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { | |||
| conn, err := upgrader.Upgrade(w, r, nil) | |||
| if err != nil { | |||
| panic(err) | |||
| } | |||
| defer conn.Close() | |||
| for { | |||
| messageType, _, err := conn.ReadMessage() | |||
| if err != nil { | |||
| return | |||
| } | |||
| h.mtx.RLock() | |||
| if h.closeConnAfterRead { | |||
| conn.Close() | |||
| } | |||
| h.mtx.RUnlock() | |||
| res := json.RawMessage(`{}`) | |||
| emptyRespBytes, _ := json.Marshal(types.RPCResponse{Result: &res}) | |||
| if err := conn.WriteMessage(messageType, emptyRespBytes); err != nil { | |||
| return | |||
| } | |||
| } | |||
| } | |||
| func TestWSClientReconnectsAfterReadFailure(t *testing.T) { | |||
| var wg sync.WaitGroup | |||
| // start server | |||
| h := &myHandler{} | |||
| s := httptest.NewServer(h) | |||
| defer s.Close() | |||
| c := startClient(t, s.Listener.Addr()) | |||
| defer c.Stop() | |||
| wg.Add(1) | |||
| go callWgDoneOnResult(t, c, &wg) | |||
| h.mtx.Lock() | |||
| h.closeConnAfterRead = true | |||
| h.mtx.Unlock() | |||
| // results in WS read error, no send retry because write succeeded | |||
| call(t, "a", c) | |||
| // expect to reconnect almost immediately | |||
| time.Sleep(10 * time.Millisecond) | |||
| h.mtx.Lock() | |||
| h.closeConnAfterRead = false | |||
| h.mtx.Unlock() | |||
| // should succeed | |||
| call(t, "b", c) | |||
| wg.Wait() | |||
| } | |||
| func TestWSClientReconnectsAfterWriteFailure(t *testing.T) { | |||
| var wg sync.WaitGroup | |||
| // start server | |||
| h := &myHandler{} | |||
| s := httptest.NewServer(h) | |||
| c := startClient(t, s.Listener.Addr()) | |||
| defer c.Stop() | |||
| wg.Add(2) | |||
| go callWgDoneOnResult(t, c, &wg) | |||
| // hacky way to abort the connection before write | |||
| c.conn.Close() | |||
| // results in WS write error, the client should resend on reconnect | |||
| call(t, "a", c) | |||
| // expect to reconnect almost immediately | |||
| time.Sleep(10 * time.Millisecond) | |||
| // should succeed | |||
| call(t, "b", c) | |||
| wg.Wait() | |||
| } | |||
| func TestWSClientReconnectFailure(t *testing.T) { | |||
| // start server | |||
| h := &myHandler{} | |||
| s := httptest.NewServer(h) | |||
| c := startClient(t, s.Listener.Addr()) | |||
| defer c.Stop() | |||
| go func() { | |||
| for { | |||
| select { | |||
| case <-c.ResultsCh: | |||
| case <-c.ErrorsCh: | |||
| case <-c.Quit: | |||
| return | |||
| } | |||
| } | |||
| }() | |||
| // hacky way to abort the connection before write | |||
| c.conn.Close() | |||
| s.Close() | |||
| // results in WS write error | |||
| // provide timeout to avoid blocking | |||
| ctx, cancel := context.WithTimeout(context.Background(), time.Second) | |||
| defer cancel() | |||
| c.Call(ctx, "a", make(map[string]interface{})) | |||
| // expect to reconnect almost immediately | |||
| time.Sleep(10 * time.Millisecond) | |||
| done := make(chan struct{}) | |||
| go func() { | |||
| // client should block on this | |||
| call(t, "b", c) | |||
| close(done) | |||
| }() | |||
| // test that client blocks on the second send | |||
| select { | |||
| case <-done: | |||
| t.Fatal("client should block on calling 'b' during reconnect") | |||
| case <-time.After(5 * time.Second): | |||
| t.Log("All good") | |||
| } | |||
| } | |||
| func startClient(t *testing.T, addr net.Addr) *WSClient { | |||
| c := NewWSClient(addr.String(), "/websocket") | |||
| _, err := c.Start() | |||
| require.Nil(t, err) | |||
| c.SetLogger(log.TestingLogger()) | |||
| return c | |||
| } | |||
| func call(t *testing.T, method string, c *WSClient) { | |||
| err := c.Call(context.Background(), method, make(map[string]interface{})) | |||
| require.NoError(t, err) | |||
| } | |||
| func callWgDoneOnResult(t *testing.T, c *WSClient, wg *sync.WaitGroup) { | |||
| for { | |||
| select { | |||
| case res := <-c.ResultsCh: | |||
| if res != nil { | |||
| wg.Done() | |||
| } | |||
| case err := <-c.ErrorsCh: | |||
| if err != nil { | |||
| t.Fatalf("unexpected error: %v", err) | |||
| } | |||
| case <-c.Quit: | |||
| return | |||
| } | |||
| } | |||
| } | |||