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