base-golang/pkg/websocket/client/client.go

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() {

}