19904408334
a66a38e074
* 新增 tools 包：统一工具定义、注册表与 Server 接口，对齐 MCP 规范 * 新增 websocket 包：泛化连接管理、心跳、并发写锁与优雅关闭 * 新增参数读取工具函数，避免类型断言静默失败 * 新增 OSS 路径识别与 JSON 扁平映射还原工具 * 修复租户 SQL 条件插入位置，正确处理 GROUP BY 与 ORDER BY 同时出现的场景
76 lines
2.0 KiB
Go
76 lines
2.0 KiB
Go
package websocket
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/encoding/gjson"
|
||
"github.com/gogf/gf/v2/os/glog"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
// WsConnection 单个WebSocket连接，Metadata 存放业务自定义数据
|
||
type WsConnection struct {
|
||
SessionId string
|
||
Conn *websocket.Conn
|
||
Headers map[string]string
|
||
Metadata sync.Map // 业务数据：FlowId, execCancel 等
|
||
|
||
writeMu sync.Mutex // 保护 websocket.Conn 并发写（WriteMessage + WriteControl）
|
||
closeCancel context.CancelFunc
|
||
closed int32
|
||
}
|
||
|
||
// SetMeta 设置业务元数据
|
||
func (c *WsConnection) SetMeta(key string, value interface{}) {
|
||
c.Metadata.Store(key, value)
|
||
}
|
||
|
||
// GetMeta 获取业务元数据
|
||
func (c *WsConnection) GetMeta(key string) (interface{}, bool) {
|
||
return c.Metadata.Load(key)
|
||
}
|
||
|
||
// GetMetaT 泛型版 GetMeta，省去外部类型断言
|
||
func GetMetaT[T any](c *WsConnection, key string) (T, bool) {
|
||
val, ok := c.Metadata.Load(key)
|
||
if !ok {
|
||
var zero T
|
||
return zero, false
|
||
}
|
||
t, ok := val.(T)
|
||
return t, ok
|
||
}
|
||
|
||
// IsClosed 连接是否已关闭
|
||
func (c *WsConnection) IsClosed() bool {
|
||
return atomic.LoadInt32(&c.closed) == 1
|
||
}
|
||
|
||
// WriteControl 带写锁保护的 WriteControl，用于心跳 Ping / Pong / Close帧
|
||
func (c *WsConnection) WriteControl(msgType int, data []byte, deadline time.Time) error {
|
||
c.writeMu.Lock()
|
||
defer c.writeMu.Unlock()
|
||
_ = c.Conn.SetWriteDeadline(deadline)
|
||
return c.Conn.WriteControl(msgType, data, deadline)
|
||
}
|
||
|
||
// WriteJSON 业务层外部写入入口，共享 writeMu 与心跳/Pong 互斥
|
||
func (c *WsConnection) WriteJSON(data interface{}) error {
|
||
jsonBytes, err := gjson.Encode(data)
|
||
if err != nil {
|
||
glog.Errorf(context.Background(), "json encode failed: %v", err)
|
||
return err
|
||
}
|
||
c.writeMu.Lock()
|
||
_ = c.Conn.SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||
err = c.Conn.WriteMessage(websocket.TextMessage, jsonBytes)
|
||
c.writeMu.Unlock()
|
||
if err != nil {
|
||
glog.Debugf(context.Background(), "websocket write failed: %v", err)
|
||
}
|
||
return err
|
||
}
|