goProject/trunk/framework/websocketServer/context.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"
}