go_gameshield/pkg/rabbitmq/rabbitmq.go

package rabbitmq

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/go-nunu/nunu-layout-advanced/pkg/log"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"
)

var (
	ErrClosed = errors.New("rabbitmq: client is closed")
)

type Config struct {
	Host              string                  `yaml:"host"`
	Port              int                     `yaml:"port"`
	Username          string                  `yaml:"username"`
	Password          string                  `yaml:"password"`
	VHost             string                  `yaml:"vhost"`
	ConnectionTimeout time.Duration           `yaml:"connection_timeout"`
	Tasks             map[string]TaskConfig   `yaml:"tasks"` // 支持多个任务配置
}

type TaskConfig struct {
	Exchange      string `mapstructure:"exchange"`
	ExchangeType  string `mapstructure:"exchange_type"`
	Queue         string `mapstructure:"queue"`
	RoutingKey    string `mapstructure:"routing_key"`
	ConsumerCount int    `mapstructure:"consumer_count"`
	PrefetchCount int    `mapstructure:"prefetch_count"`
}

type RabbitMQ struct {
	config Config
	conn   *amqp.Connection
	ch     *amqp.Channel
	logger *log.Logger

	mu     sync.RWMutex
	closed bool
}

// New 创建新的RabbitMQ客户端
func New(config Config, logger *log.Logger) (*RabbitMQ, error) {
	r := &RabbitMQ{
		config: config,
		logger: logger,
	}

	if err := r.Connect(); err != nil {
		return nil, err
	}

	if err := r.SetupAllTaskQueues(); err != nil {
		_ = r.Close() // Attempt to close the connection if setup fails
		return nil, fmt.Errorf("failed to setup task queues: %w", err)
	}

	go r.reconnectLoop()

	return r, nil
}

// Connect 连接到RabbitMQ服务器
func (r *RabbitMQ) Connect() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.conn != nil && !r.conn.IsClosed() {
		_ = r.ch.Close()
		_ = r.conn.Close()
	}

	vhost := r.config.VHost
	if vhost == "" {
		vhost = "/"
	} else if vhost[0] != '/' {
		vhost = "/" + vhost
	}

	// 构造完整的连接URL
	fullURL := fmt.Sprintf("amqp://%s:%s@%s:%d%s",
		r.config.Username,
		r.config.Password,
		r.config.Host,
		r.config.Port,
		vhost,
	)

	r.logger.Info("正在尝试连接到 RabbitMQ...", zap.String("url", fullURL))

	var err error
	r.conn, err = amqp.Dial(fullURL)

	if err != nil {
		// 记录详细的底层错误
		r.logger.Error("连接RabbitMQ失败", zap.Error(err))
		return fmt.Errorf("连接RabbitMQ失败: %w", err)
	}

	r.ch, err = r.conn.Channel()
	if err != nil {
		_ = r.conn.Close()
		return fmt.Errorf("创建通道失败: %w", err)
	}

	r.closed = false
	r.logger.Info("RabbitMQ连接成功")
	return nil
}

// reconnectLoop 监控连接状态并处理重连
func (r *RabbitMQ) reconnectLoop() {
	for {
		closeChan := make(chan *amqp.Error)
		r.mu.RLock()
		if r.conn == nil {
			r.mu.RUnlock()
			time.Sleep(5 * time.Second)
			continue
		}
		r.conn.NotifyClose(closeChan)
		isClosed := r.closed
		r.mu.RUnlock()

		if isClosed {
			r.logger.Info("RabbitMQ客户端已关闭，停止重连循环。")
			return
		}

		closeErr := <-closeChan
		if closeErr != nil {
			r.logger.Error("RabbitMQ连接断开，将尝试重新连接", zap.Error(closeErr))
		} else {
			r.logger.Info("RabbitMQ连接正常关闭。")
		}

		r.mu.RLock()
		isClosed = r.closed
		r.mu.RUnlock()
		if isClosed {
			r.logger.Info("RabbitMQ客户端已关闭，停止重连。")
			return
		}

		backoff := 1 * time.Second
		maxBackoff := 30 * time.Second

		for {
			if r.isClosed() {
				return
			}
			err := r.Connect()
			if err == nil {
				r.logger.Info("RabbitMQ重新连接成功")
				// 重新设置任务队列
				if err := r.SetupAllTaskQueues(); err != nil {
					r.logger.Error("重新设置所有任务队列失败", zap.Error(err))
				}
				break
			}

			r.logger.Error("RabbitMQ重连失败", zap.Error(err), zap.Duration("backoff", backoff))

			time.Sleep(backoff)
			backoff *= 2
			if backoff > maxBackoff {
				backoff = maxBackoff
			}
		}
	}
}

