base-golang/pkg/websocket/peer/connect/connection.go

package connect

import (
	"errors"
	"time"

	"github.com/gorilla/websocket"
	"git.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
)

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")
	ErrSendBufferLimit  = errors.New("send buffer limit")
)

type wsConnection struct {
	peer.ConnectionIdentify
	p       *peer.SessionManager
	conn    *websocket.Conn
	send    chan []byte
	running bool
}

func (ws *wsConnection) Peer() *peer.SessionManager {
	return ws.p
}

func (ws *wsConnection) Raw() interface{} {
	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) recvLoop() {
	defer func() {
		ws.p.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.p.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
}

func (ws *wsConnection) init() {
	go ws.recvLoop()
	go ws.sendLoop()
	ws.running = true
}

func newConnection(conn *websocket.Conn,
	p *peer.SessionManager) *wsConnection {
	wsc := &wsConnection{
		conn: conn,
		p:    p,
		send: make(chan []byte, maxSendBuffer),
	}
	wsc.init()
	return wsc
}