first commit

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
+}