This commit is contained in:
Cold
2025-12-04 17:39:31 +08:00
committed by 张斌
parent 4300257195
commit 3f4ac12f52
3 changed files with 332 additions and 4 deletions

View File

@@ -4,6 +4,11 @@ import (
"context" "context"
"errors" "errors"
"fmt" "fmt"
"net/http"
"reflect"
"regexp"
"strings"
_ "gitee.com/red-future---jilin-g/common/consul" _ "gitee.com/red-future---jilin-g/common/consul"
"gitee.com/red-future---jilin-g/common/jaeger" "gitee.com/red-future---jilin-g/common/jaeger"
"gitee.com/red-future---jilin-g/common/utils" "gitee.com/red-future---jilin-g/common/utils"
@@ -13,10 +18,6 @@ import (
"github.com/gogf/gf/v2/os/glog" "github.com/gogf/gf/v2/os/glog"
"github.com/gogf/gf/v2/os/gtime" "github.com/gogf/gf/v2/os/gtime"
"github.com/gogf/gf/v2/util/gconv" "github.com/gogf/gf/v2/util/gconv"
"net/http"
"reflect"
"regexp"
"strings"
) )
type ResponseEmpty struct { type ResponseEmpty struct {

165
ragflow/worker_pool.go Normal file
View File

@@ -0,0 +1,165 @@
package ragflow
import (
"context"
"gitee.com/red-future---jilin-g/common/redis"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/os/glog"
"github.com/gogf/gf/v2/os/grpool"
)
// WorkerPool RAGFlow 请求处理协程池
type WorkerPool struct {
pool *grpool.Pool
size int
}
// NewWorkerPool 创建协程池
// 参数:
// - size: 协程池大小,建议设置为 CPU 核心数的 2-4 倍
//
// 返回:
// - *WorkerPool: 协程池实例
// - error: 创建失败时返回错误
func NewWorkerPool(size int) (*WorkerPool, error) {
if size <= 0 {
return nil, gerror.New("协程池大小必须大于0")
}
pool := grpool.New(size)
return &WorkerPool{
pool: pool,
size: size,
}, nil
}
// Submit 提交任务到协程池
// 参数:
// - ctx: 上下文
// - task: 要执行的任务函数
//
// 返回error 提交失败时返回错误
func (w *WorkerPool) Submit(ctx context.Context, task func(ctx context.Context)) error {
return w.pool.Add(ctx, func(ctx context.Context) {
defer func() {
if r := recover(); r != nil {
glog.Errorf(ctx, "协程池任务执行 panic: %v", r)
}
}()
task(ctx)
})
}
// Size 获取协程池大小
func (w *WorkerPool) Size() int {
return w.size
}
// Jobs 获取当前等待执行的任务数量
func (w *WorkerPool) Jobs() int {
return w.pool.Jobs()
}
// Close 关闭协程池
func (w *WorkerPool) Close() {
w.pool.Close()
}
// WorkerStats 协程池统计信息
type WorkerStats struct {
PoolSize int // 协程池大小
Jobs int // 等待执行的任务数
}
// Stats 获取协程池统计信息
func (w *WorkerPool) Stats() WorkerStats {
return WorkerStats{
PoolSize: w.size,
Jobs: w.pool.Jobs(),
}
}
// PrintStats 打印协程池统计信息
func (w *WorkerPool) PrintStats(ctx context.Context) {
stats := w.Stats()
glog.Infof(ctx, "协程池统计 - 池大小: %d, 等待任务: %d", stats.PoolSize, stats.Jobs)
}
// QueueProcessor 队列处理器,从 Redis 队列中取出任务并提交到协程池
type QueueProcessor struct {
pool *WorkerPool
queueKey string
timeout int
stopChan chan struct{}
handleFunc func(ctx context.Context, message string) error
}
// NewQueueProcessor 创建队列处理器
// 参数:
// - pool: 协程池
// - queueKey: Redis 队列键名
// - timeout: 从队列取消息的超时时间(秒)
// - handleFunc: 消息处理函数
func NewQueueProcessor(pool *WorkerPool, queueKey string, timeout int, handleFunc func(ctx context.Context, message string) error) *QueueProcessor {
return &QueueProcessor{
pool: pool,
queueKey: queueKey,
timeout: timeout,
stopChan: make(chan struct{}),
handleFunc: handleFunc,
}
}
// Start 启动队列处理器
// 会阻塞运行,持续从 Redis 队列中取出消息并提交到协程池处理
func (q *QueueProcessor) Start(ctx context.Context) error {
glog.Infof(ctx, "队列处理器启动 - 队列: %s, 超时: %ds", q.queueKey, q.timeout)
for {
select {
case <-q.stopChan:
glog.Info(ctx, "队列处理器收到停止信号")
return nil
default:
// 从 Redis 队列中取出消息
message, err := q.fetchMessage(ctx)
if err != nil {
glog.Errorf(ctx, "从队列取消息失败: %v", err)
continue
}
// 队列为空,继续等待
if message == "" {
continue
}
// 提交到协程池处理
if err := q.submitTask(ctx, message); err != nil {
glog.Errorf(ctx, "提交任务到协程池失败: %v", err)
}
}
}
}
// Stop 停止队列处理器
func (q *QueueProcessor) Stop() {
close(q.stopChan)
}
// fetchMessage 从 Redis 队列中取出消息
func (q *QueueProcessor) fetchMessage(ctx context.Context) (string, error) {
// 调用 Redis 队列的 PopFromQueue 方法从队列中取出消息
return redis.PopFromQueue(ctx, q.queueKey, q.timeout)
}
// submitTask 将消息处理任务提交到协程池
func (q *QueueProcessor) submitTask(ctx context.Context, message string) error {
return q.pool.Submit(ctx, func(ctx context.Context) {
if err := q.handleFunc(ctx, message); err != nil {
glog.Errorf(ctx, "处理消息失败: %v, 消息: %s", err, message)
}
})
}

View File

@@ -1,6 +1,9 @@
package redis package redis
import ( import (
"context"
"time"
"github.com/gogf/gf/v2/database/gredis" "github.com/gogf/gf/v2/database/gredis"
"github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/frame/g"
) )
@@ -10,3 +13,162 @@ var RedisClient *gredis.Redis
func init() { func init() {
RedisClient = g.Redis() RedisClient = g.Redis()
} }
// 队列操作常量
const (
// RAGFlow 请求队列 Key
RAGFlowRequestQueueKey = "ragflow:request:queue"
// 会话最后活跃时间 Key 前缀
SessionLastActiveKeyPrefix = "ragflow:session:"
)
// PushToQueue 将消息推入队列LPUSH从左侧入队
// 用于 Controller 层将 RAGFlow 请求推入队列,由后台 Goroutine 处理
// 参数:
// - queueKey: 队列键名
// - message: 要推入的消息内容
//
// 返回error 推入失败时返回错误
func PushToQueue(ctx context.Context, queueKey string, message interface{}) error {
_, err := RedisClient.Do(ctx, "LPUSH", queueKey, message)
return err
}
// PopFromQueue 从队列中阻塞取出消息BRPOP从右侧出队
// 后台 Goroutine 使用此方法从队列中取出请求进行处理
// 参数:
// - queueKey: 队列键名
// - timeout: 阻塞超时时间0表示永久阻塞
//
// 返回:
// - string: 取出的消息内容,超时或队列为空返回空字符串
// - error: 操作失败时返回错误
func PopFromQueue(ctx context.Context, queueKey string, timeout int) (string, error) {
result, err := RedisClient.Do(ctx, "BRPOP", queueKey, timeout)
if err != nil {
return "", err
}
// BRPOP 返回 [key, value],我们需要取 value
if result == nil {
return "", nil // 超时返回空
}
// GoFrame gredis 返回的是 *gvar.Var 类型
arr := result.Strings()
if len(arr) >= 2 {
return arr[1], nil // arr[0] 是 keyarr[1] 是 value
}
return "", nil
}
// GetQueueLength 获取队列当前长度
// 用于监控队列积压情况
// 参数:
// - queueKey: 队列键名
//
// 返回:
// - int64: 队列中消息数量
// - error: 操作失败时返回错误
func GetQueueLength(ctx context.Context, queueKey string) (int64, error) {
result, err := RedisClient.Do(ctx, "LLEN", queueKey)
if err != nil {
return 0, err
}
return result.Int64(), nil
}
// SetSessionLastActive 设置用户最后活跃时间
// 用于控制是否发送追问:用户回复后更新活跃时间,避免重复追问
// 过期时间2小时超过2小时未活跃的记录会自动删除
// 参数:
// - userId: 用户ID
//
// 返回error 设置失败时返回错误
func SetSessionLastActive(ctx context.Context, userId string) error {
key := SessionLastActiveKeyPrefix + userId + ":last_active"
timestamp := time.Now().Unix()
// 设置过期时间为 2 小时
return RedisClient.SetEX(ctx, key, timestamp, 7200)
}
// GetSessionLastActive 获取用户最后活跃时间
// 参数:
// - userId: 用户ID
//
// 返回:
// - int64: Unix时间戳未找到返回0
// - error: 操作失败时返回错误
func GetSessionLastActive(ctx context.Context, userId string) (int64, error) {
key := SessionLastActiveKeyPrefix + userId + ":last_active"
result, err := RedisClient.Get(ctx, key)
if err != nil {
return 0, err
}
if result.IsNil() {
return 0, nil // 未找到返回 0
}
return result.Int64(), nil
}
// IsUserActive 检查用户是否在指定时间范围内活跃过
// 用于追问逻辑:如果用户最近活跃过,则不发送追问消息
// 参数:
// - userId: 用户ID
// - seconds: 时间范围例如传入300表示检查5分钟内是否活跃
//
// 返回:
// - bool: true表示用户在指定时间内活跃过
// - error: 操作失败时返回错误
func IsUserActive(ctx context.Context, userId string, seconds int64) (bool, error) {
lastActive, err := GetSessionLastActive(ctx, userId)
if err != nil {
return false, err
}
if lastActive == 0 {
return false, nil // 未找到记录,视为不活跃
}
now := time.Now().Unix()
return (now - lastActive) < seconds, nil
}
// SetSessionCache 缓存用户的 RAGFlow Session ID
// 避免每次请求都创建新 Session提高性能
// 过期时间7天超过7天未使用的Session会自动清理
// 参数:
// - userId: 用户ID
// - sessionId: RAGFlow返回的Session ID
//
// 返回error 设置失败时返回错误
func SetSessionCache(ctx context.Context, userId, sessionId string) error {
key := SessionLastActiveKeyPrefix + userId + ":session_id"
return RedisClient.SetEX(ctx, key, sessionId, 7*24*3600)
}
// GetSessionCache 获取缓存的 RAGFlow Session ID
// 如果缓存中存在则直接使用不存在则需要创建新Session
// 参数:
// - userId: 用户ID
//
// 返回:
// - string: Session ID未找到返回空字符串
// - error: 操作失败时返回错误
func GetSessionCache(ctx context.Context, userId string) (string, error) {
key := SessionLastActiveKeyPrefix + userId + ":session_id"
result, err := RedisClient.Get(ctx, key)
if err != nil {
return "", err
}
if result.IsNil() {
return "", nil
}
return result.String(), nil
}