iotlcm/internal/library/rabbitmq/rabbitmq.go

package rabbitmq

import (
	"context"
	"encoding/json"
	"errors"
	"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
}

type channelLifecycle interface {
	Cancel(consumer string, noWait bool) error
	Close() error
	IsClosed() bool
}

type managedConsumer struct {
	channel channelLifecycle
	tag     string
	done    chan struct{}
}

// 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
	consumers  []*managedConsumer
	consumerMu sync.Mutex
	closeMu    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 {
	if delay <= 0 {
		return gerror.New("rabbitmq delayed message delay must be positive")
	}
	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,
	})
}

// PublishDelayedBytesWithHeaders 发布带消息头的原始消息体到延迟队列。
func PublishDelayedBytesWithHeaders(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte, headers amqp.Table) error {
	if delay <= 0 {
		return gerror.New("rabbitmq delayed message delay must be positive")
	}
	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 publishWithConfirm(ctx, ch, delayQueue, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Expiration:   fmt.Sprintf("%d", delay.Milliseconds()),
		Timestamp:    time.Now(),
		Headers:      headers,
		Body:         body,
	})
}

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

	if err = declareQueue(ch, queue); err != nil {
		return err
	}
	return publishWithConfirm(ctx, ch, queue, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		Headers:      headers,
		Body:         body,
	})
}

func publishWithConfirm(ctx context.Context, ch *amqp.Channel, queue string, message amqp.Publishing) error {
	if err := ch.Confirm(false); err != nil {
		return gerror.Wrap(err, "rabbitmq enable publisher confirm failed")
	}
	confirms := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
	if err := ch.PublishWithContext(ctx, "", queue, false, false, message); err != nil {
		return gerror.Wrapf(err, "rabbitmq publish failed, queue:%s", queue)
	}
	select {
	case confirmation, ok := <-confirms:
		if !ok {
			return gerror.Newf("rabbitmq publisher confirm channel closed, queue:%s", queue)
		}
		if !confirmation.Ack {
			return gerror.Newf("rabbitmq message was not confirmed, queue:%s", queue)
		}
		return nil
	case <-ctx.Done():
		return gerror.Wrapf(ctx.Err(), "wait rabbitmq publisher confirm failed, queue:%s", queue)
	}
}

// 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")
	}

	consumerState := &managedConsumer{
		channel: ch,
		tag:     consumer,
		done:    make(chan struct{}),
	}
	registerConsumer(consumerState)
	go func() {
		defer close(consumerState.done)
		for {
			if ctx.Err() != nil {
				return
			}
			select {
			case <-ctx.Done():
				return
			case delivery, ok := <-deliveries:
				if !ok {
					return
				}
				consumerTag := delivery.ConsumerTag
				if consumerTag == "" {
					consumerTag = consumer
				}
				deliveryCtx := WithDeliveryMetadata(ctx, DeliveryMetadata{
					DeliveryTag: delivery.DeliveryTag,
					Redelivered: delivery.Redelivered,
					Consumer:    consumerTag,
					Queue:       queue,
					Headers:     cloneHeaders(delivery.Headers),
				})
				if handlerErr := handler(deliveryCtx, delivery.Body); handlerErr != nil {
					g.Log().Errorf(deliveryCtx, "rabbitmq message Nack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t, err:%+v", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered, handlerErr)
					_ = delivery.Nack(false, true)
					continue
				}
				// handler 内部已经完成业务处理或重新发布下一次重试消息，这里只确认当前消息。
				g.Log().Infof(deliveryCtx, "rabbitmq message Ack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered)
				_ = 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(ctx context.Context) error {
	closeMu.Lock()
	defer closeMu.Unlock()

	consumerMu.Lock()
	registeredConsumers := consumers
	consumers = nil
	consumerMu.Unlock()

	var closeErrors []error
	for _, consumer := range registeredConsumers {
		if consumer.channel.IsClosed() {
			continue
		}
		if err := consumer.channel.Cancel(consumer.tag, false); !isClosedError(err) {
			closeErrors = append(closeErrors, gerror.Wrapf(err, "cancel RabbitMQ consumer failed, consumer:%s", consumer.tag))
		}
	}

	if err := waitForConsumers(ctx, registeredConsumers); err != nil {
		closeErrors = append(closeErrors, err)
	}

	for _, consumer := range registeredConsumers {
		if consumer.channel.IsClosed() {
			continue
		}
		if err := consumer.channel.Close(); !isClosedError(err) {
			closeErrors = append(closeErrors, gerror.Wrapf(err, "close RabbitMQ channel failed, consumer:%s", consumer.tag))
		}
	}

	if err := closeConnection(); err != nil {
		closeErrors = append(closeErrors, err)
	}
	return errors.Join(closeErrors...)
}

// 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 registerConsumer(consumer *managedConsumer) {
	consumerMu.Lock()
	defer consumerMu.Unlock()
	consumers = append(consumers, consumer)
}

func waitForConsumers(ctx context.Context, registeredConsumers []*managedConsumer) error {
	for _, consumer := range registeredConsumers {
		select {
		case <-consumer.done:
		case <-ctx.Done():
			return gerror.Wrap(ctx.Err(), "wait for RabbitMQ consumers failed")
		}
	}
	return nil
}

func closeConnection() error {
	connMu.Lock()
	defer connMu.Unlock()
	if conn == nil || conn.IsClosed() {
		conn = nil
		return nil
	}
	if err := conn.Close(); !isClosedError(err) {
		return gerror.Wrap(err, "close RabbitMQ connection failed")
	}
	conn = nil
	return nil
}

func cloneHeaders(headers amqp.Table) map[string]any {
	if len(headers) == 0 {
		return nil
	}
	cloned := make(map[string]any, len(headers))
	for key, value := range headers {
		cloned[key] = value
	}
	return cloned
}

func isClosedError(err error) bool {
	return err == nil || errors.Is(err, amqp.ErrClosed)
}