我们知道我们可以直接用websocket.Conn进行ws通讯，可是为什么要对他进行二次封装，主要有两个原因。
1，不方便分模块调用
2，websocket.Conn中的很多方法是非线程安全的
线程安全的意思是当多线程调用同一模块时，相对安全的不会出错。
js对ws的简单封装
首先我们稍微优化一下测试前端。
如果不对WebSocket对象封装的话，我们无法手动控制连接的状态。
<!doctype html>
<html>
<head>
<meta charset="utf-8"/>
<title>WebSocket</title>
</head>
<body>
<button type="button" onclick="connnectWS()">建立连接</button>
<button type="button" onclick="sendMeg()">发送消息</button>
<button type="button" onclick="connClose()">关闭连接</button>
<div id="content"></div>
<script>
var text=document.getElementById("content");
var link="ws://localhost:6060/ws";
var conn={
connect:function(link) {
this.ws=new WebSocket(link);
//连接打开时触发
this.ws.onopen = function(evt) {
console.log("Connection open ...");
};
//接收到消息时触发
this.ws.onmessage = function(evt) {
text.innerHTML=text.innerHTML+"<br />"+evt.data;
};
//连接关闭时触发
this.ws.onclose = function(evt) {
console.log("Connection closed.");
};
}
}
function connnectWS(){
conn.connect(link);
}
function sendMeg(){
conn.ws.send("Hello WebSocket！");
}
function connClose(){
conn.ws.close();
}
</script>
</body>
</html>
Go对ws的简单封装
我们先看看封装结构，之后我们在一个一个看方法。
type Connection struct {
wsConn *websocket.Conn //websocket.Conn对象
inChan chan []byte //用于接收消息
outChan chan []byte //用于发送消息
closeChan chan byte //帮助内部逻辑判断连接是否被中断
mutex sync.Mutex //用于加锁
isClose bool //用于避免重复关闭中的不安全因素
}
下面是“构造方法”
const (
MAXMEGNUM = 1000 //最大消息数
)
//NewWS 构造一个WSConn
func NewWS(wsConn *websocket.Conn) (conn *Connection, err error) {
conn = &Connection{
wsConn: wsConn,
inChan: make(chan []byte, MAXMEGNUM),
outChan: make(chan []byte, MAXMEGNUM),
closeChan: make(chan byte, 1),
}
go conn.readLoop() //执行内部读逻辑
go conn.writeLoop() //执行内部写逻辑
return
}
读取和写单条消息的逻辑
//ReadMessage 读取一条消息，即接收
func (conn *Connection) ReadMessage() (data []byte, err error) {
select {
case data = <-conn.inChan:
case <-conn.closeChan:
err = errors.New("连接已关闭。")
}
return
}
//WriteMessage 写一条消息，即发送
func (conn *Connection) WriteMessage(data []byte) (err error) {
select {
case conn.outChan <- data:
case <-conn.closeChan:
err = errors.New("连接已关闭。")
}
return
}
关闭连接的方法
func (conn *Connection) Close() {
conn.wsConn.Close() //这个方法是线程安全的
conn.mutex.Lock() //上锁
//对Chan进行处理，并标记连接已关闭
if conn.isClose {
close(conn.closeChan)
conn.isClose = true
}
conn.mutex.Unlock() //解锁
}
内部实现的读写方法。避免死循环，我们使用select语句监听连接状态！
func (conn *Connection) readLoop() {
for {
megType, data, err := conn.wsConn.ReadMessage()
//TODO 消息类型处理
megType = megType
if err != nil {
util.ErrHandle(err)
conn.Close()
return
}
select {
case conn.inChan <- data:
case <-conn.closeChan:
conn.Close()
return
}
}
}
func (conn *Connection) writeLoop() {
var data []byte
for {
select {
case data = <-conn.outChan:
case <-conn.closeChan:
conn.Close()
return
}
//TODO 消息类型处理
err := conn.wsConn.WriteMessage(websocket.TextMessage, data)
if err != nil {
util.ErrHandle(err)
conn.Close()
return
}
}
}
这样我们就封装了一个相对安全的ws封装，发现内容太多，就不掩饰了。下面给一个测试代码，有兴趣的朋友可以自己试一下。
func main() {
r := gin.Default()
r.GET("/ws", func(c *gin.Context) {
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
return true
}}
wsConn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
defer wsConn.Close()
if err != nil {
fmt.Println(err)
return
}
conn, err := impl.NewWS(wsConn)
defer conn.Close()
if err != nil {
fmt.Println(err)
return
}
go func() {
for {
err := conn.WriteMessage([]byte("heartbeat"))
if err != nil {
fmt.Println(err)
return
}
time.Sleep(time.Second * 3)
}
}()
for {
data, err := conn.ReadMessage()
if err != nil {
fmt.Println(err)
return
}
err = conn.WriteMessage(data)
if err != nil {
fmt.Println(err)
return
}
}
})
r.Run(":6060")
}
有疑问加站长微信联系（非本文作者）
