first commit

pkg/websocket/peer/connect/acceptor.go

@@ -0,0 +1,98 @@
+package connect
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"git.bvbej.com/bvbej/base-golang/pkg/mux"
+	"github.com/gin-gonic/gin"
+	"go.uber.org/zap"
+	"net/http"
+	"net/url"
+	"time"
+
+	"git.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
+	"github.com/gorilla/websocket"
+)
+
+var _ WsAcceptor = (*wsAcceptor)(nil)
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
+type WsAcceptor interface {
+	Start(addr string) error
+	Stop()
+	GinHandle(ctx *gin.Context)
+	HandlerFunc() mux.HandlerFunc
+}
+
+type wsAcceptor struct {
+	server  *http.Server
+	sessMgr *peer.SessionManager
+	logger  *zap.Logger
+}
+
+func NewWsAcceptor(sessMgr *peer.SessionManager, loggers *zap.Logger) WsAcceptor {
+	return &wsAcceptor{
+		sessMgr: sessMgr,
+		logger:  loggers,
+	}
+}
+
+func (ws *wsAcceptor) Start(addr string) error {
+	urlObj, err := url.Parse(addr)
+	if err != nil {
+		return fmt.Errorf("websocket urlparse failed. url(%s) %v", addr, err)
+	}
+	if urlObj.Path == "" {
+		return fmt.Errorf("websocket start failed. expect path in url to listen addr:%s", addr)
+	}
+
+	http.HandleFunc(urlObj.Path, func(w http.ResponseWriter, r *http.Request) {
+		c, upgradeErr := upgrader.Upgrade(w, r, nil)
+		if upgradeErr != nil {
+			ws.logger.Sugar().Errorf("upgrade http failed: %s", upgradeErr)
+			return
+		}
+		ws.sessMgr.Register <- peer.NewSession(NewConnection(c, ws.sessMgr))
+	})
+
+	ws.server = &http.Server{Addr: urlObj.Host}
+	err = ws.server.ListenAndServe()
+	if err != nil && !errors.Is(err, http.ErrServerClosed) {
+		return fmt.Errorf("websocket ListenAndServe addr:%s failed:%v", addr, err)
+	}
+
+	return nil
+}
+
+func (ws *wsAcceptor) Stop() {
+	ws.sessMgr.CloseAllSession()
+
+	if ws.server != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+		defer cancel()
+		if err := ws.server.Shutdown(ctx); err != nil {
+			ws.logger.Sugar().Errorf("server shutdown err:[%s]", err)
+		}
+	}
+}
+
+func (ws *wsAcceptor) GinHandle(ctx *gin.Context) {
+	c, upgradeErr := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
+	if upgradeErr != nil {
+		ws.logger.Sugar().Errorf("upgrade http failed: %s", upgradeErr)
+		return
+	}
+	ws.sessMgr.Register <- peer.NewSession(NewConnection(c, ws.sessMgr))
+}
+
+func (ws *wsAcceptor) HandlerFunc() mux.HandlerFunc {
+	return func(c mux.Context) {
+		ws.GinHandle(c.Context())
+	}
+}

pkg/websocket/peer/connect/connection.go

@@ -0,0 +1,134 @@
+package connect
+
+import (
+	"errors"
+	"time"
+
+	"git.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
+	"github.com/gorilla/websocket"
+)
+
+const (
+	writeWait          = 20 * time.Second
+	pongWait           = 60 * time.Second
+	pingPeriod         = (pongWait * 9) / 10
+	maxFrameMessageLen = 16 * 1024 //4 * 4096
+	maxSendBuffer      = 16
+)
+
+var (
+	ErrBrokenPipe       = errors.New("send to broken pipe")
+	ErrBufferPoolExceed = errors.New("send buffer exceed")
+)
+
+type wsConnection struct {
+	peer.ConnectionIdentify
+	pm      *peer.SessionManager
+	conn    *websocket.Conn
+	send    chan []byte
+	running bool
+}
+
+func NewConnection(conn *websocket.Conn, p *peer.SessionManager) *wsConnection {
+	wsc := &wsConnection{
+		conn:    conn,
+		pm:      p,
+		send:    make(chan []byte, maxSendBuffer),
+		running: true,
+	}
+
+	go wsc.acceptLoop()
+	go wsc.sendLoop()
+
+	return wsc
+}
+
+func (ws *wsConnection) Peer() *peer.SessionManager {
+	return ws.pm
+}
+
+func (ws *wsConnection) Raw() any {
+	if ws.conn == nil {
+		return nil
+	}
+	return ws.conn
+}
+
+func (ws *wsConnection) RemoteAddr() string {
+	if ws.conn == nil {
+		return ""
+	}
+	return ws.conn.RemoteAddr().String()
+}
+
+func (ws *wsConnection) Close() {
+	_ = ws.conn.Close()
+	ws.running = false
+}
+
+func (ws *wsConnection) Send(msg []byte) (err error) {
+	defer func() {
+		if e := recover(); e != nil {
+			err = ErrBrokenPipe
+		}
+	}()
+	if !ws.running {
+		return ErrBrokenPipe
+	}
+	if len(ws.send) >= maxSendBuffer {
+		return ErrBufferPoolExceed
+	}
+	if len(msg) > maxFrameMessageLen {
+		return
+	}
+	ws.send <- msg
+	return nil
+}
+
+func (ws *wsConnection) acceptLoop() {
+	defer func() {
+		ws.pm.Unregister <- ws.ID()
+		_ = ws.conn.Close()
+		ws.running = false
+	}()
+	_ = ws.conn.SetReadDeadline(time.Now().Add(pongWait))
+	ws.conn.SetPongHandler(func(string) error {
+		_ = ws.conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+	for ws.conn != nil {
+		_, data, err := ws.conn.ReadMessage()
+		if err != nil {
+			break
+		}
+		ws.pm.ProcessMessage(ws.ID(), data)
+	}
+}
+
+func (ws *wsConnection) sendLoop() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		_ = ws.conn.Close()
+		ws.running = false
+		close(ws.send)
+	}()
+	for {
+		select {
+		case msg := <-ws.send:
+			_ = ws.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := ws.conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
+				return
+			}
+		case <-ticker.C:
+			_ = ws.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := ws.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}
+
+func (ws *wsConnection) IsClosed() bool {
+	return !ws.running
+}