huang
b7369a9c71
All checks were successful
构建并部署到测试环境（无 SSH） / build-and-deploy (push) Successful in 9m22s
99 lines
2.6 KiB
Go
99 lines
2.6 KiB
Go
package queue
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
|
||
"github.com/break/junhong_cmp_fiber/pkg/config"
|
||
"github.com/break/junhong_cmp_fiber/pkg/constants"
|
||
"github.com/bytedance/sonic"
|
||
"github.com/hibiken/asynq"
|
||
"github.com/redis/go-redis/v9"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// Client Asynq 任务提交客户端
|
||
type Client struct {
|
||
client *asynq.Client
|
||
logger *zap.Logger
|
||
}
|
||
|
||
// NewClient 创建新的 Asynq 客户端
|
||
func NewClient(redisClient *redis.Client, logger *zap.Logger) *Client {
|
||
// 从 Redis 客户端获取配置
|
||
opts := redisClient.Options()
|
||
|
||
asynqClient := asynq.NewClient(asynq.RedisClientOpt{
|
||
Addr: opts.Addr,
|
||
Password: opts.Password,
|
||
DB: opts.DB,
|
||
})
|
||
|
||
return &Client{
|
||
client: asynqClient,
|
||
logger: logger,
|
||
}
|
||
}
|
||
|
||
// EnqueueTask 提交任务到队列
|
||
// payload 必须传入 struct 或 map，禁止传入 []byte。
|
||
// 传入 []byte 会导致 sonic.Marshal 将其 base64 编码，handler 解析时类型不匹配。
|
||
func (c *Client) EnqueueTask(ctx context.Context, taskType string, payload interface{}, opts ...asynq.Option) error {
|
||
if _, ok := payload.([]byte); ok {
|
||
c.logger.Error("任务载荷类型错误，禁止传入 []byte",
|
||
zap.String("task_type", taskType))
|
||
return fmt.Errorf("task payload must be struct or map, got []byte")
|
||
}
|
||
|
||
// 内部统一序列化，调用方无需预先 Marshal
|
||
payloadBytes, err := sonic.Marshal(payload)
|
||
if err != nil {
|
||
c.logger.Error("任务载荷序列化失败",
|
||
zap.String("task_type", taskType),
|
||
zap.Error(err))
|
||
return fmt.Errorf("failed to marshal task payload: %w", err)
|
||
}
|
||
|
||
opts = append(opts, asynq.Queue(constants.QueueForTaskType(taskType)))
|
||
|
||
// 创建任务
|
||
task := asynq.NewTask(taskType, payloadBytes, opts...)
|
||
|
||
// 提交任务
|
||
info, err := c.client.EnqueueContext(ctx, task)
|
||
if err != nil {
|
||
c.logger.Error("任务提交失败",
|
||
zap.String("task_type", taskType),
|
||
zap.Error(err))
|
||
return fmt.Errorf("failed to enqueue task: %w", err)
|
||
}
|
||
|
||
c.logger.Info("任务已提交",
|
||
zap.String("task_id", info.ID),
|
||
zap.String("task_type", taskType),
|
||
zap.String("queue", info.Queue),
|
||
zap.Int("max_retry", info.MaxRetry))
|
||
|
||
return nil
|
||
}
|
||
|
||
// Close 关闭客户端
|
||
func (c *Client) Close() error {
|
||
if c.client != nil {
|
||
return c.client.Close()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ParseQueueConfig 解析队列配置为 Asynq 格式
|
||
func ParseQueueConfig(cfg *config.QueueConfig) map[string]int {
|
||
queues := constants.DefaultTaskQueueWeights()
|
||
if cfg.Queues != nil && len(cfg.Queues) > 0 {
|
||
// 配置只覆盖权重；任务队列全集由代码兜底，避免漏监听新队列。
|
||
for queueName, weight := range cfg.Queues {
|
||
queues[queueName] = weight
|
||
}
|
||
}
|
||
return queues
|
||
}
|