client.go 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. package client
  2. import (
  3. "errors"
  4. "fmt"
  5. "git.bvbej.com/bvbej/base-golang/pkg/websocket/client/service"
  6. "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec"
  7. _ "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec/json"
  8. _ "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec/protobuf"
  9. "git.bvbej.com/bvbej/base-golang/tool/assist"
  10. "github.com/gorilla/websocket"
  11. "go.uber.org/zap"
  12. "net/http"
  13. "net/url"
  14. "reflect"
  15. "sync/atomic"
  16. "time"
  17. )
  18. const (
  19. writeWait = 20 * time.Second
  20. pongWait = 60 * time.Second
  21. reconnectWait = 3 * time.Second
  22. pingPeriod = (pongWait * 9) / 10
  23. maxFrameMessageLen = 16 * 1024
  24. maxSendBuffer = 32
  25. )
  26. var (
  27. _ Client = (*client)(nil)
  28. ErrBrokenPipe = errors.New("send to broken pipe")
  29. ErrBufferPoolExceed = errors.New("send buffer exceed")
  30. )
  31. type Client interface {
  32. readLoop()
  33. writeLoop()
  34. ping()
  35. onReceive(msg []byte) error
  36. onSend(msg []byte) error
  37. connect() error
  38. Send(router string, data any) error
  39. Connect(requestHeader http.Header) error
  40. OnReceiveError() <-chan error
  41. OnReconnect(f func())
  42. }
  43. type client struct {
  44. url url.URL
  45. requestHeader http.Header
  46. logger *zap.Logger
  47. session *service.Session
  48. isConnected *atomic.Bool
  49. routerCodec codec.Codec
  50. send chan []byte
  51. handlers map[string]*service.Handler // 注册的方法列表
  52. onReceiveChan chan error
  53. pingTicker *time.Ticker
  54. onReconnect func()
  55. }
  56. func New(logger *zap.Logger, url url.URL, decoder string, handlers interface{}) (Client, error) {
  57. if !assist.InArray(url.Scheme, []string{"ws", "wss"}) {
  58. return nil, errors.New(`param: scheme not supported`)
  59. }
  60. routerCodec := codec.GetCodec(decoder)
  61. if routerCodec == nil {
  62. return nil, errors.New(`param: codec not supported`)
  63. }
  64. components := service.RegisterHandler(handlers)
  65. if len(components) == 0 {
  66. return nil, errors.New(`param: handlers unqualified`)
  67. }
  68. c := &client{
  69. logger: logger,
  70. isConnected: new(atomic.Bool),
  71. routerCodec: routerCodec,
  72. url: url,
  73. send: make(chan []byte, maxSendBuffer),
  74. handlers: components,
  75. onReceiveChan: make(chan error, maxSendBuffer),
  76. pingTicker: time.NewTicker(pingPeriod),
  77. }
  78. return c, nil
  79. }
  80. func (c *client) readLoop() {
  81. _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
  82. c.session.Conn.SetPongHandler(func(string) error {
  83. _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
  84. return nil
  85. })
  86. for {
  87. _, data, err := c.session.Conn.ReadMessage()
  88. if err != nil {
  89. c.isConnected.Store(false)
  90. for c.connect() != nil {
  91. time.Sleep(reconnectWait)
  92. }
  93. if c.onReconnect != nil {
  94. c.onReconnect()
  95. }
  96. continue
  97. }
  98. err = c.onReceive(data)
  99. if err != nil {
  100. c.onReceiveChan <- err
  101. }
  102. }
  103. }
  104. func (c *client) writeLoop() {
  105. for {
  106. msg := <-c.send
  107. _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
  108. if err := c.session.Conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
  109. continue
  110. }
  111. }
  112. }
  113. func (c *client) ping() {
  114. for range c.pingTicker.C {
  115. _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
  116. if err := c.session.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
  117. for c.connect() != nil {
  118. time.Sleep(reconnectWait)
  119. }
  120. }
  121. }
  122. }
  123. func (c *client) connect() error {
  124. conn, _, err := websocket.DefaultDialer.Dial(c.url.String(), c.requestHeader)
  125. if err != nil {
  126. return fmt.Errorf("dial: %s", err)
  127. }
  128. c.session = service.NewSession(conn)
  129. c.isConnected.Store(true)
  130. return nil
  131. }
  132. func (c *client) onReceive(msg []byte) error {
  133. _, msgPack, err := c.routerCodec.Unmarshal(msg)
  134. if err != nil {
  135. return fmt.Errorf("onreceive: %v", err)
  136. }
  137. router, ok := msgPack.Router.(string)
  138. if !ok {
  139. return fmt.Errorf("onreceive: invalid router:%v", msgPack.Router)
  140. }
  141. s, ok := c.handlers[router]
  142. if !ok {
  143. return fmt.Errorf("onreceive: function not registed router:%s err:%v", msgPack.Router, err)
  144. }
  145. if msgPack.Err != nil {
  146. return fmt.Errorf("%s:%s", router, msgPack.Err)
  147. }
  148. var args = []reflect.Value{s.Receiver, reflect.ValueOf(c.session), reflect.ValueOf(msgPack.DataPtr)}
  149. s.Method.Func.Call(args)
  150. return nil
  151. }
  152. func (c *client) onSend(msg []byte) (err error) {
  153. defer func() {
  154. if e := recover(); e != nil {
  155. err = ErrBrokenPipe
  156. }
  157. }()
  158. if !c.isConnected.Load() {
  159. return ErrBrokenPipe
  160. }
  161. if len(c.send) >= maxSendBuffer {
  162. return ErrBufferPoolExceed
  163. }
  164. if len(msg) > maxFrameMessageLen {
  165. return
  166. }
  167. c.send <- msg
  168. return nil
  169. }
  170. func (c *client) Connect(requestHeader http.Header) error {
  171. c.requestHeader = requestHeader
  172. err := c.connect()
  173. if err != nil {
  174. return err
  175. }
  176. go c.ping()
  177. go c.writeLoop()
  178. go c.readLoop()
  179. return nil
  180. }
  181. func (c *client) Send(router string, data any) error {
  182. rb, err := c.routerCodec.Marshal(router, data, nil)
  183. if err != nil {
  184. return fmt.Errorf("service: %v", err)
  185. }
  186. return c.onSend(rb)
  187. }
  188. func (c *client) OnReceiveError() <-chan error {
  189. return c.onReceiveChan
  190. }
  191. func (c *client) OnReconnect(f func()) {
  192. c.onReconnect = f
  193. }