// Close 关闭连接
func (r *RabbitMQ) Close() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.closed {
		return nil
	}
	r.closed = true

	var errs []error

	if r.ch != nil {
		if err := r.ch.Close(); err != nil {
			errs = append(errs, fmt.Errorf("关闭channel失败: %w", err))
		}
	}

	if r.conn != nil && !r.conn.IsClosed() {
		if err := r.conn.Close(); err != nil {
			errs = append(errs, fmt.Errorf("关闭connection失败: %w", err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("关闭RabbitMQ时发生错误: %v", errs)
	}

	return nil
}

func (r *RabbitMQ) isClosed() bool {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.closed
}

// GetTaskConfig retrieves a specific task's configuration.
func (r *RabbitMQ) GetTaskConfig(name string) (TaskConfig, bool) {
	taskCfg, ok := r.config.Tasks[name]
	return taskCfg, ok
}

func (r *RabbitMQ) withChannel(fn func(*amqp.Channel) error) error {
	if r.isClosed() {
		return ErrClosed
	}

	r.mu.RLock()
	defer r.mu.RUnlock()

	if r.ch == nil || r.conn.IsClosed() {
		return errors.New("rabbitmq: channel or connection is not available")
	}
	return fn(r.ch)
}

// Publish sends a message to the specified exchange with the given routing key.
// This is a convenience wrapper around PublishWithCh.
func (r *RabbitMQ) Publish(exchange, routingKey string, body []byte) error {
	return r.PublishWithCh(exchange, routingKey, amqp.Publishing{
		ContentType:  "text/plain",
		Body:         body,
		DeliveryMode: amqp.Persistent, // Default to persistent
	})
}

// PublishWithCh sends a message to the specified exchange with the given routing key using a custom amqp.Publishing struct.
func (r *RabbitMQ) PublishWithCh(exchange, routingKey string, msg amqp.Publishing) error {
	if r.isClosed() {
		return errors.New("rabbitmq connection is closed")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	return r.withChannel(func(ch *amqp.Channel) error {
		return ch.PublishWithContext(ctx,
			exchange,
			routingKey,
			false, // mandatory
			false, // immediate
			msg,
		)
	})
}

// Consume 获取消息消费通道. 注意: Qos的设置需要调用方在获取channel后自行处理，或者为Consume方法增加prefetchCount参数
func (r *RabbitMQ) Consume(queue, consumer string, prefetchCount int) (<-chan amqp.Delivery, error) {
	var deliveries <-chan amqp.Delivery
	err := r.withChannel(func(ch *amqp.Channel) error {
		if err := ch.Qos(prefetchCount, 0, false); err != nil {
			return fmt.Errorf("设置Qos失败: %w", err)
		}

		var err error
		deliveries, err = ch.Consume(
			queue,
			consumer,
			false, // auto-ack: false, 手动确认
			false, // exclusive
			false, // no-local
			false, // no-wait
			nil,   // args
		)
		return err
	})
	return deliveries, err
}

// SetupAllTaskQueues 遍历配置中的所有任务，并为每个任务设置队列
func (r *RabbitMQ) SetupAllTaskQueues() error {
	if len(r.config.Tasks) == 0 {
		r.logger.Info("在配置中未找到任何任务队列定义。")
		return nil
	}

	for name, taskCfg := range r.config.Tasks {
		r.logger.Info("正在设置任务队列", zap.String("task_name", name))
		if err := r.setupQueue(taskCfg); err != nil {
			return fmt.Errorf("为任务 '%s' 设置队列失败: %w", name, err)
		}
	}
	return nil
}

// setupQueue 为单个任务配置设置交换机、队列和绑定
func (r *RabbitMQ) setupQueue(taskCfg TaskConfig) error {
	if taskCfg.Exchange == "" {
		r.logger.Warn("任务队列的交换机名称为空，将使用默认交换机。这在多任务场景下可能导致问题。", zap.String("queue", taskCfg.Queue))
		return r.withChannel(func(ch *amqp.Channel) error {
			_, err := ch.QueueDeclare(taskCfg.Queue, true, false, false, false, nil)
			if err != nil {
				return fmt.Errorf("声明队列失败 (默认交换机): %w", err)
			}
			r.logger.Info("成功声明队列并绑定到默认交换机", zap.String("queue", taskCfg.Queue))
			return nil
		})
	}

	return r.withChannel(func(ch *amqp.Channel) error {
		// 声明主交换机
		exchangeType := taskCfg.ExchangeType
	if exchangeType == "" {
		exchangeType = "direct" // 默认为 direct 类型，兼容旧配置
	}
	err := ch.ExchangeDeclare(
		taskCfg.Exchange, // name
		exchangeType,     // type
			true,               // durable
			false,              // autoDelete
			false,              // internal
			false,              // noWait
			nil,                // args
		)
		if err != nil {
			return fmt.Errorf("声明主交换机 '%s' 失败: %w", taskCfg.Exchange, err)
		}

		// 为主队列设置死信交换机参数
		dlxExchange := taskCfg.Exchange + ".dlx"
		args := amqp.Table{
			"x-dead-letter-exchange": dlxExchange,
		}

		// 声明主队列
				_, err = ch.QueueDeclare(taskCfg.Queue, true, false, false, false, args)
		if err != nil {
			return fmt.Errorf("声明主队列 '%s' 失败: %w", taskCfg.Queue, err)
		}

		// 绑定主队列到主交换机
		if err := ch.QueueBind(taskCfg.Queue, taskCfg.RoutingKey, taskCfg.Exchange, false, nil); err != nil {
			return fmt.Errorf("绑定主队列失败: %w", err)
		}

		// --- 设置死信队列 ---
		// 声明死信交换机 (DLX)
		if err := ch.ExchangeDeclare(dlxExchange, "direct", true, false, false, false, nil); err != nil {
			return fmt.Errorf("声明死信交换机 '%s' 失败: %w", dlxExchange, err)
		}

		// 声明死信队列 (DLQ)
		dlq := taskCfg.Queue + ".dlq"
		_, err = ch.QueueDeclare(dlq, true, false, false, false, nil)
		if err != nil {
			return fmt.Errorf("声明死信队列 '%s' 失败: %w", dlq, err)
		}

		// 绑定DLQ到DLX，使用与主队列相同的路由键
		if err := ch.QueueBind(dlq, taskCfg.RoutingKey, dlxExchange, false, nil); err != nil {
			return fmt.Errorf("绑定死信队列失败: %w", err)
		}

		r.logger.Info("成功设置任务队列及其死信队列",
			zap.String("exchange", taskCfg.Exchange),
			zap.String("queue", taskCfg.Queue),
			zap.String("routing_key", taskCfg.RoutingKey),
		)
		return nil
	})
}