164 lines
2.7 KiB
Go
164 lines
2.7 KiB
Go
// ************************************
|
||
// @package: websocketServer
|
||
// @description: websocket环境
|
||
// @author:
|
||
// @revision history:
|
||
// @create date: 2022-02-18 15:59:08
|
||
// ************************************
|
||
package websocketServer
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gorilla/websocket"
|
||
webServer "Framework/webServer"
|
||
)
|
||
|
||
// Context
|
||
// @description: websocket环境
|
||
type Context struct {
|
||
ctx context.Context
|
||
|
||
// web_server环境
|
||
webServerCtx *webServer.Context
|
||
|
||
// 存放用户自定义数据
|
||
userData interface{}
|
||
|
||
// websocket连接
|
||
conn *websocket.Conn
|
||
|
||
// 防止多协程同时写/关闭
|
||
mu sync.Mutex
|
||
|
||
// 最近一次收到心跳包的时间
|
||
heartbeat time.Time
|
||
|
||
// 指示是否已经关闭
|
||
isClosed bool
|
||
|
||
// 关闭连接信号
|
||
cls chan<- struct{}
|
||
}
|
||
|
||
// GetWebServerContext
|
||
// @description: 获取web_server环境
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
//
|
||
// @*web_server.Context: web_server环境
|
||
func (ctx *Context) GetWebServerContext() *webServer.Context {
|
||
return ctx.webServerCtx
|
||
}
|
||
|
||
// GetUserData
|
||
// @description: 获取用户自定义数据
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
//
|
||
// @interface{}:
|
||
func (ctx *Context) GetUserData() interface{} {
|
||
return ctx.userData
|
||
}
|
||
|
||
// SetUserData
|
||
// @description: 设置用户自定义数据
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
// @userData:
|
||
//
|
||
// return:
|
||
func (ctx *Context) SetUserData(userData interface{}) {
|
||
ctx.userData = userData
|
||
}
|
||
|
||
// SendMessage
|
||
// @description: 发送websocket数据
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
// @messageType:
|
||
// @data:
|
||
//
|
||
// return:
|
||
//
|
||
// @err:
|
||
func (ctx *Context) SendMessage(messageType int, data []byte) (err error) {
|
||
ctx.mu.Lock()
|
||
defer ctx.mu.Unlock()
|
||
|
||
return ctx.conn.WriteMessage(messageType, data)
|
||
}
|
||
|
||
// Close
|
||
// @description: 关闭websocket连接
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
func (ctx *Context) Close() {
|
||
ctx.mu.Lock()
|
||
defer ctx.mu.Unlock()
|
||
|
||
if !ctx.isClosed {
|
||
// 设置已关闭标志
|
||
ctx.isClosed = true
|
||
// 发送退出信号
|
||
// 所有向cls写入，都使用select超时结构，以保证这儿不会一直阻塞，确保此协程能退出
|
||
select {
|
||
case <-time.After(time.Millisecond * 10):
|
||
case ctx.cls <- struct{}{}:
|
||
}
|
||
//ctx.cls <- struct{}{}
|
||
}
|
||
}
|
||
|
||
// Ctx
|
||
// @description: 返回Context
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
func (ctx *Context) Ctx() context.Context {
|
||
ctx.mu.Lock()
|
||
defer ctx.mu.Unlock()
|
||
|
||
return ctx.ctx
|
||
}
|
||
|
||
// SetCtx
|
||
// @description: 设置新的Ctx
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
func (ctx *Context) SetCtx(c context.Context) {
|
||
ctx.mu.Lock()
|
||
defer ctx.mu.Unlock()
|
||
|
||
ctx.ctx = c
|
||
}
|
||
|
||
// GetConnType
|
||
// @description: 获取连接类型
|
||
// parameter:
|
||
//
|
||
// @receiver ctx:
|
||
//
|
||
// return:
|
||
func (ctx *Context) GetConnType() string {
|
||
return "websocket"
|
||
}
|