增加es归档 分布式和constants变量
This commit is contained in:
Cold
@@ -2,189 +2,111 @@ package ragflow
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitee.com/red-future---jilin-g/common/redis"
|
||||
"github.com/gogf/gf/v2/os/glog"
|
||||
"github.com/gogf/gf/v2/os/grpool"
|
||||
)
|
||||
|
||||
// 默认协程池大小
|
||||
const defaultPoolSize = 200
|
||||
// 默认批量大小（每次从 Redis 读取并发送的消息数）
|
||||
const defaultBatchSize = 200
|
||||
|
||||
// workerPool 协程池单例（grpool.New 是原型模式，需要变量引用）
|
||||
var workerPool = grpool.New(defaultPoolSize)
|
||||
|
||||
// WorkerPool RAGFlow 请求处理协程池（封装 grpool）
|
||||
type WorkerPool struct {
|
||||
pool *grpool.Pool
|
||||
size int
|
||||
}
|
||||
|
||||
// Pool 协程池单例实例（直接引用使用）
|
||||
var Pool = &WorkerPool{
|
||||
pool: workerPool,
|
||||
size: defaultPoolSize,
|
||||
}
|
||||
|
||||
// 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 Stream 处理器，从 Redis Stream 中取出任务并提交到协程池
|
||||
// QueueProcessor Stream 处理器，批量读取消息并发送到 RAGFlow
|
||||
type QueueProcessor struct {
|
||||
pool *WorkerPool
|
||||
streamKey string // Stream 键名
|
||||
groupName string // 消费者组名称
|
||||
consumerName string // 消费者名称
|
||||
timeout int64 // 阻塞超时时间（毫秒）
|
||||
batchSize int64 // 每次读取的消息数量
|
||||
stopChan chan struct{}
|
||||
streamKey string // Stream 键名
|
||||
groupName string // 消费者组名称
|
||||
consumerName string // 消费者名称
|
||||
timeout int64 // 阻塞超时时间（毫秒）
|
||||
batchSize int64 // 最大并发数（信号量容量）
|
||||
stopChan chan struct{} // 停止信号
|
||||
semaphore chan struct{} // 并发信号量（控制最大并发）
|
||||
handleFunc func(ctx context.Context, message map[string]interface{}) error
|
||||
}
|
||||
|
||||
// NewQueueProcessor 创建 Stream 处理器
|
||||
// 参数：
|
||||
// - pool: 协程池
|
||||
// - streamKey: Redis Stream 键名
|
||||
// - groupName: 消费者组名称
|
||||
// - consumerName: 消费者名称（唯一标识）
|
||||
// - timeout: 从 Stream 取消息的超时时间（毫秒）
|
||||
// - batchSize: 每次读取的消息数量
|
||||
// - handleFunc: 消息处理函数
|
||||
func NewQueueProcessor(pool *WorkerPool, streamKey, groupName, consumerName string, timeout int64, batchSize int64, handleFunc func(ctx context.Context, message map[string]interface{}) error) *QueueProcessor {
|
||||
func NewQueueProcessor(streamKey, groupName, consumerName string, timeout, batchSize int64, handleFunc func(ctx context.Context, message map[string]interface{}) error) *QueueProcessor {
|
||||
return &QueueProcessor{
|
||||
pool: pool,
|
||||
streamKey: streamKey,
|
||||
groupName: groupName,
|
||||
consumerName: consumerName,
|
||||
timeout: timeout,
|
||||
batchSize: batchSize,
|
||||
stopChan: make(chan struct{}),
|
||||
semaphore: make(chan struct{}, batchSize), // 信号量容量 = 最大并发数
|
||||
handleFunc: handleFunc,
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动 Stream 处理器
|
||||
// 会阻塞运行，持续从 Redis Stream 中取出消息并提交到协程池处理
|
||||
// 削峰填谷：每次读取 batchSize 条消息，并发发送，发完立刻读下一批
|
||||
func (q *QueueProcessor) Start(ctx context.Context) error {
|
||||
glog.Infof(ctx, "Stream 处理器启动 - Stream: %s, 消费者组: %s, 消费者: %s, 超时: %dms",
|
||||
q.streamKey, q.groupName, q.consumerName, q.timeout)
|
||||
glog.Infof(ctx, "Stream 处理器启动 - Stream: %s, 消费者组: %s, 消费者: %s, 批量大小: %d",
|
||||
q.streamKey, q.groupName, q.consumerName, q.batchSize)
|
||||
|
||||
// 确保 Consumer Group 存在（重试直到成功）
|
||||
for {
|
||||
if err := redis.CreateConsumerGroup(ctx, q.streamKey, q.groupName); err != nil {
|
||||
// BUSYGROUP 表示已存在，不是错误
|
||||
if strings.Contains(err.Error(), "BUSYGROUP") {
|
||||
glog.Debugf(ctx, "Consumer Group 已存在")
|
||||
break
|
||||
}
|
||||
glog.Warningf(ctx, "创建 Consumer Group 失败: %v，1秒后重试", err)
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
glog.Infof(ctx, "Consumer Group 创建成功")
|
||||
break
|
||||
}
|
||||
|
||||
loopCount := 0
|
||||
for {
|
||||
select {
|
||||
case <-q.stopChan:
|
||||
glog.Info(ctx, "Stream 处理器收到停止信号")
|
||||
return nil
|
||||
default:
|
||||
loopCount++
|
||||
if loopCount%10 == 1 {
|
||||
glog.Debugf(ctx, "[DEBUG] 第 %d 次循环，准备读取消息...", loopCount)
|
||||
}
|
||||
|
||||
// 从 Redis Stream 中读取消息
|
||||
messages, err := q.fetchMessages(ctx)
|
||||
// 1. 从 Redis Stream 读取一批消息
|
||||
messages, err := redis.ReadFromStream(ctx, q.streamKey, q.groupName, q.consumerName, q.batchSize, q.timeout)
|
||||
if err != nil {
|
||||
glog.Errorf(ctx, "从 Stream 读取消息失败: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 没有新消息，继续等待
|
||||
if len(messages) == 0 {
|
||||
if loopCount%10 == 1 {
|
||||
glog.Debugf(ctx, "[DEBUG] 第 %d 次循环，无新消息", loopCount)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
glog.Infof(ctx, "[DEBUG] 收到 %d 条消息", len(messages))
|
||||
glog.Debugf(ctx, "读取 %d 条消息，开始发送", len(messages))
|
||||
|
||||
// 处理每条消息
|
||||
// 2. 用信号量控制并发：获取信号量后发送，完成后释放
|
||||
for _, msg := range messages {
|
||||
glog.Infof(ctx, "[DEBUG] 处理消息 ID: %s, Values: %+v", msg.ID, msg.Values)
|
||||
// 提交到协程池处理
|
||||
if err := q.submitTask(ctx, msg); err != nil {
|
||||
glog.Errorf(ctx, "提交任务到协程池失败: %v, 消息ID: %s", err, msg.ID)
|
||||
}
|
||||
// 获取信号量（阻塞直到有空位）
|
||||
q.semaphore <- struct{}{}
|
||||
go func(m redis.StreamMessage) {
|
||||
defer func() { <-q.semaphore }() // 完成后释放信号量
|
||||
q.processMessage(ctx, m)
|
||||
}(msg)
|
||||
}
|
||||
// 3. 立刻读下一批（不等待，信号量自动控制并发数）
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processMessage 处理单条消息（异步执行）
|
||||
func (q *QueueProcessor) processMessage(ctx context.Context, message redis.StreamMessage) {
|
||||
// 调用处理函数发送到 RAGFlow
|
||||
if err := q.handleFunc(ctx, message.Values); err != nil {
|
||||
glog.Errorf(ctx, "消息处理失败: %v, 消息ID: %s", err, message.ID)
|
||||
}
|
||||
|
||||
// 无论成功失败都 ACK（避免重复消费）
|
||||
if err := redis.AckMessage(ctx, q.streamKey, q.groupName, message.ID); err != nil {
|
||||
glog.Errorf(ctx, "确认消息失败: %v, 消息ID: %s", err, message.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop 停止队列处理器
|
||||
func (q *QueueProcessor) Stop() {
|
||||
close(q.stopChan)
|
||||
}
|
||||
|
||||
// fetchMessages 从 Redis Stream 中读取消息
|
||||
func (q *QueueProcessor) fetchMessages(ctx context.Context) ([]redis.StreamMessage, error) {
|
||||
// 从消费者组读取消息
|
||||
return redis.ReadFromStream(ctx, q.streamKey, q.groupName, q.consumerName, q.batchSize, q.timeout)
|
||||
}
|
||||
|
||||
// submitTask 将消息处理任务提交到协程池
|
||||
func (q *QueueProcessor) submitTask(ctx context.Context, message redis.StreamMessage) error {
|
||||
return q.pool.Submit(ctx, func(ctx context.Context) {
|
||||
// 处理消息
|
||||
if err := q.handleFunc(ctx, message.Values); err != nil {
|
||||
glog.Errorf(ctx, "处理消息失败: %v, 消息ID: %s", err, message.ID)
|
||||
return
|
||||
}
|
||||
|
||||
// 处理成功后确认消息
|
||||
if err := redis.AckMessage(ctx, q.streamKey, q.groupName, message.ID); err != nil {
|
||||
glog.Errorf(ctx, "确认消息失败: %v, 消息ID: %s", err, message.ID)
|
||||
} else {
|
||||
glog.Debugf(ctx, "消息处理完成并已确认: %s", message.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user