backend/pkg/mq/rabbitmq_mq.go

package mq

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/XingfenD/yoresee_doc/pkg/errs"
	"github.com/rabbitmq/amqp091-go"
	"github.com/sirupsen/logrus"
)

type RabbitMQ struct {
	conn    *amqp091.Connection
	channel *amqp091.Channel
}

type RabbitMQConfig struct {
	URL string
}

func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
	conn, err := amqp091.Dial(config.URL)
	if err != nil {
		return nil, errs.Wrap(errs.ErrRabbitMQConnectFailed, err)
	}

	channel, err := conn.Channel()
	if err != nil {
		conn.Close()
		return nil, errs.Wrap(errs.ErrRabbitMQOpenChannelFailed, err)
	}

	return &RabbitMQ{
		conn:    conn,
		channel: channel,
	}, nil
}

func (rmq *RabbitMQ) Publish(ctx context.Context, msg PublishMessage) error {
	topic := strings.TrimSpace(msg.Topic)
	if topic == "" {
		return errs.ErrTopicEmpty
	}
	if err := rmq.declareTopicExchange(topic); err != nil {
		return errs.Wrap(errs.ErrRabbitMQDeclareExchange, err)
	}

	if err := rmq.channel.PublishWithContext(
		ctx,
		topic,
		"",
		false,
		false,
		amqp091.Publishing{
			ContentType: "application/json",
			Body:        msg.Body,
			MessageId:   strings.TrimSpace(msg.Key),
			Headers:     toAMQPHeaders(msg.Headers),
			Timestamp:   time.Now(),
		},
	); err != nil {
		return errs.Wrap(errs.ErrRabbitMQPublishMessage, err)
	}
	return nil
}

func (rmq *RabbitMQ) Consume(ctx context.Context, opts ConsumeOptions, handler MessageHandler) error {
	opts = opts.normalize()
	topic := strings.TrimSpace(opts.Topic)
	if topic == "" {
		return errs.ErrTopicEmpty
	}

	if err := rmq.declareTopicExchange(topic); err != nil {
		return errs.Wrap(errs.ErrRabbitMQDeclareExchange, err)
	}

	queueName, durable, autoDelete, exclusive := consumeQueueConfig(topic, opts)
	queue, err := rmq.channel.QueueDeclare(
		queueName,
		durable,
		autoDelete,
		exclusive,
		false,
		nil,
	)
	if err != nil {
		return errs.Wrap(errs.ErrRabbitMQDeclareQueue, err)
	}

	if err := rmq.channel.QueueBind(
		queue.Name,
		"",
		topic,
		false,
		nil,
	); err != nil {
		return errs.Wrap(errs.ErrRabbitMQBindQueue, err)
	}

	consumerTag := buildConsumerTag(opts.Consumer, topic)
	msgs, err := rmq.channel.Consume(
		queue.Name,
		consumerTag,
		opts.AutoAck,
		false,
		false,
		false,
		nil,
	)
	if err != nil {
		return errs.Wrap(errs.ErrRabbitMQRegisterConsumer, err)
	}

	logrus.Infof("RabbitMQ consume started, topic=%s, queue=%s, mode=%s, group=%s, auto_ack=%v",
		topic, queue.Name, opts.Mode, opts.Group, opts.AutoAck,
	)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case d, ok := <-msgs:
			if !ok {
				return errs.Detailf(errs.ErrRabbitMQConsumerChanClosed, "topic=%s, queue=%s", topic, queue.Name)
			}

			message := Message{
				ID:        d.MessageId,
				Topic:     topic,
				Key:       d.MessageId,
				Body:      d.Body,
				Headers:   fromAMQPHeaders(d.Headers),
				Timestamp: d.Timestamp,
			}
			if message.Timestamp.IsZero() {
				message.Timestamp = time.Now()
			}

			handleErr := handler(context.Background(), message)
			if opts.AutoAck {
				if handleErr != nil {
					logrus.Errorf("RabbitMQ consume handler failed (auto-ack), topic=%s, queue=%s, err=%v", topic, queue.Name, handleErr)
				}
				continue
			}

			if handleErr != nil {
				requeue := opts.OnError == ErrorActionRequeue
				if nackErr := d.Nack(false, requeue); nackErr != nil {
					logrus.Errorf("RabbitMQ nack failed, topic=%s, queue=%s, requeue=%v, err=%v", topic, queue.Name, requeue, nackErr)
				}
				logrus.Errorf("RabbitMQ consume handler failed, topic=%s, queue=%s, requeue=%v, err=%v", topic, queue.Name, requeue, handleErr)
				continue
			}

			if ackErr := d.Ack(false); ackErr != nil {
				logrus.Errorf("RabbitMQ ack failed, topic=%s, queue=%s, err=%v", topic, queue.Name, ackErr)
			}
		}
	}
}

func (rmq *RabbitMQ) declareTopicExchange(topic string) error {
	return rmq.channel.ExchangeDeclare(
		topic,
		"topic",
		true,
		false,
		false,
		false,
		nil,
	)
}

func consumeQueueConfig(topic string, opts ConsumeOptions) (queueName string, durable, autoDelete, exclusive bool) {
	if opts.Mode != ConsumeModeGroup {
		return "", false, true, true
	}
	group := strings.TrimSpace(opts.Group)
	if group == "" {
		group = "default"
	}
	return fmt.Sprintf("%s.%s", topic, group), true, false, false
}

func buildConsumerTag(consumer, topic string) string {
	consumer = strings.TrimSpace(consumer)
	if consumer == "" {
		hostName, _ := os.Hostname()
		consumer = fmt.Sprintf("%s-%s-%d", strings.TrimSpace(topic), strings.TrimSpace(hostName), os.Getpid())
	}
	return consumer
}

func toAMQPHeaders(headers map[string]string) amqp091.Table {
	if len(headers) == 0 {
		return nil
	}
	table := amqp091.Table{}
	for key, val := range headers {
		key = strings.TrimSpace(key)
		if key == "" {
			continue
		}
		table[key] = val
	}
	if len(table) == 0 {
		return nil
	}
	return table
}

func fromAMQPHeaders(table amqp091.Table) map[string]string {
	if len(table) == 0 {
		return nil
	}
	headers := make(map[string]string, len(table))
	for key, val := range table {
		headers[key] = fmt.Sprint(val)
	}
	return headers
}

func (rmq *RabbitMQ) Close() error {
	if rmq.channel != nil {
		rmq.channel.Close()
	}
	if rmq.conn != nil {
		rmq.conn.Close()
	}
	return nil
}