backend/internal/service/mq_service/mq_service.go

package mq_service

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

	"github.com/XingfenD/yoresee_doc/internal/config"
	svc_iface "github.com/XingfenD/yoresee_doc/internal/service/interface"
	"github.com/XingfenD/yoresee_doc/internal/status"
	"github.com/XingfenD/yoresee_doc/internal/utils"
	"github.com/XingfenD/yoresee_doc/pkg/mq"
	"github.com/bytedance/sonic"
	"github.com/sirupsen/logrus"
)

type MQService struct {
	backend mq.Backend
	queues  map[mq.Backend]mq.MessageQueue
}

var MQSvc = &MQService{
	backend: mq.BackendRedis,
	queues:  map[mq.Backend]mq.MessageQueue{},
}

func (srvc *MQService) IsInitialized() bool {
	return srvc != nil && len(srvc.queues) > 0
}

func (srvc *MQService) Init(cfg *config.MQConfig) error {
	if srvc == nil {
		return status.StatusMQNotInitialized
	}

	var options *mq.Options
	if cfg != nil {
		options = &mq.Options{RabbitMQURL: cfg.RabbitMQ.URL}
	}
	queues, err := mq.NewMessageQueues(options)
	if err != nil {
		logrus.Errorf("[Service layer: MQService] Init failed, err=%+v", err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "init message queue failed")
	}
	srvc.queues = queues
	return nil
}

func (srvc *MQService) Close() error {
	if srvc == nil {
		return status.StatusMQNotInitialized
	}

	var firstErr error
	for backend, q := range srvc.queues {
		if q == nil {
			continue
		}
		if err := q.Close(); err != nil {
			logrus.Errorf("[Service layer: MQService] Close backend failed, backend=%s, err=%+v", backend, err)
			if firstErr == nil {
				firstErr = err
			}
		}
	}
	srvc.queues = map[mq.Backend]mq.MessageQueue{}
	return firstErr
}

func (srvc *MQService) Publish(ctx context.Context, backend mq.Backend, msg mq.PublishMessage) error {
	if !srvc.IsInitialized() {
		return status.StatusMQNotInitialized
	}

	msg.Topic = strings.TrimSpace(msg.Topic)
	if msg.Topic == "" {
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "publish message failed")
	}

	q, err := srvc.getQueue(backend)
	if err != nil {
		logrus.Errorf("[Service layer: MQService] Publish getQueue failed, backend=%s, err=%+v", backend, err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "publish message failed")
	}

	if err := q.Publish(ctx, msg); err != nil {
		logrus.Errorf("[Service layer: MQService] Publish failed, backend=%s, topic=%s, err=%+v", backend, msg.Topic, err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "publish message failed")
	}
	return nil
}

func (srvc *MQService) Consume(ctx context.Context, backend mq.Backend, opts mq.ConsumeOptions, handler func(context.Context, mq.Message) error) error {
	if !srvc.IsInitialized() {
		return status.StatusMQNotInitialized
	}
	opts = normalizeConsumeOptions(opts)
	if opts.Topic == "" {
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "subscribe message failed")
	}
	if err := validateConsumeOptions(opts); err != nil {
		logrus.Errorf("[Service layer: MQService] Consume options invalid, backend=%s, topic=%s, err=%+v", backend, opts.Topic, err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "subscribe message failed")
	}

	q, err := srvc.getQueue(backend)
	if err != nil {
		logrus.Errorf("[Service layer: MQService] Consume getQueue failed, backend=%s, topic=%s, err=%+v", backend, opts.Topic, err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "subscribe message failed")
	}

	if err := q.Consume(ctx, opts, handler); err != nil {
		logrus.Errorf("[Service layer: MQService] Consume failed, backend=%s, topic=%s, err=%+v", backend, opts.Topic, err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "subscribe message failed")
	}
	return nil
}

func PublishByProducer[T any](ctx context.Context, producer svc_iface.TopicProducer, taskData T) error {
	if !MQSvc.IsInitialized() {
		return status.StatusMQNotInitialized
	}
	data, err := sonic.Marshal(taskData)
	if err != nil {
		logrus.Errorf("[Service layer: MQService] marshal failed, producer=%s, err=%+v", producer.Topic(), err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "marshal message failed")
	}
	return MQSvc.Publish(ctx, MQSvc.backend, mq.PublishMessage{
		Topic: producer.Topic(),
		Body:  data,
	})
}

func PublishByProducerTo[T any](ctx context.Context, backend mq.Backend, producer svc_iface.TopicProducer, taskData T) error {
	if !MQSvc.IsInitialized() {
		return status.StatusMQNotInitialized
	}
	data, err := sonic.Marshal(taskData)
	if err != nil {
		logrus.Errorf("[Service layer: MQService] marshal failed, backend=%s, producer=%s, err=%+v", backend, producer.Topic(), err)
		return status.GenErrWithCustomMsg(status.StatusServiceInternalError, "marshal message failed")
	}
	return MQSvc.Publish(ctx, backend, mq.PublishMessage{
		Topic: producer.Topic(),
		Body:  data,
	})
}

func normalizeConsumeOptions(opts mq.ConsumeOptions) mq.ConsumeOptions {
	n := opts
	n.Topic = strings.TrimSpace(n.Topic)
	n.Group = strings.TrimSpace(n.Group)
	n.Consumer = strings.TrimSpace(n.Consumer)
	if n.Mode == "" {
		n.Mode = mq.ConsumeModeFanout
	}
	if n.OnError == "" {
		n.OnError = mq.ErrorActionRequeue
	}

	if n.Mode == mq.ConsumeModeGroup {
		n.Group = utils.NormalizeToken(n.Group, "default")
	}

	consumer := n.Consumer
	if consumer == "" {
		hostName, _ := os.Hostname()
		consumer = fmt.Sprintf(
			"%s-%s-%d",
			utils.NormalizeToken(n.Topic, "default"),
			utils.NormalizeToken(hostName, "default"),
			os.Getpid(),
		)
	}
	n.Consumer = utils.NormalizeToken(consumer, "default")

	return n
}

func validateConsumeOptions(opts mq.ConsumeOptions) error {
	switch opts.Mode {
	case mq.ConsumeModeFanout, mq.ConsumeModeGroup:
	default:
		return fmt.Errorf("invalid consume mode: %s", opts.Mode)
	}

	switch opts.OnError {
	case mq.ErrorActionDrop, mq.ErrorActionRequeue:
	default:
		return fmt.Errorf("invalid error action: %s", opts.OnError)
	}

	return nil
}

func (srvc *MQService) getQueue(backend mq.Backend) (mq.MessageQueue, error) {
	if srvc == nil || len(srvc.queues) == 0 {
		return nil, fmt.Errorf("message queue not initialized")
	}
	q, ok := srvc.queues[backend]
	if !ok || q == nil {
		return nil, fmt.Errorf("message queue backend not initialized: %s", backend)
	}
	return q, nil
}