294 lines
7.3 KiB
Go
294 lines
7.3 KiB
Go
package message
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
amqp "github.com/rabbitmq/amqp091-go"
|
||
)
|
||
|
||
type RabbitMQPublishMsgConfig struct {
|
||
QueueName string
|
||
Durable bool
|
||
DelayTime int
|
||
Data any
|
||
}
|
||
|
||
type RabbitMQSubscribeMsgConfig struct {
|
||
QueueName string
|
||
Durable bool
|
||
DelayTime int
|
||
ConsumerName string
|
||
AutoAck bool
|
||
PrefetchCount int
|
||
HandleFunc func(ctx context.Context, message map[string]interface{}) error
|
||
}
|
||
|
||
func (*RabbitMQPublishMsgConfig) GetPublishMsgType() {
|
||
|
||
}
|
||
|
||
func (*RabbitMQSubscribeMsgConfig) GetSubscribeMsgType() {
|
||
|
||
}
|
||
|
||
func init() {
|
||
// 注册 RabbitMQ 插件，必须使用 RegisterPlugin 确保连接检测
|
||
//registerPlugin(MessageRabbitMQ, func() messageUtil {
|
||
// return &rabbitMQ{}
|
||
//})
|
||
}
|
||
|
||
type rabbitMQ struct{}
|
||
|
||
// Ping 检测 RabbitMQ 连接状态
|
||
func (c *rabbitMQ) ping(ctx context.Context) bool {
|
||
return rabbitmqPing()
|
||
}
|
||
|
||
// Reconnect 重连 RabbitMQ
|
||
func (c *rabbitMQ) reconnect(ctx context.Context) error {
|
||
return rabbitmqReconnect(ctx)
|
||
}
|
||
|
||
// Close 关闭 RabbitMQ 连接
|
||
func (c *rabbitMQ) close(ctx context.Context) error {
|
||
return rabbitmqClose(ctx)
|
||
}
|
||
|
||
// Publish 发布消息
|
||
func (c *rabbitMQ) Publish(ctx context.Context, msgConfig messagePublishConfig) error {
|
||
cfg, ok := msgConfig.(*RabbitMQPublishMsgConfig)
|
||
if !ok {
|
||
return fmt.Errorf("无效的 RabbitMQ 配置类型")
|
||
}
|
||
if g.IsEmpty(cfg.QueueName) {
|
||
return fmt.Errorf("队列名称不能为空")
|
||
}
|
||
if cfg.Data == nil {
|
||
return fmt.Errorf("数据不能为空")
|
||
}
|
||
return c.publishMessageInternal(ctx, cfg.QueueName, cfg.Durable, cfg.DelayTime, cfg.Data)
|
||
}
|
||
|
||
// publishMessage 发布消息内部实现
|
||
func (c *rabbitMQ) publishMessageInternal(ctx context.Context, queueName string, durable bool, delayTime int, data interface{}) error {
|
||
delayMsg := delayTime > 0
|
||
|
||
// 1. 决定 Exchange 类型
|
||
exchangeType := "fanout"
|
||
exchangeName := queueName
|
||
routingKey := queueName
|
||
args := amqp.Table{}
|
||
if delayMsg {
|
||
exchangeType = "x-delayed-message"
|
||
exchangeName = queueName + ".delayed"
|
||
args["x-delayed-type"] = "fanout" // 底层用 topic
|
||
}
|
||
|
||
// 2. 声明 Exchange（只声明一次）
|
||
if err := channel.ExchangeDeclare(
|
||
queueName, // exchange 交换机名称
|
||
exchangeType,
|
||
durable,
|
||
false, // autoDelete
|
||
false, // internal
|
||
false, // noWait
|
||
args,
|
||
); err != nil {
|
||
return fmt.Errorf("声明 Exchange 失败: %w", err)
|
||
}
|
||
|
||
// 3. 声明队列
|
||
if _, err := channel.QueueDeclare(
|
||
queueName,
|
||
durable,
|
||
false, // autoDelete
|
||
false, // exclusive
|
||
false, // noWait
|
||
nil, // args
|
||
); err != nil {
|
||
return fmt.Errorf("声明队列失败: %w", err)
|
||
}
|
||
|
||
// 4. 绑定队列
|
||
if err := channel.QueueBind(
|
||
queueName,
|
||
routingKey, // routingKey 路由键
|
||
exchangeName, // exchange 交换机名称
|
||
false, // noWait
|
||
nil, // args
|
||
); err != nil {
|
||
return fmt.Errorf("绑定队列失败: %w", err)
|
||
}
|
||
|
||
// 5. 序列化数据
|
||
body, err := json.Marshal(data)
|
||
if err != nil {
|
||
return fmt.Errorf("序列化数据失败: %w", err)
|
||
}
|
||
// 6. 发布消息
|
||
deliveryMode := amqp.Transient
|
||
if durable {
|
||
deliveryMode = amqp.Persistent
|
||
}
|
||
publishing := amqp.Publishing{
|
||
ContentType: "application/json",
|
||
Body: body,
|
||
DeliveryMode: deliveryMode,
|
||
Timestamp: time.Now(),
|
||
}
|
||
if delayMsg {
|
||
duration := time.Duration(delayTime) * time.Minute
|
||
publishing.Headers = amqp.Table{
|
||
"x-delay": duration, // 延迟时间（毫秒）
|
||
}
|
||
}
|
||
err = channel.PublishWithContext(
|
||
ctx,
|
||
exchangeName,
|
||
routingKey,
|
||
false, false,
|
||
publishing,
|
||
)
|
||
return err
|
||
}
|
||
|
||
// Subscribe 订阅消息
|
||
func (c *rabbitMQ) Subscribe(ctx context.Context, msgConfig messageSubscribeConfig) error {
|
||
cfg, ok := msgConfig.(*RabbitMQSubscribeMsgConfig)
|
||
if !ok {
|
||
return fmt.Errorf("无效的 RabbitMQ 配置类型")
|
||
}
|
||
if g.IsEmpty(cfg.QueueName) {
|
||
return fmt.Errorf("队列名称不能为空")
|
||
}
|
||
if g.IsEmpty(cfg.ConsumerName) {
|
||
return fmt.Errorf("消费者名称不能为空")
|
||
}
|
||
if g.IsEmpty(cfg.PrefetchCount) {
|
||
cfg.PrefetchCount = 1
|
||
}
|
||
if g.IsEmpty(cfg.HandleFunc) {
|
||
return fmt.Errorf("必须提供处理函数")
|
||
}
|
||
return c.createSubscribeInternal(ctx, cfg.QueueName, cfg.ConsumerName, cfg.PrefetchCount, cfg.AutoAck, cfg.HandleFunc)
|
||
}
|
||
|
||
// createSubscribe 内部订阅消息
|
||
func (c *rabbitMQ) createSubscribeInternal(ctx context.Context, queueName, consumerName string, prefetchCount int, autoAck bool, handler func(ctx context.Context, message map[string]interface{}) error) error {
|
||
g.Log().Infof(ctx, "🔔 RabbitMQ 开始订阅: queueName=%s, consumerName=%s", queueName, consumerName)
|
||
|
||
if err := channel.Qos(prefetchCount, 0, false); err != nil {
|
||
return fmt.Errorf("设置 Qos 失败: %w", err)
|
||
}
|
||
g.Log().Infof(ctx, "📊 设置 Prefetch Count: %d", prefetchCount)
|
||
|
||
msg, err := channel.Consume(
|
||
queueName, // queue
|
||
consumerName, // consumer
|
||
autoAck, // auto-ack (根据配置决定)
|
||
false, // exclusive
|
||
false, // no-local
|
||
false, // no-wait
|
||
nil, // args
|
||
)
|
||
if err != nil {
|
||
return fmt.Errorf("注册消费者失败: %w", err)
|
||
}
|
||
|
||
go func() {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
g.Log().Errorf(ctx, "❌ RabbitMQ 消费者 panic: %v", r)
|
||
}
|
||
}()
|
||
|
||
// 并发控制信号量
|
||
semaphore := make(chan struct{}, 10) // 限制最大并发数为 10
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
g.Log().Infof(ctx, "🔕 RabbitMQ 消费者停止: queueName=%s, consumerName=%s", queueName, consumerName)
|
||
return
|
||
case msg, ok := <-msg:
|
||
if !ok {
|
||
g.Log().Warningf(ctx, "⚠️ RabbitMQ 消息通道关闭")
|
||
return
|
||
}
|
||
|
||
// 获取并发控制槽位
|
||
semaphore <- struct{}{}
|
||
|
||
go func(m amqp.Delivery) {
|
||
defer func() {
|
||
<-semaphore // 释放槽位
|
||
if r := recover(); r != nil {
|
||
g.Log().Errorf(ctx, "❌ 消息处理 panic: %v", r)
|
||
}
|
||
}()
|
||
|
||
if err := c.handleMessageWithRetryInternal(ctx, m, handler, autoAck); err != nil {
|
||
g.Log().Errorf(ctx, "❌ 消息处理失败(重试次数耗尽): %v", err)
|
||
|
||
// 仅在手动 ACK 模式下拒绝消息
|
||
if !autoAck {
|
||
// 拒绝消息不再重新入队（避免死循环）
|
||
m.Nack(false, false)
|
||
}
|
||
return
|
||
}
|
||
|
||
// 仅在手动 ACK 模式下确认消息
|
||
if autoAck {
|
||
if err := m.Ack(false); err != nil {
|
||
g.Log().Errorf(ctx, "❌ ACK 消息失败: %v", err)
|
||
}
|
||
}
|
||
}(msg)
|
||
}
|
||
}
|
||
}()
|
||
|
||
return nil
|
||
}
|
||
|
||
// handleMessageWithRetry 处理消息（支持重试）
|
||
func (c *rabbitMQ) handleMessageWithRetryInternal(ctx context.Context, msg amqp.Delivery, handler func(ctx context.Context, message map[string]interface{}) error, autoAck bool) error {
|
||
var data map[string]interface{}
|
||
|
||
if err := json.Unmarshal(msg.Body, &data); err != nil {
|
||
// 如果不是 JSON，直接使用原始内容
|
||
data = map[string]interface{}{
|
||
"data": string(msg.Body),
|
||
}
|
||
}
|
||
|
||
// 重试逻辑
|
||
const maxRetry = 3
|
||
for attempt := 0; attempt <= maxRetry; attempt++ {
|
||
if attempt > 0 {
|
||
g.Log().Infof(ctx, "🔄 消息处理重试 (第%d次)", attempt)
|
||
// 指数退避
|
||
time.Sleep(time.Duration(attempt) * time.Second)
|
||
}
|
||
|
||
err := handler(ctx, data)
|
||
if err == nil {
|
||
return nil // 成功
|
||
}
|
||
|
||
g.Log().Warningf(ctx, "⚠️ 消息处理失败 (第%d次): %v", attempt+1, err)
|
||
|
||
if attempt == maxRetry {
|
||
return fmt.Errorf("达到最大重试次数 %d: %w", maxRetry, err)
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|