iotlcm/admin-svr/server/internal/library/rabbitmq/rabbitmq.go

package rabbitmq

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	amqp "github.com/rabbitmq/amqp091-go"
)

type Config struct {
	Enabled bool
	Dsn     string
}

// Handler 是 RabbitMQ 消费处理函数。
// 返回 nil 时消息 Ack；返回 error 时消息 Nack 且不重新入队。
type Handler func(ctx context.Context, body []byte) error

// JSONHandler 是 RabbitMQ JSON 消息的业务处理函数。
type JSONHandler[T any] func(ctx context.Context, message T) error

var (
	conn     *amqp.Connection
	connMu   sync.Mutex
	closers  = make(map[uint64]func() error)
	closerID uint64
	closerMu sync.Mutex
)

func GetConfig(ctx context.Context) Config {
	return Config{
		Enabled: g.Cfg().MustGet(ctx, "rabbitmq.enabled", false).Bool(),
		Dsn:     g.Cfg().MustGet(ctx, "rabbitmq.dsn", "").String(),
	}
}

// Enabled 判断 RabbitMQ 是否启用且 DSN 是否已配置。
func Enabled(ctx context.Context) bool {
	cfg := GetConfig(ctx)
	return cfg.Enabled && cfg.Dsn != ""
}

// Conn 懒加载 RabbitMQ 连接，供多个生产者/消费者复用。
func Conn(ctx context.Context) (*amqp.Connection, error) {
	cfg := GetConfig(ctx)
	if !cfg.Enabled {
		return nil, gerror.New("rabbitmq is disabled")
	}
	if cfg.Dsn == "" {
		return nil, gerror.New("rabbitmq dsn is empty")
	}

	connMu.Lock()
	defer connMu.Unlock()
	if conn != nil && !conn.IsClosed() {
		return conn, nil
	}

	c, err := amqp.Dial(cfg.Dsn)
	if err != nil {
		return nil, gerror.Wrap(err, "rabbitmq connect failed")
	}
	conn = c
	return conn, nil
}

// Channel 基于共享连接创建独立 channel。
func Channel(ctx context.Context) (*amqp.Channel, error) {
	c, err := Conn(ctx)
	if err != nil {
		return nil, err
	}
	ch, err := c.Channel()
	if err != nil {
		return nil, gerror.Wrap(err, "rabbitmq open channel failed")
	}
	return ch, nil
}

// DeclareQueue 声明一个持久化队列。
func DeclareQueue(ctx context.Context, queue string) error {
	ch, err := Channel(ctx)
	if err != nil {
		return err
	}
	defer ch.Close()
	return declareQueue(ch, queue)
}

// Publish 将结构体序列化为 JSON 后发布到指定队列。
func Publish(ctx context.Context, queue string, data any) error {
	body, err := json.Marshal(data)
	if err != nil {
		return gerror.Wrap(err, "rabbitmq marshal message failed")
	}
	return PublishBytes(ctx, queue, body)
}

// PublishDelayed 将结构体序列化为 JSON 后发布到延迟队列。
// 消息在 delayQueue 中等待 delay 后，通过死信路由进入 targetQueue。
func PublishDelayed(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, data any) error {
	body, err := json.Marshal(data)
	if err != nil {
		return gerror.Wrap(err, "rabbitmq marshal delayed message failed")
	}
	return PublishDelayedBytes(ctx, delayQueue, targetQueue, delay, body)
}

// PublishBytes 发布原始消息体到指定队列。
func PublishBytes(ctx context.Context, queue string, body []byte) error {
	ch, err := Channel(ctx)
	if err != nil {
		return err
	}
	defer ch.Close()

	// 发布前也声明一次队列，避免生产者先启动时队列不存在。
	if err = declareQueue(ch, queue); err != nil {
		return err
	}
	return ch.PublishWithContext(ctx, "", queue, false, false, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		Body:         body,
	})
}

// PublishDelayedBytes 发布原始消息体到延迟队列。
func PublishDelayedBytes(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte) error {
	ch, err := Channel(ctx)
	if err != nil {
		return err
	}
	defer ch.Close()

	if err = declareQueue(ch, targetQueue); err != nil {
		return err
	}
	if err = declareDelayQueue(ch, delayQueue, targetQueue); err != nil {
		return err
	}
	return ch.PublishWithContext(ctx, "", delayQueue, false, false, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Expiration:   fmt.Sprintf("%d", delay.Milliseconds()),
		Timestamp:    time.Now(),
		Body:         body,
	})
}

// Consume 启动消费者，使用手动 Ack，handler 成功才确认消息。
func Consume(ctx context.Context, queue, consumer string, handler Handler) error {
	ch, err := Channel(ctx)
	if err != nil {
		return err
	}
	if err = declareQueue(ch, queue); err != nil {
		_ = ch.Close()
		return err
	}
	if err = ch.Qos(1, 0, false); err != nil {
		_ = ch.Close()
		return gerror.Wrap(err, "rabbitmq set qos failed")
	}

	// autoAck=false，业务处理成功后手动 Ack，处理失败时 Nack 且不重新入队。
	deliveries, err := ch.Consume(queue, consumer, false, false, false, false, nil)
	if err != nil {
		_ = ch.Close()
		return gerror.Wrap(err, "rabbitmq consume failed")
	}

	registerCloser(ch.Close)
	go func() {
		for {
			select {
			case <-ctx.Done():
				_ = ch.Close()
				return
			case delivery, ok := <-deliveries:
				if !ok {
					return
				}
				if err = handler(ctx, delivery.Body); err != nil {
					g.Log().Errorf(ctx, "rabbitmq handle message failed, queue:%s, err:%+v", queue, err)
					_ = delivery.Nack(false, false)
					continue
				}
				// handler 内部已经完成业务处理或重新发布下一次重试消息，这里只确认当前消息。
				_ = delivery.Ack(false)
			}
		}
	}()
	return nil
}

