15 Star 90 Fork 25

konyshe/gogo

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
WSConnection.go 3.20 KB
一键复制 编辑 原始数据 按行查看 历史
konyshe 提交于 2019-07-11 11:03 . 修改websocket客户端分组接口
package gogo
import (
"errors"
"net"
"sync"
"github.com/gorilla/websocket"
)
type ImplConnection struct {
// 存放websocket连接
wsConn *websocket.Conn
// 用于存放数据
inChan chan []byte
// 用于读取数据
outChan chan []byte
closeChan chan byte
mutex sync.Mutex
// chan是否被关闭
isClosed bool
// 会话存储数据
session interface{}
groupName string
}
// SetGroup 设置指定websocket客户端的分组
// conn websocket客户端连接句柄
// group 分组名称
func (conn *ImplConnection) SetGroup(groupName string) {
conn.groupName = groupName
}
// GetGroup 获取websocket客户端的分组信息
func (conn *ImplConnection) GetGroup() string {
return conn.groupName
}
// OutGroup 设置指定websocket客户端的分组
// conn websocket客户端连接句柄
// group 分组名称
func (conn *ImplConnection) OutGroup() {
conn.groupName = ""
}
// RemoteAddr 获取客户端IP
func (conn *ImplConnection) RemoteAddr() net.Addr {
return conn.wsConn.RemoteAddr()
}
// ReadMessage 读取数据
func (conn *ImplConnection) ReadMessage() (data []byte, err error) {
//select是Go中的一个控制结构,类似于用于通信的switch语句。每个case必须是一个通信操作,要么是发送要么是接收。 select随机执行一个可运行的case。如果没有case可运行,它将阻塞,直到有case可运行。一个默认的子句应该总是可运行的。
select {
case data = <-conn.inChan:
case <-conn.closeChan:
err = errors.New("connection is closed")
}
return
}
// WriteMessage 发送数据
func (conn *ImplConnection) WriteMessage(data []byte) (err error) {
select {
case conn.outChan <- data:
case <-conn.closeChan:
err = errors.New("connection is closed")
}
return
}
// Close 关闭连接
func (conn *ImplConnection) Close() {
// 线程安全的Close,可以并发多次调用也叫做可重入的Close
conn.wsConn.Close()
conn.mutex.Lock()
if !conn.isClosed {
// 关闭chan,但是chan只能关闭一次
close(conn.closeChan)
conn.isClosed = true
}
conn.mutex.Unlock()
}
// InitImplConnection 初始化长连接
func InitImplConnection(wsConn *websocket.Conn) (conn *ImplConnection, err error) {
conn = &ImplConnection{
wsConn: 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 *ImplConnection) readLoop() {
var (
data []byte
err error
)
for {
if _, data, err = conn.wsConn.ReadMessage(); err != nil {
goto ERR
}
// 容易阻塞到这里,等待inChan有空闲的位置
select {
case conn.inChan <- data:
case <-conn.closeChan: // closeChan关闭的时候执行
goto ERR
}
}
ERR:
conn.Close()
}
func (conn *ImplConnection) writeLoop() {
var (
data []byte
err error
)
for {
select {
case data = <-conn.outChan:
case <-conn.closeChan:
goto ERR
}
if err = conn.wsConn.WriteMessage(websocket.TextMessage, data); err != nil {
goto ERR
}
}
ERR:
conn.Close()
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Go
1
https://gitee.com/konyshe/gogo.git
git@gitee.com:konyshe/gogo.git
konyshe
gogo
gogo
v2

搜索帮助