iotlcm/internal/library/rabbitmq/rabbitmq.go

package rabbitmq

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

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

// Config 保存 RabbitMQ 最小连接配置。
type Config struct {
	Enabled bool
	Dsn     string
}

// Handler 是原始字节消息的消费回调。
type Handler func(ctx context.Context, body []byte) error

// JSONHandler 是 JSON 消息的消费回调，框架会先反序列化再交给业务。
type JSONHandler[T any] func(ctx context.Context, message T) error

var (
	conn     *amqp.Connection
	connMu   sync.Mutex
	closers  []func() error
	closerMu sync.Mutex
)

// GetConfig 从配置中心读取 RabbitMQ 开关和 DSN。
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。
func Enabled(ctx context.Context) bool {
	cfg := GetConfig(ctx)
	return cfg.Enabled && cfg.Dsn != ""
}

// Conn 返回共享连接，首次调用时懒加载创建。
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)
}

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

// Consume 启动普通消费者，业务成功后才 Ack。
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")
	}

	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
				}
				_ = delivery.Ack(false)
			}
		}
	}()
	return nil
}

// ConsumeJSON 启动 JSON 消费者，自动完成反序列化。
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 = nil
	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
}

// registerCloser 记录 channel 关闭函数，供服务退出时统一清理。
func registerCloser(closer func() error) {
	closerMu.Lock()
	defer closerMu.Unlock()
	closers = append(closers, closer)
}