common/message/rabbitmq_msg.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
}