client.go 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226
  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. }
  42. type client struct {
  43. url url.URL
  44. requestHeader http.Header
  45. logger *zap.Logger
  46. session *service.Session
  47. isConnected *atomic.Bool
  48. routerCodec codec.Codec
  49. send chan []byte
  50. handlers map[string]*service.Handler // 注册的方法列表
  51. onReceiveChan chan error
  52. pingTicker *time.Ticker
  53. }
  54. func New(logger *zap.Logger, url url.URL, decoder string, handlers interface{}) (Client, error) {
  55. if !assist.InArray(url.Scheme, []string{"ws", "wss"}) {
  56. return nil, errors.New(`param: scheme not supported`)
  57. }
  58. routerCodec := codec.GetCodec(decoder)
  59. if routerCodec == nil {
  60. return nil, errors.New(`param: codec not supported`)
  61. }
  62. components := service.RegisterHandler(handlers)
  63. if len(components) == 0 {
  64. return nil, errors.New(`param: handlers unqualified`)
  65. }
  66. c := &client{
  67. logger: logger,
  68. isConnected: new(atomic.Bool),
  69. routerCodec: routerCodec,
  70. url: url,
  71. send: make(chan []byte, maxSendBuffer),
  72. handlers: components,
  73. onReceiveChan: make(chan error, maxSendBuffer),
  74. pingTicker: time.NewTicker(pingPeriod),
  75. }
  76. return c, nil
  77. }
  78. func (c *client) readLoop() {
  79. defer func() {
  80. c.logger.Sugar().Info("readLoop closed")
  81. }()
  82. _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
  83. c.session.Conn.SetPongHandler(func(string) error {
  84. _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
  85. return nil
  86. })
  87. c.logger.Sugar().Info("readLoop running")
  88. for {
  89. _, data, err := c.session.Conn.ReadMessage()
  90. if err != nil {
  91. c.isConnected.Store(false)
  92. for c.connect() != nil {
  93. time.Sleep(reconnectWait)
  94. }
  95. continue
  96. }
  97. err = c.onReceive(data)
  98. if err != nil {
  99. c.onReceiveChan <- err
  100. }
  101. }
  102. }
  103. func (c *client) writeLoop() {
  104. defer func() {
  105. c.logger.Sugar().Info("writeLoop closed")
  106. }()
  107. c.logger.Sugar().Info("writeLoop running")
  108. for {
  109. select {
  110. case msg := <-c.send:
  111. _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
  112. if err := c.session.Conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
  113. break
  114. }
  115. }
  116. }
  117. }
  118. func (c *client) ping() {
  119. for range c.pingTicker.C {
  120. _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
  121. if err := c.session.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
  122. for c.connect() != nil {
  123. time.Sleep(reconnectWait)
  124. }
  125. }
  126. }
  127. }
  128. func (c *client) connect() error {
  129. conn, _, err := websocket.DefaultDialer.Dial(c.url.String(), c.requestHeader)
  130. if err != nil {
  131. return fmt.Errorf("dial: %s", err)
  132. }
  133. c.session = service.NewSession(conn)
  134. c.isConnected.Store(true)
  135. return nil
  136. }
  137. func (c *client) onReceive(msg []byte) error {
  138. _, msgPack, err := c.routerCodec.Unmarshal(msg)
  139. if err != nil {
  140. return fmt.Errorf("onreceive: %v", err)
  141. }
  142. router, ok := msgPack.Router.(string)
  143. if !ok {
  144. return fmt.Errorf("onreceive: invalid router:%v", msgPack.Router)
  145. }
  146. s, ok := c.handlers[router]
  147. if !ok {
  148. return fmt.Errorf("onreceive: function not registed router:%s err:%v", msgPack.Router, err)
  149. }
  150. var args = []reflect.Value{s.Receiver, reflect.ValueOf(c.session), reflect.ValueOf(msgPack.DataPtr)}
  151. s.Method.Func.Call(args)
  152. return nil
  153. }
  154. func (c *client) onSend(msg []byte) (err error) {
  155. defer func() {
  156. if e := recover(); e != nil {
  157. err = ErrBrokenPipe
  158. }
  159. }()
  160. if !c.isConnected.Load() {
  161. return ErrBrokenPipe
  162. }
  163. if len(c.send) >= maxSendBuffer {
  164. return ErrBufferPoolExceed
  165. }
  166. if len(msg) > maxFrameMessageLen {
  167. return
  168. }
  169. c.send <- msg
  170. return nil
  171. }
  172. func (c *client) Send(router string, data any) error {
  173. rb, err := c.routerCodec.Marshal(router, data, nil)
  174. if err != nil {
  175. return fmt.Errorf("service: %v", err)
  176. }
  177. return c.onSend(rb)
  178. }
  179. func (c *client) OnReceiveError() <-chan error {
  180. return c.onReceiveChan
  181. }
  182. func (c *client) Connect(requestHeader http.Header) error {
  183. c.requestHeader = requestHeader
  184. err := c.connect()
  185. if err != nil {
  186. return err
  187. }
  188. go c.ping()
  189. go c.writeLoop()
  190. go c.readLoop()
  191. return nil
  192. }