multimedia/server/webrtc_easy_signal/ws_conection.go

package main

import (
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// 允许等待的写入时间
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 0
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var maxConnId int64

// 用于广播
var WsConnAll map[int64]*WsConnection

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许所有的CORS 跨域请求，正式环境可以关闭
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 客户端读写消息
type wsMessage struct {
	// websocket.TextMessage 消息类型
	messageType int
	data        []byte
}

// 客户端连接
type WsConnection struct {
	wsSocket *websocket.Conn // 底层websocket
	inChan   chan *wsMessage // 读队列
	outChan  chan *wsMessage // 写队列

	mutex             sync.Mutex // 避免重复关闭管道,加锁处理
	isClosed          bool
	closeChan         chan byte // 关闭通知
	id                int64
	lastHeartBeatTime time.Time // 上次心跳包时间
	needHeartBeat     bool      // 上次心跳
}

func init() {
	maxConnId = 0
}

// 读取消息队列中的消息
func WsHandler(resp http.ResponseWriter, req *http.Request) {
	// 应答客户端告知升级连接为websocket
	wsSocket, err := upgrader.Upgrade(resp, req, nil)
	if err != nil {
		log.Println("升级为websocket失败", err.Error())
		return
	}
	maxConnId++
	// TODO 如果要控制连接数可以计算，wsConnAll长度
	// 连接数保持一定数量，超过的部分不提供服务
	wsConn := &WsConnection{
		wsSocket:          wsSocket,
		inChan:            make(chan *wsMessage, 50),
		outChan:           make(chan *wsMessage, 50),
		closeChan:         make(chan byte),
		isClosed:          false,
		id:                maxConnId,
		lastHeartBeatTime: time.Now(),
		needHeartBeat:     false,
	}
	WsConnAll[maxConnId] = wsConn
	log.Println("当前在线人数", len(WsConnAll))
	// 处理器,发送定时信息，避免意外关闭
	go wsConn.processLoop()
	// 读协程
	go wsConn.wsReadLoop()
	// 写协程
	go wsConn.wsWriteLoop()
}

// 读取消息队列中的消息
func (wsConn *WsConnection) wsRead() (*wsMessage, error) {
	tick := time.Tick(time.Second * 1)
	select {
	case msg := <-wsConn.inChan:
		// 获取到消息队列中的消息
		return msg, nil
	case <-wsConn.closeChan:
	case <-tick:
		return nil, errors.New("超时")
	}
	return nil, errors.New("连接已经关闭")
}

// 处理队列中的消息
func (this *WsConnection) PayloadParseAndCallback(dat []byte) error {
	var payload Request
	e := json.Unmarshal(dat, &payload)
	if nil != e {
		log.Print(e.Error())
		return e
	}
	log.Print(payload)
	return nil
}

// 处理队列中的消息
func (wsConn *WsConnection) processLoop() {
	// 处理消息队列中的消息
	// 获取到消息队列中的消息，处理完成后，发送消息给客户端
	for {
		msg, err := wsConn.wsRead()
		if err != nil {
			if err.Error() == "超时" {
				// log.Print(wsConn.lastHeartBeatTime.String())
				if wsConn.needHeartBeat {
					if time.Now().Sub(wsConn.lastHeartBeatTime) > time.Second*15 {
						log.Print("心跳超时")
						wsConn.close()
					}
				}
				continue
			}
			break
		}
		ProtoCallBack(wsConn, msg.data)
	}
}

type Response struct {
	Type    int         `json:"type"`
	Payload interface{} `json:"data"`
}

// 发送信息
func (wsConn *WsConnection) SendPayload(ptype int64, v interface{}) error {
	var resp Response
	resp.Type = int(ptype)
	resp.Payload = v
	bytes, e := json.Marshal(resp)
	if nil != e {
		log.Print(e.Error())
	}
	wsConn.wsWrite(1, bytes)
	return nil
}

// 写入消息到队列中
func (wsConn *WsConnection) wsWrite(messageType int, data []byte) error {
	select {
	case wsConn.outChan <- &wsMessage{messageType, data}:
	case <-wsConn.closeChan:
		return errors.New("连接已经关闭")
	}
	return nil
}

// 处理消息队列中的消息
func (wsConn *WsConnection) wsReadLoop() {
	// 设置消息的最大长度
	wsConn.wsSocket.SetReadLimit(maxMessageSize)
	for {
		// 读一个message
		msgType, data, err := wsConn.wsSocket.ReadMessage()
		req := &wsMessage{
			msgType,
			data,
		}
		if err != nil {
			websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure)
			log.Println("消息读取出现错误", err.Error())

			wsConn.close()
		}
		log.Print(string(data))
		// 放入请求队列,消息入栈
		select {
		case wsConn.inChan <- req:
		case <-wsConn.closeChan:
			return
		}
	}
}

// 发送消息给客户端
func (wsConn *WsConnection) wsWriteLoop() {
	for {
		select {
		// 取一个应答
		case msg := <-wsConn.outChan:
			// 写给websocket
			if err := wsConn.wsSocket.WriteMessage(msg.messageType, msg.data); err != nil {
				log.Println("发送消息给客户端发生错误", err.Error())
				// 切断服务
				wsConn.close()
				return
			}
		case <-wsConn.closeChan:
			// 获取到关闭通知
			return
		}
	}
}

// 关闭连接
func (wsConn *WsConnection) close() {
	log.Println("关闭连接被调用了")
	wsConn.wsSocket.Close()
	wsConn.mutex.Lock()

	defer wsConn.mutex.Unlock()
	if wsConn.isClosed == false {
		wsConn.isClosed = true
		// 删除这个连接的变量
		delete(WsConnAll, wsConn.id)
		close(wsConn.closeChan)
	}
}