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