// ConsumeJSON 启动 JSON 消费者，统一完成反序列化和 Ack/Nack 处理。
func ConsumeJSON[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error {
	return Consume(ctx, queue, consumer, func(ctx context.Context, body []byte) error {
		var message T
		if err := json.Unmarshal(body, &message); err != nil {
			g.Log().Errorf(ctx, "rabbitmq message dropped, invalid json, queue:%s, body:%s, err:%+v", queue, string(body), err)
			return nil
		}
		return handler(ctx, message)
	})
}

// Close 关闭消费者 channel 和共享连接。
func Close() error {
	closerMu.Lock()
	for _, closer := range closers {
		_ = closer()
	}
	closers = make(map[uint64]func() error)
	closerMu.Unlock()

	connMu.Lock()
	defer connMu.Unlock()
	if conn == nil {
		return nil
	}
	err := conn.Close()
	conn = nil
	return err
}

// declareQueue 统一声明持久化队列，生产者和消费者都会调用，保证队列存在。
func declareQueue(ch *amqp.Channel, queue string) error {
	_, err := ch.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		return gerror.Wrapf(err, "rabbitmq declare queue failed, queue:%s", queue)
	}
	return nil
}

// declareDelayQueue 声明 TTL 延迟队列，到期后死信转发到目标队列。
func declareDelayQueue(ch *amqp.Channel, delayQueue, targetQueue string) error {
	_, err := ch.QueueDeclare(delayQueue, true, false, false, false, amqp.Table{
		"x-dead-letter-exchange":    "",
		"x-dead-letter-routing-key": targetQueue,
	})
	if err != nil {
		return gerror.Wrapf(err, "rabbitmq declare delay queue failed, queue:%s", delayQueue)
	}
	return nil
}

// registerCloser 记录消费者 channel，服务退出时统一关闭。
func registerCloser(closer func() error) func() {
	closerMu.Lock()
	defer closerMu.Unlock()
	closerID++
	id := closerID
	closers[id] = closer
	return func() {
		closerMu.Lock()
		defer closerMu.Unlock()
		delete(closers, id)
	}
}

func ConsumeJSONRequeue[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error {
	ctx, cancel := context.WithCancel(ctx)
	open := func() (*amqp.Channel, <-chan amqp.Delivery, error) {
		channel, err := Channel(ctx)
		if err != nil {
			return nil, nil, err
		}
		if err = declareQueue(channel, queue); err == nil {
			err = channel.Qos(1, 0, false)
		}
		if err != nil {
			_ = channel.Close()
			return nil, nil, err
		}
		deliveries, err := channel.Consume(queue, consumer, false, false, false, false, nil)
		if err != nil {
			_ = channel.Close()
			return nil, nil, err
		}
		return channel, deliveries, nil
	}
	channel, deliveries, err := open()
	if err != nil {
		cancel()
		return err
	}
	var channelMu sync.Mutex
	unregister := registerCloser(func() error {
		cancel()
		channelMu.Lock()
		defer channelMu.Unlock()
		return channel.Close()
	})
	go func() {
		defer func() {
			cancel()
			unregister()
			_ = channel.Close()
		}()
		for {
			select {
			case <-ctx.Done():
				return
			case delivery, ok := <-deliveries:
				if !ok {
					_ = channel.Close()
					for {
						select {
						case <-ctx.Done():
							return
						case <-time.After(3 * time.Second):
						}
						next, stream, reconnectErr := open()
						if reconnectErr != nil {
							g.Log().Errorf(ctx, "role sync consumer reconnect failed, queue:%s, err:%v", queue, reconnectErr)
							continue
						}
						channelMu.Lock()
						if ctx.Err() != nil {
							channelMu.Unlock()
							_ = next.Close()
							return
						}
						channel, deliveries = next, stream
						channelMu.Unlock()
						break
					}
					continue
				}
				var message T
				if decodeErr := json.Unmarshal(delivery.Body, &message); decodeErr != nil {
					g.Log().Errorf(ctx, "rabbitmq invalid role sync message, queue:%s, err:%v", queue, decodeErr)
					_ = delivery.Ack(false)
					continue
				}
				if handleErr := handler(ctx, message); handleErr != nil {
					g.Log().Errorf(ctx, "role sync message retained, queue:%s, err:%v", queue, handleErr)
					select {
					case <-ctx.Done():
						return
					case <-time.After(time.Second):
					}
					_ = delivery.Nack(false, true)
					continue
				}
				_ = delivery.Ack(false)
			}
		}
	}()
	return nil
}