package rabbitmq import ( "context" "encoding/json" "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 } // 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 closers = make(map[uint64]func() error) closerID uint64 closerMu 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 { 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, }) } // 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") } 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 } // handler 内部已经完成业务处理或重新发布下一次重试消息,这里只确认当前消息。 _ = 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() error { closerMu.Lock() for _, closer := range closers { _ = closer() } closers = make(map[uint64]func() error) 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 } // 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 registerCloser(closer func() error) func() { closerMu.Lock() defer closerMu.Unlock() closerID++ id := closerID closers[id] = closer return func() { closerMu.Lock() defer closerMu.Unlock() delete(closers, id) } } func ConsumeJSONRequeue[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error { ctx, cancel := context.WithCancel(ctx) open := func() (*amqp.Channel, <-chan amqp.Delivery, error) { channel, err := Channel(ctx) if err != nil { return nil, nil, err } if err = declareQueue(channel, queue); err == nil { err = channel.Qos(1, 0, false) } if err != nil { _ = channel.Close() return nil, nil, err } deliveries, err := channel.Consume(queue, consumer, false, false, false, false, nil) if err != nil { _ = channel.Close() return nil, nil, err } return channel, deliveries, nil } channel, deliveries, err := open() if err != nil { cancel() return err } var channelMu sync.Mutex unregister := registerCloser(func() error { cancel() channelMu.Lock() defer channelMu.Unlock() return channel.Close() }) go func() { defer func() { cancel() unregister() _ = channel.Close() }() for { select { case <-ctx.Done(): return case delivery, ok := <-deliveries: if !ok { _ = channel.Close() for { select { case <-ctx.Done(): return case <-time.After(3 * time.Second): } next, stream, reconnectErr := open() if reconnectErr != nil { g.Log().Errorf(ctx, "role sync consumer reconnect failed, queue:%s, err:%v", queue, reconnectErr) continue } channelMu.Lock() if ctx.Err() != nil { channelMu.Unlock() _ = next.Close() return } channel, deliveries = next, stream channelMu.Unlock() break } continue } var message T if decodeErr := json.Unmarshal(delivery.Body, &message); decodeErr != nil { g.Log().Errorf(ctx, "rabbitmq invalid role sync message, queue:%s, err:%v", queue, decodeErr) _ = delivery.Ack(false) continue } if handleErr := handler(ctx, message); handleErr != nil { g.Log().Errorf(ctx, "role sync message retained, queue:%s, err:%v", queue, handleErr) select { case <-ctx.Done(): return case <-time.After(time.Second): } _ = delivery.Nack(false, true) continue } _ = delivery.Ack(false) } } }() return nil }