|
@@ -6,7 +6,6 @@ import (
|
|
|
"fmt"
|
|
|
"github.com/go-nunu/nunu-layout-advanced/internal/service"
|
|
|
"github.com/go-nunu/nunu-layout-advanced/internal/service/api/waf"
|
|
|
- "github.com/google/uuid"
|
|
|
"github.com/rabbitmq/amqp091-go"
|
|
|
"go.uber.org/zap"
|
|
|
"strconv"
|
|
@@ -14,9 +13,7 @@ import (
|
|
|
"sync"
|
|
|
)
|
|
|
|
|
|
-// taskHandler 定义了处理单个消息的函数签名
|
|
|
-// 它负责业务逻辑的执行,并返回一个 error 来告知调用者处理是否成功。
|
|
|
-type taskHandler func(ctx context.Context, logger *zap.Logger, delivery amqp091.Delivery) error
|
|
|
+// 使用公共的 TaskHandler 定义
|
|
|
|
|
|
// WhitelistJob 定义了处理白名单相关任务的接口
|
|
|
type WhitelistJob interface {
|
|
@@ -55,66 +52,9 @@ func (j *whitelistJob) IpConsumer(ctx context.Context) {
|
|
|
j.consume(ctx, "ip_white", "ip_white_consumer", j.handleIpMessage)
|
|
|
}
|
|
|
|
|
|
-// consume 是一个通用的 RabbitMQ 消费者方法,封装了重复的逻辑
|
|
|
-func (j *whitelistJob) consume(ctx context.Context, taskName, consumerName string, handler taskHandler) {
|
|
|
- taskCfg, ok := j.Rabbitmq.GetTaskConfig(taskName)
|
|
|
- if !ok {
|
|
|
- j.logger.Error(fmt.Sprintf("未找到任务 '%s' 的配置", taskName))
|
|
|
- return
|
|
|
- }
|
|
|
-
|
|
|
- j.logger.Info("正在启动消费者...",
|
|
|
- zap.String("task", taskName),
|
|
|
- zap.String("queue", taskCfg.Queue),
|
|
|
- zap.String("consumer", consumerName),
|
|
|
- )
|
|
|
-
|
|
|
- msgs, err := j.Rabbitmq.Consume(taskCfg.Queue, consumerName, taskCfg.PrefetchCount)
|
|
|
- if err != nil {
|
|
|
- j.logger.Error("启动消费者失败", zap.String("task", taskName), zap.Error(err))
|
|
|
- return
|
|
|
- }
|
|
|
-
|
|
|
- for {
|
|
|
- select {
|
|
|
- case <-ctx.Done():
|
|
|
- j.logger.Info("消费者正在关闭...", zap.String("task", taskName))
|
|
|
- return
|
|
|
- case d, ok := <-msgs:
|
|
|
- if !ok {
|
|
|
- j.logger.Warn("消息通道已关闭,消费者退出。", zap.String("task", taskName))
|
|
|
- return
|
|
|
- }
|
|
|
-
|
|
|
- // 尝试从消息头获取 trace_id,如果不存在则生成一个新的
|
|
|
- traceID, ok := d.Headers["trace_id"].(string)
|
|
|
- if !ok || traceID == "" {
|
|
|
- traceID = uuid.New().String()
|
|
|
- }
|
|
|
-
|
|
|
- // 创建一个带有 trace_id 的 logger,用于本次任务的所有日志记录
|
|
|
- scopedLogger := j.logger.With(zap.String("trace_id", traceID))
|
|
|
-
|
|
|
- // 创建一个带有 trace_id 的 context,用于传递给下游服务
|
|
|
- ctxWithTrace := context.WithValue(ctx, "trace_id", traceID)
|
|
|
-
|
|
|
- // 调用具体的业务处理器
|
|
|
- processingErr := handler(ctxWithTrace, scopedLogger, d)
|
|
|
-
|
|
|
- // 根据处理结果统一进行 Ack/Nack
|
|
|
- if processingErr != nil {
|
|
|
- // 业务失败,拒绝消息并不重新入队
|
|
|
- if err := d.Nack(false, false); err != nil {
|
|
|
- scopedLogger.Error("消息 Nack 失败", zap.Error(err), zap.String("task", taskName))
|
|
|
- }
|
|
|
- } else {
|
|
|
- // 业务处理成功,手动发送确认
|
|
|
- if err := d.Ack(false); err != nil {
|
|
|
- scopedLogger.Error("消息 Ack 失败", zap.Error(err), zap.String("task", taskName))
|
|
|
- }
|
|
|
- }
|
|
|
- }
|
|
|
- }
|
|
|
+// consume 调用公共的 Consume 方法
|
|
|
+func (j *whitelistJob) consume(ctx context.Context, taskName, consumerName string, handler TaskHandler) {
|
|
|
+ j.Job.Consume(ctx, taskName, consumerName, handler)
|
|
|
}
|
|
|
|
|
|
// handleDomainMessage 是域名白名单任务的具体处理器
|