From a66a38e07419b04ffe3b7898a888d3990552de44 Mon Sep 17 00:00:00 2001 From: qhd <1766646056@qq.com> Date: Fri, 21 Aug 2026 09:26:35 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=E9=80=9A=E7=94=A8?= =?UTF-8?q?=E5=B7=A5=E5=85=B7=E6=A1=86=E6=9E=B6=E4=B8=8EWebSocket=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 新增 tools 包：统一工具定义、注册表与 Server 接口，对齐 MCP 规范 * 新增 websocket 包：泛化连接管理、心跳、并发写锁与优雅关闭 * 新增参数读取工具函数，避免类型断言静默失败 * 新增 OSS 路径识别与 JSON 扁平映射还原工具 * 修复租户 SQL 条件插入位置，正确处理 GROUP BY 与 ORDER BY 同时出现的场景 --- db/gfdb/gfdb.go | 10 +- go.mod | 8 +- go.sum | 2 + tools/args.go | 47 ++++++ tools/doc.go | 20 +++ tools/result.go | 29 ++++ tools/server.go | 47 ++++++ tools/tool.go | 26 ++++ utils/json_flatten.go | 43 ++++++ utils/oss.go | 41 +++++ utils/utils.go | 21 --- websocket/connection.go | 75 +++++++++ websocket/const.go | 12 ++ websocket/option.go | 72 +++++++++ websocket/server.go | 326 ++++++++++++++++++++++++++++++++++++++++ 15 files changed, 752 insertions(+), 27 deletions(-) create mode 100644 tools/args.go create mode 100644 tools/doc.go create mode 100644 tools/result.go create mode 100644 tools/server.go create mode 100644 tools/tool.go create mode 100644 utils/json_flatten.go create mode 100644 utils/oss.go create mode 100644 websocket/connection.go create mode 100644 websocket/const.go create mode 100644 websocket/option.go create mode 100644 websocket/server.go diff --git a/db/gfdb/gfdb.go b/db/gfdb/gfdb.go index e9782e5..f97556a 100644 --- a/db/gfdb/gfdb.go +++ b/db/gfdb/gfdb.go @@ -318,14 +318,16 @@ func selectHook(ctx context.Context, in *gdb.HookSelectInput) (result gdb.Result return nil, err } tenantId = user.TenantId - // 【关键修复】找到 SQL 中第一个出现的 ORDER BY / GROUP BY / LIMIT 等关键字位置 + // 【关键修复】找到 SQL 中最靠前出现的关键字位置：SQL 子句顺序固定为 + // GROUP BY → HAVING → ORDER BY → LIMIT，必须取最早出现的位置，而非关键字列表里先命中的那个。 + // 否则同时含 GROUP BY 与 ORDER BY 的查询会命中靠后的 ORDER BY，把 tenant_id 条件拼进 GROUP BY + // （如 GROUP BY DATE(created_at) AND tenant_id = 94），导致 pq: argument of AND must be type boolean。 sql := in.Sql insertPos := len(sql) - keywords := []string{" ORDER BY ", " GROUP BY ", " HAVING ", " LIMIT ", " FOR UPDATE"} + keywords := []string{" GROUP BY ", " HAVING ", " ORDER BY ", " LIMIT ", " FOR UPDATE"} for _, kw := range keywords { - if idx := gstr.PosI(sql, kw); idx != -1 { + if idx := gstr.PosI(sql, kw); idx != -1 && idx < insertPos { insertPos = idx - break } } diff --git a/go.mod b/go.mod index 2f93a38..7e3ac1e 100644 --- a/go.mod +++ b/go.mod @@ -9,12 +9,15 @@ require ( github.com/gogf/gf/contrib/registry/consul/v2 v2.9.5 github.com/gogf/gf/contrib/trace/otlphttp/v2 v2.9.5 github.com/gogf/gf/v2 v2.9.5 + github.com/google/uuid v1.6.0 + github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 github.com/hashicorp/consul/api v1.26.1 github.com/meilisearch/meilisearch-go v0.36.1 github.com/olivere/elastic/v7 v7.0.32 github.com/r3labs/diff/v2 v2.15.1 github.com/rpcxio/rpcx-consul v0.1.1 github.com/smallnest/rpcx v1.9.1 + github.com/tidwall/sjson v1.2.5 github.com/tiger1103/gfast-token v1.0.10 go.mongodb.org/mongo-driver/v2 v2.4.0 go.opentelemetry.io/otel v1.38.0 @@ -70,8 +73,6 @@ require ( github.com/google/flatbuffers v1.12.1 // indirect github.com/google/gnostic-models v0.7.0 // indirect github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 // indirect - github.com/google/uuid v1.6.0 // indirect - github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grandcat/zeroconf v1.0.0 // indirect github.com/grokify/html-strip-tags-go v0.1.0 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 // indirect @@ -127,6 +128,9 @@ require ( github.com/spf13/pflag v1.0.9 // indirect github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect github.com/tinylib/msgp v1.3.0 // indirect github.com/tjfoc/gmsm v1.4.1 // indirect github.com/tklauser/go-sysconf v0.3.6 // indirect diff --git a/go.sum b/go.sum index e08e9ad..e3bebed 100644 --- a/go.sum +++ b/go.sum @@ -608,10 +608,12 @@ github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 h1:89CEmDvlq/F7S github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161/go.mod h1:wM7WEvslTq+iOEAMDLSzhVuOt5BRZ05WirO+b09GHQU= github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b h1:fj5tQ8acgNUr6O8LEplsxDhUIe2573iLkJc+PqnzZTI= github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b/go.mod h1:5XA7W9S6mni3h5uvOC75dA3m9CCCaS83lltmc0ukdi4= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= diff --git a/tools/args.go b/tools/args.go new file mode 100644 index 0000000..ac5fec4 --- /dev/null +++ b/tools/args.go @@ -0,0 +1,47 @@ +package tools + +import "github.com/gogf/gf/v2/util/gconv" + +// 标准入参读取。LLM/外部传入的 map[string]any 中数值可能是 float64、int 或字符串， +// 直接类型断言极易静默失败，统一走 gconv 强转。 + +// HasArg 判断参数是否存在且非 nil +func HasArg(args map[string]any, key string) bool { + v, ok := args[key] + return ok && v != nil +} + +// ArgString 读取字符串参数 +func ArgString(args map[string]any, key string) string { + return gconv.String(args[key]) +} + +// ArgInt 读取整数参数（自动兼容 float64/int/字符串） +func ArgInt(args map[string]any, key string) int { + return gconv.Int(args[key]) +} + +// ArgFloat64 读取浮点参数 +func ArgFloat64(args map[string]any, key string) float64 { + return gconv.Float64(args[key]) +} + +// ArgBool 读取布尔参数 +func ArgBool(args map[string]any, key string) bool { + return gconv.Bool(args[key]) +} + +// ArgSlice 读取切片参数 +func ArgSlice(args map[string]any, key string) []any { + return gconv.SliceAny(args[key]) +} + +// ArgStrings 读取字符串切片参数 +func ArgStrings(args map[string]any, key string) []string { + return gconv.Strings(args[key]) +} + +// ArgMap 读取对象参数 +func ArgMap(args map[string]any, key string) map[string]any { + return gconv.Map(args[key]) +} diff --git a/tools/doc.go b/tools/doc.go new file mode 100644 index 0000000..05e4ac8 --- /dev/null +++ b/tools/doc.go @@ -0,0 +1,20 @@ +// Package tools 提供共享的模型工具框架：工具定义、注册表与结构化返回。 +// +// 数据模型对齐 MCP 的 Tool 定义（name + description + inputSchema）。 +// 工具是**给模型 function calling 用**的通用能力：模型在推理中动态决定是否调用， +// 消费方依赖 Server 接口（List/Call）而非注册表本身，未来可无缝替换为远程 MCP Server 客户端。 +// +// # 分层 +// +// - tools（本包）：Tool 定义 + 全局注册表 Register + Server 接口，纯框架，不依赖任何业务包 +// - 业务侧：各服务在自身代码中定义**通用**工具（含执行实现 Func），通过 init() 注册进共享注册表； +// 工具的"使用"（如 ReAct 执行循环）也由业务服务基于 Server.Call 自行编排 +// +// 非通用、绑定到固定业务场景的处理逻辑（如工作流节点的前后置钩子）**不属于工具**， +// 应由各业务在自己的编排层维护，不进入本注册表。 +// +// # 消费路径 +// +// 注册表内工具由各业务服务的模型调用方消费：模型 function calling 模式下 +// 经 Server.List 获取全部工具定义交给模型，按返回的 ToolCall 经 Server.Call 执行。 +package tools diff --git a/tools/result.go b/tools/result.go new file mode 100644 index 0000000..748e932 --- /dev/null +++ b/tools/result.go @@ -0,0 +1,29 @@ +package tools + +import "fmt" + +// ToolResult 工具统一返回结构。Code=0 表示成功，非 0 为业务/内部错误码； +// 调用方（工作流前置/后置钩子、模型 function calling）统一按 Code 判断结果。 +type ToolResult struct { + Code int `json:"code"` // 0=成功，非 0=失败 + Message string `json:"message"` // 成功说明或错误信息 + Data any `json:"data,omitempty"` +} + +// 标准错误码。业务工具可自定义扩展 >500 的错误码。 +const ( + CodeOK = 0 // 成功 + CodeInvalidArgs = 400 // 入参缺失或格式错误 + CodeNotFound = 404 // 工具/数据不存在 + CodeInternal = 500 // 内部执行错误 +) + +// OK 构造成功结果 +func OK(data any) ToolResult { + return ToolResult{Code: CodeOK, Data: data} +} + +// Fail 构造失败结果，message 支持 fmt 格式化 +func Fail(code int, format string, args ...any) ToolResult { + return ToolResult{Code: code, Message: fmt.Sprintf(format, args...)} +} diff --git a/tools/server.go b/tools/server.go new file mode 100644 index 0000000..3dde179 --- /dev/null +++ b/tools/server.go @@ -0,0 +1,47 @@ +package tools + +import ( + "context" + "sort" +) + +// Server 提供工具的发现与执行能力，对应 MCP 中 Server 侧的 tools/list 与 tools/call。 +// 消费方（工作流前置/后置钩子、/tool/list 接口、模型 function calling）依赖该接口而非注册表本身， +// 未来可无缝替换为远程 MCP Server 客户端实现。 +type Server interface { + // List 返回全部已注册工具定义，按名称排序 + List(ctx context.Context) ([]*Tool, error) + // Call 按名称调用工具。业务失败通过返回结果的 Code 表达， + // error 仅用于基础设施异常（如 ctx 取消）。 + Call(ctx context.Context, name string, args map[string]any) (ToolResult, error) +} + +// localServer 基于本地注册表的 Server 实现 +type localServer struct{} + +// Default 默认本地 Server。未来接远程 MCP 时，替换该变量即可，业务侧零改动。 +var Default Server = localServer{} + +func (localServer) List(ctx context.Context) ([]*Tool, error) { + list := make([]*Tool, 0, len(registry)) + for _, t := range registry { + list = append(list, t) + } + sort.Slice(list, func(i, j int) bool { return list[i].Name < list[j].Name }) + return list, nil +} + +func (localServer) Call(ctx context.Context, name string, args map[string]any) (ToolResult, error) { + if err := ctx.Err(); err != nil { + return ToolResult{}, err + } + tool := registry[name] + if tool == nil || tool.Func == nil { + return Fail(CodeNotFound, "工具[%s]不存在或未实现", name), nil + } + res, err := tool.Func(ctx, args) + if err != nil { + return Fail(CodeInternal, "工具[%s]执行失败: %v", name, err), nil + } + return res, nil +} diff --git a/tools/tool.go b/tools/tool.go new file mode 100644 index 0000000..d2da4a4 --- /dev/null +++ b/tools/tool.go @@ -0,0 +1,26 @@ +package tools + +import "context" + +// Tool 对应 MCP 的 Tool 数据结构：name + description + inputSchema。 +// 模型 function calling 通过 InputSchema 感知入参，推理中按需调用 Func。 +// 若未来接远程 MCP Server，Func 由客户端统一实现，本结构不变。 +type Tool struct { + Name string // 工具唯一标识 + Description string // 用途说明（模型必读，用于决定何时调用） + Parameters map[string]any // 入参 JSON Schema + Func func(ctx context.Context, args map[string]any) (ToolResult, error) +} + +// registry 工具注册表 +var registry = make(map[string]*Tool) + +// Register 注册工具，同名覆盖 +func Register(list ...*Tool) { + for _, t := range list { + if t == nil || t.Name == "" { + continue + } + registry[t.Name] = t + } +} diff --git a/utils/json_flatten.go b/utils/json_flatten.go new file mode 100644 index 0000000..29af6c5 --- /dev/null +++ b/utils/json_flatten.go @@ -0,0 +1,43 @@ +package utils + +import ( + "encoding/json" + "fmt" + + "github.com/tidwall/sjson" +) + +// IsFlatMap 递归判断 map 是否扁平化 +func IsFlatMap(m map[string]interface{}) bool { + for _, v := range m { + switch val := v.(type) { + case map[string]interface{}: + return false + case []interface{}: + for _, item := range val { + if _, ok := item.(map[string]interface{}); ok { + return false + } + } + } + } + return true +} + +// UnFlatBySjson 将扁平路径映射还原为嵌套 JSON +func UnFlatBySjson(flatMap map[string]interface{}) (map[string]interface{}, error) { + raw := "{}" + for path, val := range flatMap { + var err error + raw, err = sjson.Set(raw, path, val) + if err != nil { + return nil, fmt.Errorf("sjson set path %s failed: %w", path, err) + } + } + + var result map[string]interface{} + if err := json.Unmarshal([]byte(raw), &result); err != nil { + return nil, fmt.Errorf("parse final json failed: %w", err) + } + return result, nil +} diff --git a/utils/oss.go b/utils/oss.go new file mode 100644 index 0000000..7ee0921 --- /dev/null +++ b/utils/oss.go @@ -0,0 +1,41 @@ +package utils + +import ( + "context" + "fmt" + "regexp" + + "github.com/gogf/gf/v2/frame/g" +) + +// ossObjectPathPattern 匹配 MinIO 上传生成的对象路径（不带 http 前缀的相对路径）： +// /YYYY-MM-DD/32位uuid.扩展名，如 /2026-08-19/1e9d9e48-3f6b-4a2c-8d5e-1f2a3b4c.png +var ossObjectPathPattern = regexp.MustCompile(`^/\d{4}-\d{2}-\d{2}/[0-9a-fA-F-]{32}\.[a-zA-Z0-9]{1,10}$`) + +// IsOSSPath 判断字符串是否为 MinIO 对象路径（无 http(s) 前缀）。 +// 模型网关把结果转存 OSS 后返回该裸路径，消费方据此识别"已是文件路径"而不再重复上传。 +// 对象命名规则见 oss/minio 的 ensureBucketAndObjectName，格式变化只需改这一处。 +func IsOSSPath(s string) bool { + return ossObjectPathPattern.MatchString(s) +} + +// GetFileAddressPrefix 拼接图片前缀地址 +func GetFileAddressPrefix(ctx context.Context) (imageUrl string, err error) { + // 拼接图片前缀地址 + bucketName, err := GetBucketName(ctx) + if err != nil { + return + } + imageUrl = fmt.Sprintf("%s/%s", g.Cfg().MustGet(ctx, "filePrefix").String(), bucketName) + return +} + +// GetBucketName 获取bucket名称 +func GetBucketName(ctx context.Context) (bucketName string, err error) { + user, err := GetUserInfo(ctx) + if err != nil { + return + } + bucketName = fmt.Sprintf("tenantid-%d", user.TenantId) + return +} diff --git a/utils/utils.go b/utils/utils.go index ba23cba..5c34214 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -389,27 +389,6 @@ func intPow10(n int) int { return result } -// GetFileAddressPrefix 拼接图片前缀地址 -func GetFileAddressPrefix(ctx context.Context) (imageUrl string, err error) { - // 拼接图片前缀地址 - bucketName, err := GetBucketName(ctx) - if err != nil { - return - } - imageUrl = fmt.Sprintf("%s/%s", g.Cfg().MustGet(ctx, "filePrefix").String(), bucketName) - return -} - -// GetBucketName 获取bucket名称 -func GetBucketName(ctx context.Context) (bucketName string, err error) { - user, err := GetUserInfo(ctx) - if err != nil { - return - } - bucketName = fmt.Sprintf("tenantid-%d", user.TenantId) - return -} - // Lock 分布式锁 func Lock(ctx context.Context, key string, expireSeconds int64, fn func(ctx context.Context) error) (success bool, err error) { limit := 3 diff --git a/websocket/connection.go b/websocket/connection.go new file mode 100644 index 0000000..fa2bb51 --- /dev/null +++ b/websocket/connection.go @@ -0,0 +1,75 @@ +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 +} diff --git a/websocket/const.go b/websocket/const.go new file mode 100644 index 0000000..d69b9e5 --- /dev/null +++ b/websocket/const.go @@ -0,0 +1,12 @@ +package websocket + +import "time" + +const ( + DefaultReadTimeout = 90 * time.Second + DefaultWriteTimeout = 10 * time.Second + DefaultHeartbeatInterval = 30 * time.Second + DefaultWorkerPoolSize = 50 + DefaultMaxConnections = 2000 + DefaultConnKeyPrefix = "ws:" +) diff --git a/websocket/option.go b/websocket/option.go new file mode 100644 index 0000000..5871cae --- /dev/null +++ b/websocket/option.go @@ -0,0 +1,72 @@ +package websocket + +import ( + "context" + netHttp "net/http" + "time" +) + +// MessageHandler 业务消息处理函数 +type MessageHandler func(ctx context.Context, conn *WsConnection, payload interface{}) + +// WsMessage 通用入站消息 +type WsMessage struct { + Type string `json:"type"` + Payload interface{} `json:"payload,omitempty"` +} + +// WsPushMsg 通用出站推送消息（业务可自行扩展字段） +type WsPushMsg struct { + Type string `json:"type"` + Message string `json:"message,omitempty"` + Data interface{} `json:"data,omitempty"` + Error string `json:"error,omitempty"` +} + +// ServerOptions 服务配置 +type ServerOptions struct { + readTimeout time.Duration + writeTimeout time.Duration + heartbeatInterval time.Duration + workerPoolSize int + maxConnections int + connKeyPrefix string + checkOrigin func(r *netHttp.Request) bool +} + +// ServerOption 配置函数 +type ServerOption func(*ServerOptions) + +func WithReadTimeout(d time.Duration) ServerOption { + return func(o *ServerOptions) { o.readTimeout = d } +} +func WithWriteTimeout(d time.Duration) ServerOption { + return func(o *ServerOptions) { o.writeTimeout = d } +} +func WithHeartbeatInterval(d time.Duration) ServerOption { + return func(o *ServerOptions) { o.heartbeatInterval = d } +} +func WithWorkerPoolSize(n int) ServerOption { + return func(o *ServerOptions) { o.workerPoolSize = n } +} +func WithMaxConnections(n int) ServerOption { + return func(o *ServerOptions) { o.maxConnections = n } +} +func WithConnKeyPrefix(p string) ServerOption { + return func(o *ServerOptions) { o.connKeyPrefix = p } +} +func WithCheckOrigin(fn func(r *netHttp.Request) bool) ServerOption { + return func(o *ServerOptions) { o.checkOrigin = fn } +} + +func defaultOptions() ServerOptions { + return ServerOptions{ + readTimeout: DefaultReadTimeout, + writeTimeout: DefaultWriteTimeout, + heartbeatInterval: DefaultHeartbeatInterval, + workerPoolSize: DefaultWorkerPoolSize, + maxConnections: DefaultMaxConnections, + connKeyPrefix: DefaultConnKeyPrefix, + checkOrigin: func(r *netHttp.Request) bool { return true }, + } +} diff --git a/websocket/server.go b/websocket/server.go new file mode 100644 index 0000000..07819ec --- /dev/null +++ b/websocket/server.go @@ -0,0 +1,326 @@ +package websocket + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/gogf/gf/v2/container/gmap" + "github.com/gogf/gf/v2/encoding/gjson" + "github.com/gogf/gf/v2/frame/g" + "github.com/gogf/gf/v2/net/ghttp" + "github.com/gogf/gf/v2/os/glog" + "github.com/gogf/gf/v2/os/grpool" + "github.com/google/uuid" + "github.com/gorilla/websocket" +) + +// WsServer 泛化 WebSocket 服务器 +type WsServer struct { + connections *gmap.StrAnyMap + upgrader websocket.Upgrader + workerPool *grpool.Pool + handlers map[string]MessageHandler + handlerMu sync.RWMutex + opts ServerOptions + + svcClosed int32 + closeOnce sync.Once +} + +// NewWsServer 创建泛化 WebSocket 服务器 +func NewWsServer(opts ...ServerOption) *WsServer { + o := defaultOptions() + for _, opt := range opts { + opt(&o) + } + + return &WsServer{ + connections: gmap.NewStrAnyMap(true), + upgrader: websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + CheckOrigin: o.checkOrigin, + }, + workerPool: grpool.New(o.workerPoolSize), + handlers: make(map[string]MessageHandler), + opts: o, + svcClosed: 0, + } +} + +// OnMessage 注册业务消息处理器 +func (s *WsServer) OnMessage(msgType string, handler MessageHandler) { + s.handlerMu.Lock() + defer s.handlerMu.Unlock() + s.handlers[msgType] = handler +} + +// Upgrade 将 HTTP 连接升级为 WebSocket 并注册到连接池 +func (s *WsServer) Upgrade(ctx context.Context, r *ghttp.Request, sessionId string) (*WsConnection, error) { + if g.IsEmpty(sessionId) { + sessionId = uuid.NewString() + } + if atomic.LoadInt32(&s.svcClosed) == 1 { + return nil, errors.New("websocket server is closed") + } + if s.connections.Size() >= s.opts.maxConnections { + return nil, errors.New("too many online websocket connections") + } + + wsConn, err := s.upgrader.Upgrade(r.Response.Writer, r.Request, nil) + if err != nil { + return nil, fmt.Errorf("upgrade failed: %w", err) + } + + headers := make(map[string]string) + for k, v := range r.Request.Header { + if len(v) > 0 { + headers[k] = v[0] + } + } + + key := s.opts.connKeyPrefix + sessionId + + // 踢下线旧连接 + s.kickOld(key) + + baseCtx := context.WithoutCancel(ctx) + closeCtx, closeCancel := context.WithCancel(baseCtx) + + wc := &WsConnection{ + SessionId: sessionId, + Conn: wsConn, + Headers: headers, + closeCancel: closeCancel, + closed: 0, + } + + s.connections.Set(key, wc) + + // 连接成功回执 + _ = s.writeJSON(closeCtx, wc, &WsPushMsg{Type: "ack", Message: "WebSocket连接成功", Data: map[string]any{ + "sessionId": sessionId, + }}) + + // Pong 心跳回调，重置读超时 + wsConn.SetPongHandler(func(string) error { + _ = wsConn.SetReadDeadline(time.Now().Add(s.opts.readTimeout)) + return nil + }) + + go s.handleConnection(closeCtx, key, wc) + return wc, nil +} + +// PushToSession 向指定会话推送消息 +func (s *WsServer) PushToSession(ctx context.Context, sessionId string, msg *WsPushMsg) { + key := s.opts.connKeyPrefix + sessionId + val := s.connections.Get(key) + if val == nil { + return + } + wc, ok := val.(*WsConnection) + if !ok || wc.IsClosed() { + return + } + _ = s.writeJSON(ctx, wc, msg) +} + +// GetOnlineSessions 获取在线会话列表 +func (s *WsServer) GetOnlineSessions() []string { + var sessions []string + prefixLen := len(s.opts.connKeyPrefix) + s.connections.Iterator(func(key string, _ interface{}) bool { + if len(key) > prefixLen { + sessions = append(sessions, key[prefixLen:]) + } + return true + }) + return sessions +} + +// Close 全局优雅关闭 +func (s *WsServer) Close() { + s.closeOnce.Do(func() { + atomic.StoreInt32(&s.svcClosed, 1) + s.workerPool.Close() + + s.connections.LockFunc(func(m map[string]interface{}) { + for _, val := range m { + wc, ok := val.(*WsConnection) + if !ok { + continue + } + if atomic.CompareAndSwapInt32(&wc.closed, 0, 1) { + if wc.closeCancel != nil { + wc.closeCancel() + } + _ = wc.Conn.Close() + } + } + }) + s.connections.Clear() + }) +} + +// ====================== 内部方法 ====================== + +// kickOld 踢掉同session旧连接，不再主动remove，由旧连接defer清理 +func (s *WsServer) kickOld(key string) { + val := s.connections.Get(key) + if val == nil { + return + } + old, ok := val.(*WsConnection) + if !ok { + return + } + if atomic.CompareAndSwapInt32(&old.closed, 0, 1) { + if old.closeCancel != nil { + old.closeCancel() + } + _ = old.Conn.Close() + } +} + +// heartbeatLoop 心跳发送协程，入参改为 *WsConnection，复用写锁 +func (s *WsServer) heartbeatLoop(ctx context.Context, wc *WsConnection, done <-chan struct{}) { + ticker := time.NewTicker(s.opts.heartbeatInterval) + defer ticker.Stop() + conn := wc.Conn + + for { + select { + case <-ticker.C: + wc.writeMu.Lock() + _ = conn.SetWriteDeadline(time.Now().Add(s.opts.writeTimeout)) + err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(s.opts.writeTimeout)) + wc.writeMu.Unlock() + if err != nil { + glog.Debugf(ctx, "heartbeat ping failed: %v", err) + return + } + case <-done: + return + case <-ctx.Done(): + return + } + } +} + +func (s *WsServer) handleConnection(ctx context.Context, key string, wc *WsConnection) { + conn := wc.Conn + + defer func() { + if atomic.CompareAndSwapInt32(&wc.closed, 0, 1) { + if wc.closeCancel != nil { + wc.closeCancel() + } + _ = conn.Close() + } + // 关键修复：只删除自身实例，防止旧连接误删新连接 + s.connections.LockFunc(func(m map[string]interface{}) { + if v, exist := m[key]; exist && v == wc { + delete(m, key) + } + }) + }() + + done := make(chan struct{}) + defer close(done) + go s.heartbeatLoop(ctx, wc, done) + + _ = conn.SetReadDeadline(time.Now().Add(s.opts.readTimeout)) + + for { + select { + case <-ctx.Done(): + return + default: + } + + msgType, data, err := conn.ReadMessage() + if err != nil { + // 正常关闭不打error日志 + if !websocket.IsUnexpectedCloseError(err, + websocket.CloseNormalClosure, + websocket.CloseGoingAway, + websocket.CloseNoStatusReceived, + ) { + glog.Debugf(ctx, "normal close: %s, err: %v", key, err) + } else { + glog.Infof(ctx, "unexpected close: %s, err: %v", key, err) + } + break + } + + _ = conn.SetReadDeadline(time.Now().Add(s.opts.readTimeout)) + + switch msgType { + case websocket.PingMessage: + wc.writeMu.Lock() + _ = conn.SetWriteDeadline(time.Now().Add(s.opts.writeTimeout)) + _ = conn.WriteMessage(websocket.PongMessage, nil) + wc.writeMu.Unlock() + continue + case websocket.CloseMessage: + return + case websocket.BinaryMessage, websocket.TextMessage: + default: + continue + } + + if len(data) == 0 { + continue + } + + var msg WsMessage + if err := gjson.Unmarshal(data, &msg); err != nil { + _ = s.writeJSON(ctx, wc, &WsPushMsg{Type: "error", Message: "消息格式错误", Error: err.Error()}) + continue + } + + s.handlerMu.RLock() + handler, exists := s.handlers[msg.Type] + s.handlerMu.RUnlock() + + if !exists { + _ = s.writeJSON(ctx, wc, &WsPushMsg{Type: "error", Message: fmt.Sprintf("未知消息类型: %s", msg.Type)}) + continue + } + + // 【重要修复】投递到workerPool，避免业务阻塞读循环 + taskCtx := ctx + payload := msg.Payload + if err := s.workerPool.Add(taskCtx, func(ctx context.Context) { + handler(ctx, wc, payload) + }); err != nil { + _ = s.writeJSON(ctx, wc, &WsPushMsg{ + Type: "error", + Message: "服务繁忙，任务队列已满", + }) + } + } +} + +// writeJSON 统一写入消息，入参改为 *WsConnection，带并发写锁 +func (s *WsServer) writeJSON(ctx context.Context, wc *WsConnection, data interface{}) error { + wc.writeMu.Lock() + defer wc.writeMu.Unlock() + + jsonBytes, err := gjson.Encode(data) + if err != nil { + glog.Errorf(ctx, "json encode failed: %v", err) + return err + } + _ = wc.Conn.SetWriteDeadline(time.Now().Add(s.opts.writeTimeout)) + if err = wc.Conn.WriteMessage(websocket.TextMessage, jsonBytes); err != nil { + glog.Debugf(ctx, "websocket write failed: %v", err) + return err + } + return nil +}