feat(database): 添加数据库连接初始化和ORM支持
- 实现MySQL和SQLite数据库连接初始化功能 - 集成GORM ORM框架支持 - 添加数据库连接状态检测和错误处理 - 实现数据库自动迁移功能 - 添加字段重命名和删除功能 - 更新配置文件以支持数据库类型选择 - 添加TCP服务器示例代码 - 添加WebSocket服务示例代码 - 集成MQTT客户端功能 - 更新项目依赖包列表
This commit is contained in:
@@ -0,0 +1,73 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
)
|
||||
|
||||
var manager = NewWs()
|
||||
|
||||
func NewWs() *Manager {
|
||||
// 1. 自定义配置（可选，也可使用默认配置）
|
||||
customConfig := &Config{
|
||||
AllowAllOrigins: true,
|
||||
HeartbeatInterval: 20 * time.Second, // 20秒发一次心跳
|
||||
HeartbeatTimeout: 40 * time.Second, // 40秒超时
|
||||
}
|
||||
|
||||
// 2. 创建管理器
|
||||
m := NewManager(customConfig)
|
||||
|
||||
// 3. 覆盖业务回调（核心：自定义消息处理逻辑）
|
||||
// 连接建立回调
|
||||
m.OnConnect = func(connID string) {
|
||||
log.Printf("业务回调：连接[%s]上线，当前在线数：%d", connID, m.GetOnlineCount())
|
||||
// 欢迎消息
|
||||
_ = m.SendToConn(connID, []byte("欢迎连接WebSocket服务！"))
|
||||
}
|
||||
|
||||
// 收到消息回调
|
||||
m.OnMessage = func(connID string, msgType int, data any) {
|
||||
log.Printf("业务回调：收到连接[%s]消息：%s", connID, gconv.String(data))
|
||||
// 示例：echo回复
|
||||
reply := []byte("服务端回复：" + gconv.String(data))
|
||||
_ = m.SendToConn(connID, reply)
|
||||
|
||||
// 示例：广播消息给所有连接
|
||||
_ = m.Broadcast([]byte("广播：" + connID + "说：" + gconv.String(data)))
|
||||
}
|
||||
|
||||
// 连接断开回调
|
||||
m.OnDisconnect = func(connID string, err error) {
|
||||
log.Printf("业务回调：连接[%s]下线，原因：%v，当前在线数：%d", connID, err, m.GetOnlineCount())
|
||||
}
|
||||
return m
|
||||
}
|
||||
func Upgrade(w http.ResponseWriter, r *http.Request, connID string) {
|
||||
_, err := manager.Upgrade(w, r, connID)
|
||||
if err != nil {
|
||||
log.Printf("升级连接失败：%v", err)
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
}
|
||||
func main() {
|
||||
// 4. 注册WebSocket路由
|
||||
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
|
||||
// 自定义连接ID（示例：使用请求参数中的user_id）
|
||||
connID := r.URL.Query().Get("user_id")
|
||||
if connID == "" {
|
||||
http.Error(w, "user_id不能为空", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
// 升级连接
|
||||
Upgrade(w, r, connID)
|
||||
})
|
||||
|
||||
// 5. 启动服务
|
||||
log.Println("WebSocket服务启动：http://localhost:8080/ws")
|
||||
log.Fatal(http.ListenAndServe(":8080", nil))
|
||||
}
|
||||
+488
@@ -0,0 +1,488 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gogf/gf/v2/encoding/gjson"
|
||||
"github.com/gogf/gf/v2/os/gctx"
|
||||
"github.com/gogf/gf/v2/os/gtime"
|
||||
"github.com/gogf/gf/v2/os/gtimer"
|
||||
"github.com/gogf/gf/v2/text/gstr"
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// 常量定义：默认配置
|
||||
const (
|
||||
// DefaultReadBufferSize 默认读写缓冲区大小（字节）
|
||||
DefaultReadBufferSize = 1024
|
||||
DefaultWriteBufferSize = 1024
|
||||
// DefaultHeartbeatInterval 默认心跳间隔（秒）：每30秒发送一次心跳
|
||||
DefaultHeartbeatInterval = 30 * time.Second
|
||||
// DefaultHeartbeatTimeout 默认心跳超时（秒）：60秒未收到客户端心跳响应则关闭连接
|
||||
DefaultHeartbeatTimeout = 60 * time.Second
|
||||
// DefaultReadTimeout 默认读写超时（秒）
|
||||
DefaultReadTimeout = 60 * time.Second
|
||||
DefaultWriteTimeout = 10 * time.Second
|
||||
// MessageTypeText 消息类型
|
||||
MessageTypeText = websocket.TextMessage
|
||||
MessageTypeBinary = websocket.BinaryMessage
|
||||
// HeartbeatMaxRetry 心跳最大重试次数
|
||||
HeartbeatMaxRetry = 3
|
||||
)
|
||||
|
||||
// Config WebSocket服务端配置
|
||||
type Config struct {
|
||||
// 读写缓冲区大小
|
||||
ReadBufferSize int
|
||||
WriteBufferSize int
|
||||
// 跨域配置：是否允许所有跨域（生产环境建议指定Origin）
|
||||
AllowAllOrigins bool
|
||||
// 允许的跨域Origin列表（AllowAllOrigins=false时生效）
|
||||
AllowedOrigins []string
|
||||
// 心跳配置
|
||||
HeartbeatInterval time.Duration // 心跳发送间隔
|
||||
HeartbeatTimeout time.Duration // 心跳超时时间
|
||||
// 读写超时
|
||||
ReadTimeout time.Duration
|
||||
WriteTimeout time.Duration
|
||||
MsgType int // 发送消息的默认类型
|
||||
HeartbeatValue string // 心跳消息的标识字段值（如"heartbeat"、"pong"）
|
||||
HeartbeatKey string // 心跳消息的标识字段名（如"type"）
|
||||
}
|
||||
|
||||
// 默认配置
|
||||
func DefaultConfig() *Config {
|
||||
return &Config{
|
||||
ReadBufferSize: DefaultReadBufferSize,
|
||||
WriteBufferSize: DefaultWriteBufferSize,
|
||||
AllowAllOrigins: true,
|
||||
AllowedOrigins: []string{},
|
||||
HeartbeatInterval: DefaultHeartbeatInterval,
|
||||
HeartbeatTimeout: DefaultHeartbeatTimeout,
|
||||
ReadTimeout: DefaultReadTimeout,
|
||||
WriteTimeout: DefaultWriteTimeout,
|
||||
MsgType: MessageTypeText,
|
||||
HeartbeatValue: "heartbeat",
|
||||
HeartbeatKey: "type", // 心跳消息的标识字段名，默认"type"
|
||||
}
|
||||
}
|
||||
|
||||
// Connection WebSocket连接结构体
|
||||
type Connection struct {
|
||||
conn *websocket.Conn // 底层连接
|
||||
connID string // 唯一连接ID
|
||||
manager *Manager // 所属管理器
|
||||
createTime time.Time // 连接创建时间
|
||||
heartbeatChan time.Time // 心跳通道（用于检测客户端响应）
|
||||
heartbeatTime *gtimer.Entry
|
||||
ctx context.Context // 上下文
|
||||
cancel context.CancelFunc // 上下文取消函数
|
||||
writeMutex sync.Mutex // 写消息互斥锁（防止并发写）
|
||||
heartbeatRetry int // 心跳发送重试次数
|
||||
}
|
||||
|
||||
// Manager WebSocket连接管理器
|
||||
type Manager struct {
|
||||
config *Config // 配置
|
||||
upgrader *websocket.Upgrader // HTTP升级器
|
||||
connections map[string]*Connection // 所有在线连接（connID -> Connection）
|
||||
mutex sync.RWMutex // 读写锁（保护connections）
|
||||
// 业务回调：收到消息时触发（用户自定义处理逻辑）
|
||||
OnMessage func(connID string, msgType int, data any)
|
||||
// 业务回调：连接建立时触发
|
||||
OnConnect func(connID string)
|
||||
// 业务回调：连接关闭时触发
|
||||
OnDisconnect func(connID string, err error)
|
||||
}
|
||||
|
||||
// Merge 合并配置，用传入的配置覆盖非零值部分
|
||||
func (c *Config) Merge(other *Config) *Config {
|
||||
result := *c // 复制当前配置
|
||||
|
||||
if other == nil {
|
||||
return &result
|
||||
}
|
||||
|
||||
if other.ReadBufferSize > 0 {
|
||||
result.ReadBufferSize = other.ReadBufferSize
|
||||
}
|
||||
if other.WriteBufferSize > 0 {
|
||||
result.WriteBufferSize = other.WriteBufferSize
|
||||
}
|
||||
if other.HeartbeatInterval > 0 {
|
||||
result.HeartbeatInterval = other.HeartbeatInterval
|
||||
}
|
||||
if other.HeartbeatTimeout > 0 {
|
||||
result.HeartbeatTimeout = other.HeartbeatTimeout
|
||||
}
|
||||
if other.ReadTimeout > 0 {
|
||||
result.ReadTimeout = other.ReadTimeout
|
||||
}
|
||||
if other.WriteTimeout > 0 {
|
||||
result.WriteTimeout = other.WriteTimeout
|
||||
}
|
||||
if other.AllowAllOrigins {
|
||||
result.AllowAllOrigins = other.AllowAllOrigins
|
||||
}
|
||||
if other.HeartbeatValue != "" {
|
||||
result.HeartbeatValue = other.HeartbeatValue
|
||||
}
|
||||
if other.HeartbeatKey != "" {
|
||||
result.HeartbeatKey = other.HeartbeatKey
|
||||
}
|
||||
if len(other.AllowedOrigins) > 0 {
|
||||
result.AllowedOrigins = other.AllowedOrigins
|
||||
}
|
||||
if other.MsgType != 0 {
|
||||
result.MsgType = other.MsgType
|
||||
}
|
||||
|
||||
return &result
|
||||
}
|
||||
|
||||
// NewManager 创建连接管理器
|
||||
func NewManager(config *Config) *Manager {
|
||||
defaultConfig := DefaultConfig()
|
||||
finalConfig := defaultConfig.Merge(config)
|
||||
// 初始化升级器
|
||||
upgrader := &websocket.Upgrader{
|
||||
ReadBufferSize: config.ReadBufferSize,
|
||||
WriteBufferSize: config.WriteBufferSize,
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
// 跨域检查
|
||||
if config.AllowAllOrigins {
|
||||
return true
|
||||
}
|
||||
origin := r.Header.Get("Origin")
|
||||
for _, allowed := range finalConfig.AllowedOrigins {
|
||||
if origin == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
},
|
||||
}
|
||||
|
||||
return &Manager{
|
||||
config: finalConfig,
|
||||
upgrader: upgrader,
|
||||
connections: make(map[string]*Connection),
|
||||
mutex: sync.RWMutex{},
|
||||
// 默认回调（用户可覆盖）
|
||||
OnMessage: func(connID string, msgType int, data any) {
|
||||
log.Printf("[默认回调] 收到连接[%s]消息：%s", connID, gconv.String(data))
|
||||
},
|
||||
OnConnect: func(connID string) {
|
||||
log.Printf("[默认回调] 连接[%s]已建立", connID)
|
||||
},
|
||||
OnDisconnect: func(connID string, err error) {
|
||||
log.Printf("[默认回调] 连接[%s]已关闭：%v", connID, err)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Upgrade HTTP升级为WebSocket连接
|
||||
// connID：自定义连接唯一ID（如用户ID、设备ID）
|
||||
func (m *Manager) Upgrade(w http.ResponseWriter, r *http.Request, connID string) (*Connection, error) {
|
||||
if connID == "" {
|
||||
return nil, errors.New("连接ID不能为空")
|
||||
}
|
||||
|
||||
// 检查连接ID是否已存在
|
||||
m.mutex.RLock()
|
||||
_, exists := m.connections[connID]
|
||||
m.mutex.RUnlock()
|
||||
if exists {
|
||||
return nil, fmt.Errorf("连接ID[%s]已存在", connID)
|
||||
}
|
||||
|
||||
// 升级HTTP连接
|
||||
conn, err := m.upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("升级WebSocket失败：%w", err)
|
||||
}
|
||||
|
||||
// 创建上下文（用于优雅关闭）
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
// 创建连接实例
|
||||
wsConn := &Connection{
|
||||
conn: conn,
|
||||
connID: connID,
|
||||
manager: m,
|
||||
createTime: time.Now(),
|
||||
heartbeatChan: time.Now(), // 缓冲1，防止阻塞
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
writeMutex: sync.Mutex{},
|
||||
heartbeatRetry: 0,
|
||||
}
|
||||
wsConn.heartbeatTime = gtimer.AddSingleton(gctx.New(), m.config.HeartbeatTimeout, func(ctx context.Context) {
|
||||
log.Printf("[心跳检测] 连接[%s]已关闭：心跳超时", wsConn.connID)
|
||||
wsConn.heartbeatTime.Close()
|
||||
wsConn.heartbeatTime.Stop()
|
||||
wsConn.heartbeatTime = nil
|
||||
wsConn.ctx.Done()
|
||||
wsConn.Close(fmt.Errorf("心跳超时"))
|
||||
})
|
||||
// 添加到管理器
|
||||
m.mutex.Lock()
|
||||
m.connections[connID] = wsConn
|
||||
m.mutex.Unlock()
|
||||
|
||||
// 触发连接建立回调
|
||||
m.OnConnect(connID)
|
||||
|
||||
// 启动读消息协程
|
||||
go wsConn.ReadPump()
|
||||
// 启动写消息协程（处理异步发送）
|
||||
go wsConn.WritePump()
|
||||
// 启动心跳检测协程
|
||||
go wsConn.Heartbeat()
|
||||
|
||||
return wsConn, nil
|
||||
}
|
||||
|
||||
// ReadPump 读取客户端消息（持续运行）
|
||||
func (c *Connection) ReadPump() {
|
||||
defer func() {
|
||||
// 发生panic时关闭连接
|
||||
if err := recover(); err != nil {
|
||||
log.Printf("连接[%s]读消息协程panic：%v", c.connID, err)
|
||||
}
|
||||
// 关闭连接并清理
|
||||
c.Close(fmt.Errorf("读消息协程退出"))
|
||||
}()
|
||||
|
||||
// 循环读取消息
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return // 上下文已取消，退出
|
||||
default:
|
||||
// 设置读超时（每次读取前重置，防止长时间无消息超时）
|
||||
c.conn.SetReadDeadline(time.Now().Add(c.manager.config.ReadTimeout))
|
||||
// 读取客户端消息
|
||||
msgType, data, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
// 区分正常关闭和异常错误
|
||||
var closeErr *websocket.CloseError
|
||||
if errors.As(err, &closeErr) {
|
||||
c.Close(fmt.Errorf("客户端主动关闭：%s（代码：%d）", closeErr.Text, closeErr.Code))
|
||||
} else {
|
||||
c.Close(fmt.Errorf("读取消息失败：%w", err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 尝试解析JSON格式的心跳消息（精准判断，替代包含判断）
|
||||
isHeartbeat := false
|
||||
// 先尝试解析为JSON对象
|
||||
var msgMap map[string]interface{}
|
||||
if err := gjson.DecodeTo(data, &msgMap); err == nil {
|
||||
// 获取心跳标识字段的值
|
||||
heartbeatValue := gconv.String(msgMap[c.manager.config.HeartbeatKey])
|
||||
if heartbeatValue == c.manager.config.HeartbeatValue {
|
||||
isHeartbeat = true
|
||||
}
|
||||
} else {
|
||||
// 非JSON格式，降级为包含判断（兼容纯文本心跳）
|
||||
str := gconv.String(data)
|
||||
if gstr.Contains(str, c.manager.config.HeartbeatValue) {
|
||||
isHeartbeat = true
|
||||
}
|
||||
}
|
||||
if isHeartbeat {
|
||||
log.Printf("[心跳] 收到连接[%s]心跳消息：%s", c.connID, string(data))
|
||||
// 心跳消息：重置重试次数 + 发送心跳信号 + 重置读超时
|
||||
js, err := gjson.Encode(&Msg[any]{c.manager.config.HeartbeatValue, nil, gtime.Timestamp()})
|
||||
if err != nil {
|
||||
log.Printf("[心跳] 客户端[%s]json编码失败", c.connID)
|
||||
continue
|
||||
}
|
||||
err = c.Send(js)
|
||||
if err != nil {
|
||||
log.Printf("[心跳] 客户端[%s]发送心跳消息失败", c.connID)
|
||||
continue
|
||||
}
|
||||
c.heartbeatTime.Reset()
|
||||
continue // 跳过业务回调
|
||||
}
|
||||
|
||||
// 非心跳消息：触发业务回调
|
||||
c.manager.OnMessage(c.connID, msgType, data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Msg[T any] struct {
|
||||
Type string `json:"type"`
|
||||
Data T `json:"data"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
}
|
||||
|
||||
// WritePump 处理异步写消息（持续运行）
|
||||
// 扩展为监听写队列，防止消息丢失
|
||||
func (c *Connection) WritePump() {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
log.Printf("连接[%s]写消息协程panic：%v", c.connID, err)
|
||||
}
|
||||
}()
|
||||
|
||||
// 暂时保持简化，实际可扩展为带缓冲的写队列
|
||||
<-c.ctx.Done()
|
||||
}
|
||||
|
||||
// Heartbeat 心跳检测（持续运行）
|
||||
func (c *Connection) Heartbeat() {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
log.Printf("连接[%s]心跳协程panic：%v", c.connID, err)
|
||||
}
|
||||
}()
|
||||
c.heartbeatTime.Start()
|
||||
}
|
||||
|
||||
// Send 发送消息到客户端（线程安全）
|
||||
func (c *Connection) Send(data []byte) error {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return errors.New("连接已关闭，无法发送消息")
|
||||
default:
|
||||
// 加锁防止并发写
|
||||
c.writeMutex.Lock()
|
||||
defer c.writeMutex.Unlock()
|
||||
|
||||
// 设置写超时
|
||||
c.conn.SetWriteDeadline(time.Now().Add(c.manager.config.WriteTimeout))
|
||||
|
||||
// 发送消息（使用连接的默认类型，支持动态调整）
|
||||
err := c.conn.WriteMessage(c.manager.config.MsgType, data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("发送消息失败：%w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Close 关闭连接（优雅清理）
|
||||
func (c *Connection) Close(err error) {
|
||||
// 防止重复关闭
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// 取消上下文（终止所有协程）
|
||||
c.cancel()
|
||||
|
||||
// 关闭底层连接（友好关闭）
|
||||
_ = c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, err.Error()))
|
||||
_ = c.conn.Close()
|
||||
|
||||
// 从管理器移除
|
||||
c.manager.mutex.Lock()
|
||||
delete(c.manager.connections, c.connID)
|
||||
c.manager.mutex.Unlock()
|
||||
|
||||
// 触发断开回调
|
||||
c.manager.OnDisconnect(c.connID, err)
|
||||
|
||||
log.Printf("连接[%s]已关闭，当前在线数：%d，原因：%v", c.connID, c.manager.GetOnlineCount(), err)
|
||||
}
|
||||
|
||||
// GetOnlineCount 获取在线连接数
|
||||
func (m *Manager) GetOnlineCount() int {
|
||||
m.mutex.RLock()
|
||||
defer m.mutex.RUnlock()
|
||||
return len(m.connections)
|
||||
}
|
||||
|
||||
// Broadcast 广播消息到所有在线连接
|
||||
func (m *Manager) Broadcast(data []byte) error {
|
||||
m.mutex.RLock()
|
||||
defer m.mutex.RUnlock()
|
||||
|
||||
if len(m.connections) == 0 {
|
||||
return errors.New("无在线连接")
|
||||
}
|
||||
|
||||
// 并发发送（非阻塞）
|
||||
var wg sync.WaitGroup
|
||||
var errMsg string
|
||||
|
||||
for _, conn := range m.connections {
|
||||
wg.Add(1)
|
||||
go func(c *Connection) {
|
||||
defer wg.Done()
|
||||
if err := c.Send(data); err != nil {
|
||||
errMsg += fmt.Sprintf("连接[%s]广播失败：%v；", c.connID, err)
|
||||
}
|
||||
}(conn)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
if errMsg != "" {
|
||||
return errors.New(errMsg)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SendToConn 定向发送消息到指定连接
|
||||
func (m *Manager) SendToConn(connID string, data []byte) error {
|
||||
m.mutex.RLock()
|
||||
conn, exists := m.connections[connID]
|
||||
m.mutex.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return fmt.Errorf("连接[%s]不存在", connID)
|
||||
}
|
||||
|
||||
return conn.Send(data)
|
||||
}
|
||||
|
||||
func (m *Manager) GetAllConn() map[string]*Connection {
|
||||
m.mutex.RLock()
|
||||
defer m.mutex.RUnlock()
|
||||
// 返回副本，防止外部修改
|
||||
connCopy := make(map[string]*Connection, len(m.connections))
|
||||
for k, v := range m.connections {
|
||||
connCopy[k] = v
|
||||
}
|
||||
return connCopy
|
||||
}
|
||||
|
||||
func (m *Manager) GetConn(connID string) *Connection {
|
||||
m.mutex.RLock()
|
||||
defer m.mutex.RUnlock()
|
||||
return m.connections[connID]
|
||||
}
|
||||
|
||||
// CloseAll 关闭所有连接
|
||||
func (m *Manager) CloseAll() {
|
||||
m.mutex.RLock()
|
||||
connIDs := make([]string, 0, len(m.connections))
|
||||
for connID := range m.connections {
|
||||
connIDs = append(connIDs, connID)
|
||||
}
|
||||
m.mutex.RUnlock()
|
||||
|
||||
for _, connID := range connIDs {
|
||||
m.mutex.RLock()
|
||||
conn := m.connections[connID]
|
||||
m.mutex.RUnlock()
|
||||
if conn != nil {
|
||||
conn.Close(errors.New("服务端主动关闭所有连接"))
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user