You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

224 lines
4.5 KiB

7 years ago
  1. package rpcclient
  2. import (
  3. "context"
  4. "encoding/json"
  5. "net"
  6. "net/http"
  7. "net/http/httptest"
  8. "sync"
  9. "testing"
  10. "time"
  11. "github.com/gorilla/websocket"
  12. "github.com/stretchr/testify/require"
  13. "github.com/tendermint/tmlibs/log"
  14. types "github.com/tendermint/tendermint/rpc/lib/types"
  15. )
  16. var wsCallTimeout = 5 * time.Second
  17. type myHandler struct {
  18. closeConnAfterRead bool
  19. mtx sync.RWMutex
  20. }
  21. var upgrader = websocket.Upgrader{
  22. ReadBufferSize: 1024,
  23. WriteBufferSize: 1024,
  24. }
  25. func (h *myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  26. conn, err := upgrader.Upgrade(w, r, nil)
  27. if err != nil {
  28. panic(err)
  29. }
  30. defer conn.Close() // nolint: errcheck
  31. for {
  32. messageType, _, err := conn.ReadMessage()
  33. if err != nil {
  34. return
  35. }
  36. h.mtx.RLock()
  37. if h.closeConnAfterRead {
  38. if err := conn.Close(); err != nil {
  39. panic(err)
  40. }
  41. }
  42. h.mtx.RUnlock()
  43. res := json.RawMessage(`{}`)
  44. emptyRespBytes, _ := json.Marshal(types.RPCResponse{Result: res})
  45. if err := conn.WriteMessage(messageType, emptyRespBytes); err != nil {
  46. return
  47. }
  48. }
  49. }
  50. func TestWSClientReconnectsAfterReadFailure(t *testing.T) {
  51. var wg sync.WaitGroup
  52. // start server
  53. h := &myHandler{}
  54. s := httptest.NewServer(h)
  55. defer s.Close()
  56. c := startClient(t, s.Listener.Addr())
  57. defer c.Stop()
  58. wg.Add(1)
  59. go callWgDoneOnResult(t, c, &wg)
  60. h.mtx.Lock()
  61. h.closeConnAfterRead = true
  62. h.mtx.Unlock()
  63. // results in WS read error, no send retry because write succeeded
  64. call(t, "a", c)
  65. // expect to reconnect almost immediately
  66. time.Sleep(10 * time.Millisecond)
  67. h.mtx.Lock()
  68. h.closeConnAfterRead = false
  69. h.mtx.Unlock()
  70. // should succeed
  71. call(t, "b", c)
  72. wg.Wait()
  73. }
  74. func TestWSClientReconnectsAfterWriteFailure(t *testing.T) {
  75. var wg sync.WaitGroup
  76. // start server
  77. h := &myHandler{}
  78. s := httptest.NewServer(h)
  79. c := startClient(t, s.Listener.Addr())
  80. defer c.Stop()
  81. wg.Add(2)
  82. go callWgDoneOnResult(t, c, &wg)
  83. // hacky way to abort the connection before write
  84. if err := c.conn.Close(); err != nil {
  85. t.Error(err)
  86. }
  87. // results in WS write error, the client should resend on reconnect
  88. call(t, "a", c)
  89. // expect to reconnect almost immediately
  90. time.Sleep(10 * time.Millisecond)
  91. // should succeed
  92. call(t, "b", c)
  93. wg.Wait()
  94. }
  95. func TestWSClientReconnectFailure(t *testing.T) {
  96. // start server
  97. h := &myHandler{}
  98. s := httptest.NewServer(h)
  99. c := startClient(t, s.Listener.Addr())
  100. defer c.Stop()
  101. go func() {
  102. for {
  103. select {
  104. case <-c.ResponsesCh:
  105. case <-c.Quit():
  106. return
  107. }
  108. }
  109. }()
  110. // hacky way to abort the connection before write
  111. if err := c.conn.Close(); err != nil {
  112. t.Error(err)
  113. }
  114. s.Close()
  115. // results in WS write error
  116. // provide timeout to avoid blocking
  117. ctx, cancel := context.WithTimeout(context.Background(), wsCallTimeout)
  118. defer cancel()
  119. if err := c.Call(ctx, "a", make(map[string]interface{})); err != nil {
  120. t.Error(err)
  121. }
  122. // expect to reconnect almost immediately
  123. time.Sleep(10 * time.Millisecond)
  124. done := make(chan struct{})
  125. go func() {
  126. // client should block on this
  127. call(t, "b", c)
  128. close(done)
  129. }()
  130. // test that client blocks on the second send
  131. select {
  132. case <-done:
  133. t.Fatal("client should block on calling 'b' during reconnect")
  134. case <-time.After(5 * time.Second):
  135. t.Log("All good")
  136. }
  137. }
  138. func TestNotBlockingOnStop(t *testing.T) {
  139. timeout := 2 * time.Second
  140. s := httptest.NewServer(&myHandler{})
  141. c := startClient(t, s.Listener.Addr())
  142. c.Call(context.Background(), "a", make(map[string]interface{}))
  143. // Let the readRoutine get around to blocking
  144. time.Sleep(time.Second)
  145. passCh := make(chan struct{})
  146. go func() {
  147. // Unless we have a non-blocking write to ResponsesCh from readRoutine
  148. // this blocks forever ont the waitgroup
  149. c.Stop()
  150. passCh <- struct{}{}
  151. }()
  152. select {
  153. case <-passCh:
  154. // Pass
  155. case <-time.After(timeout):
  156. t.Fatalf("WSClient did failed to stop within %v seconds - is one of the read/write routines blocking?",
  157. timeout.Seconds())
  158. }
  159. }
  160. func startClient(t *testing.T, addr net.Addr) *WSClient {
  161. c := NewWSClient(addr.String(), "/websocket")
  162. err := c.Start()
  163. require.Nil(t, err)
  164. c.SetLogger(log.TestingLogger())
  165. return c
  166. }
  167. func call(t *testing.T, method string, c *WSClient) {
  168. err := c.Call(context.Background(), method, make(map[string]interface{}))
  169. require.NoError(t, err)
  170. }
  171. func callWgDoneOnResult(t *testing.T, c *WSClient, wg *sync.WaitGroup) {
  172. for {
  173. select {
  174. case resp := <-c.ResponsesCh:
  175. if resp.Error != nil {
  176. t.Fatalf("unexpected error: %v", resp.Error)
  177. }
  178. if resp.Result != nil {
  179. wg.Done()
  180. }
  181. case <-c.Quit():
  182. return
  183. }
  184. }
  185. }