123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159 |
- package client
- import (
- "errors"
- "fmt"
- "git.bvbej.com/bvbej/base-golang/pkg/websocket/client/service"
- "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec"
- "git.bvbej.com/bvbej/base-golang/tool/assist"
- "github.com/gorilla/websocket"
- "go.uber.org/zap"
- "net/http"
- "net/url"
- "sync/atomic"
- "time"
- )
- const (
- writeWait = 20 * time.Second
- pongWait = 60 * time.Second
- pingPeriod = (pongWait * 9) / 10
- maxFrameMessageLen = 16 * 1024
- maxSendBuffer = 32
- )
- var (
- _ Client = (*client)(nil)
- ErrBrokenPipe = errors.New("send to broken pipe")
- ErrBufferPoolExceed = errors.New("send buffer exceed")
- )
- type Client interface {
- Connect(requestHeader http.Header) error
- }
- type client struct {
- logger *zap.Logger
- session service.Session
- isSurviving *atomic.Bool
- routerCodec codec.Codec
- url url.URL
- send chan []byte
- handlers map[string]*service.Handler // 注册的方法列表
- }
- func New(logger *zap.Logger, url url.URL, decoder string, handlers ...interface{}) (Client, error) {
- if !assist.InArray(url.Scheme, []string{"ws", "wss"}) {
- return nil, errors.New(`param: scheme not supported`)
- }
- routerCodec := codec.GetCodec(decoder)
- if routerCodec == nil {
- return nil, errors.New(`param: codec not supported`)
- }
- components := service.RegisterHandler(handlers)
- if len(components) == 0 {
- return nil, errors.New(`param: handlers unqualified`)
- }
- isSurviving := new(atomic.Bool)
- isSurviving.Store(false)
- c := &client{
- logger: logger,
- isSurviving: isSurviving,
- routerCodec: routerCodec,
- url: url,
- send: make(chan []byte, maxSendBuffer),
- handlers: components,
- }
- return c, nil
- }
- func (c *client) acceptLoop() {
- defer func() {
- _ = c.session.Conn.Close()
- c.isSurviving.Store(false)
- }()
- _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
- c.session.Conn.SetPongHandler(func(string) error {
- _ = c.session.Conn.SetReadDeadline(time.Now().Add(pongWait))
- return nil
- })
- for c.session.Conn != nil {
- _, data, err := c.session.Conn.ReadMessage()
- if err != nil {
- break
- }
- c.onReceive(data)
- }
- }
- func (c *client) sendLoop() {
- ticker := time.NewTicker(pingPeriod)
- defer func() {
- ticker.Stop()
- _ = c.session.Conn.Close()
- c.isSurviving.Store(false)
- close(c.send)
- }()
- for {
- select {
- case msg := <-c.send:
- _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
- if err := c.session.Conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
- return
- }
- case <-ticker.C:
- _ = c.session.Conn.SetWriteDeadline(time.Now().Add(writeWait))
- if err := c.session.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
- return
- }
- }
- }
- }
- func (c *client) onReceive(msg []byte) {
- }
- func (c *client) Send(msg []byte) (err error) {
- defer func() {
- if e := recover(); e != nil {
- err = ErrBrokenPipe
- }
- }()
- if !c.isSurviving.Load() {
- return ErrBrokenPipe
- }
- if len(c.send) >= maxSendBuffer {
- return ErrBufferPoolExceed
- }
- if len(msg) > maxFrameMessageLen {
- return
- }
- c.send <- msg
- return nil
- }
- func (c *client) Connect(requestHeader http.Header) (err error) {
- c.session.Conn, _, err = websocket.DefaultDialer.Dial(c.url.String(), requestHeader)
- if err != nil {
- return fmt.Errorf("dial: %s", err)
- }
- go c.acceptLoop()
- go c.sendLoop()
- c.isSurviving.Store(true)
- return nil
- }
- func (c *client) Disconnect() {
- }
|