代码拉取完成,页面将自动刷新
package websocket
import (
"errors"
"fmt"
"github.com/gorilla/websocket"
"sync"
)
type Connection struct {
wsConnect *websocket.Conn
inChan chan []byte
outChan chan []byte
closeChan chan byte
mutex sync.Mutex // 对closeChan关闭上锁
isClosed bool // 防止closeChan被关闭多次
}
func InitConnection(wsConn *websocket.Conn) (conn *Connection, err error) {
conn = &Connection{
wsConnect: wsConn,
inChan: make(chan []byte, 1000),
outChan: make(chan []byte, 1000),
closeChan: make(chan byte, 1),
}
// 启动读协程
go conn.readLoop()
// 启动写协程
go conn.writeLoop()
return
}
func (conn *Connection) ReadMessage() (data []byte, err error) {
select {
case data = <-conn.inChan:
case <-conn.closeChan:
err = errors.New("connection is closeed")
}
return
}
func (conn *Connection) WriteMessage(data []byte) (err error) {
select {
case conn.outChan <- data:
case <-conn.closeChan:
err = errors.New("connection is closeed")
}
return
}
func (conn *Connection) Close() {
// 线程安全,可多次调用
conn.wsConnect.Close()
// 利用标记,让closeChan只关闭一次
conn.mutex.Lock()
if !conn.isClosed {
close(conn.closeChan)
conn.isClosed = true
}
conn.mutex.Unlock()
}
// 内部实现
func (conn *Connection) readLoop() {
var (
data []byte
err error
)
for {
if _, data, err = conn.wsConnect.ReadMessage(); err != nil {
goto ERR
}
//阻塞在这里,等待inChan有空闲位置
select {
case conn.inChan <- data:
case <-conn.closeChan: // closeChan 感知 conn断开
goto ERR
}
}
ERR:
conn.Close()
}
func (conn *Connection) writeLoop() {
var (
data []byte
err error
)
for {
select {
case data = <-conn.outChan:
case <-conn.closeChan:
fmt.Println("发送失败")
goto ERR
}
fmt.Println("开始发送")
if err = conn.wsConnect.WriteMessage(websocket.TextMessage, data); err != nil {
goto ERR
}
}
ERR:
conn.Close()
}
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。