You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
1.8 KiB
100 lines
1.8 KiB
package websocketpack
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
/*
|
|
*
|
|
@ 作者: 秦东
|
|
@ 时间: 2025-06-12 16:19:04
|
|
@ 功能: 初始化设置
|
|
*/
|
|
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
|
|
}
|
|
|
|
/**
|
|
@ 作者: 秦东
|
|
@ 时间: 2025-06-12 16:21:26
|
|
@ 功能: 读取消息
|
|
*/
|
|
|
|
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:
|
|
goto ERR
|
|
}
|
|
if err = conn.WsConnect.WriteMessage(websocket.TextMessage, data); err != nil {
|
|
goto ERR
|
|
}
|
|
}
|
|
ERR:
|
|
conn.Close()
|
|
|
|
}
|
|
|