package rabbitmq import ( "context" "encoding/json" "errors" "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 } type channelLifecycle interface { Cancel(consumer string, noWait bool) error Close() error IsClosed() bool } type managedConsumer struct { channel channelLifecycle tag string done chan struct{} } // 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 consumers []*managedConsumer consumerMu sync.Mutex closeMu 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 { if delay <= 0 { return gerror.New("rabbitmq delayed message delay must be positive") } 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, }) } // PublishDelayedBytesWithHeaders 发布带消息头的原始消息体到延迟队列。 func PublishDelayedBytesWithHeaders(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte, headers amqp.Table) error { if delay <= 0 { return gerror.New("rabbitmq delayed message delay must be positive") } 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 publishWithConfirm(ctx, ch, delayQueue, amqp.Publishing{ ContentType: "application/json", DeliveryMode: amqp.Persistent, Expiration: fmt.Sprintf("%d", delay.Milliseconds()), Timestamp: time.Now(), Headers: headers, Body: body, }) } // PublishBytesWithHeaders 发布带消息头的原始消息体到指定队列。 func PublishBytesWithHeaders(ctx context.Context, queue string, body []byte, headers amqp.Table) error { ch, err := Channel(ctx) if err != nil { return err } defer ch.Close() if err = declareQueue(ch, queue); err != nil { return err } return publishWithConfirm(ctx, ch, queue, amqp.Publishing{ ContentType: "application/json", DeliveryMode: amqp.Persistent, Timestamp: time.Now(), Headers: headers, Body: body, }) } func publishWithConfirm(ctx context.Context, ch *amqp.Channel, queue string, message amqp.Publishing) error { if err := ch.Confirm(false); err != nil { return gerror.Wrap(err, "rabbitmq enable publisher confirm failed") } confirms := ch.NotifyPublish(make(chan amqp.Confirmation, 1)) if err := ch.PublishWithContext(ctx, "", queue, false, false, message); err != nil { return gerror.Wrapf(err, "rabbitmq publish failed, queue:%s", queue) } select { case confirmation, ok := <-confirms: if !ok { return gerror.Newf("rabbitmq publisher confirm channel closed, queue:%s", queue) } if !confirmation.Ack { return gerror.Newf("rabbitmq message was not confirmed, queue:%s", queue) } return nil case <-ctx.Done(): return gerror.Wrapf(ctx.Err(), "wait rabbitmq publisher confirm failed, queue:%s", queue) } } // Consume 启动消费者,使用手动 Ack,handler 成功才确认消息。 func Consume(ctx context.Context, queue, consumer string, handler Handler) error { _, err := ConsumeWithDone(ctx, queue, consumer, handler) return err } func ConsumeWithDone(ctx context.Context, queue, consumer string, handler Handler) (<-chan struct{}, error) { ch, err := Channel(ctx) if err != nil { return nil, err } if err = declareQueue(ch, queue); err != nil { _ = ch.Close() return nil, err } if err = ch.Qos(1, 0, false); err != nil { _ = ch.Close() return nil, 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 nil, gerror.Wrap(err, "rabbitmq consume failed") } consumerState := &managedConsumer{ channel: ch, tag: consumer, done: make(chan struct{}), } registerConsumer(consumerState) go func() { defer close(consumerState.done) defer ch.Close() defer unregisterConsumer(consumerState) for { if ctx.Err() != nil { return } select { case <-ctx.Done(): return case delivery, ok := <-deliveries: if !ok { return } consumerTag := delivery.ConsumerTag if consumerTag == "" { consumerTag = consumer } deliveryCtx := WithDeliveryMetadata(ctx, DeliveryMetadata{ DeliveryTag: delivery.DeliveryTag, Redelivered: delivery.Redelivered, Consumer: consumerTag, Queue: queue, Headers: cloneHeaders(delivery.Headers), }) if handlerErr := handler(deliveryCtx, delivery.Body); handlerErr != nil { g.Log().Errorf(deliveryCtx, "rabbitmq message Nack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t, err:%+v", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered, handlerErr) _ = delivery.Nack(false, true) continue } // handler 内部已经完成业务处理或重新发布下一次重试消息,这里只确认当前消息。 g.Log().Infof(deliveryCtx, "rabbitmq message Ack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered) _ = delivery.Ack(false) } } }() return consumerState.done, 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(ctx context.Context) error { closeMu.Lock() defer closeMu.Unlock() consumerMu.Lock() registeredConsumers := consumers consumers = nil consumerMu.Unlock() var closeErrors []error for _, consumer := range registeredConsumers { if consumer.channel.IsClosed() { continue } if err := consumer.channel.Cancel(consumer.tag, false); !isClosedError(err) { closeErrors = append(closeErrors, gerror.Wrapf(err, "cancel RabbitMQ consumer failed, consumer:%s", consumer.tag)) } } if err := waitForConsumers(ctx, registeredConsumers); err != nil { closeErrors = append(closeErrors, err) } for _, consumer := range registeredConsumers { if consumer.channel.IsClosed() { continue } if err := consumer.channel.Close(); !isClosedError(err) { closeErrors = append(closeErrors, gerror.Wrapf(err, "close RabbitMQ channel failed, consumer:%s", consumer.tag)) } } if err := closeConnection(); err != nil { closeErrors = append(closeErrors, err) } return errors.Join(closeErrors...) } // 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 registerConsumer(consumer *managedConsumer) { consumerMu.Lock() defer consumerMu.Unlock() consumers = append(consumers, consumer) } func waitForConsumers(ctx context.Context, registeredConsumers []*managedConsumer) error { for _, consumer := range registeredConsumers { select { case <-consumer.done: case <-ctx.Done(): return gerror.Wrap(ctx.Err(), "wait for RabbitMQ consumers failed") } } return nil } func closeConnection() error { connMu.Lock() defer connMu.Unlock() if conn == nil || conn.IsClosed() { conn = nil return nil } if err := conn.Close(); !isClosedError(err) { return gerror.Wrap(err, "close RabbitMQ connection failed") } conn = nil return nil } func cloneHeaders(headers amqp.Table) map[string]any { if len(headers) == 0 { return nil } cloned := make(map[string]any, len(headers)) for key, value := range headers { cloned[key] = value } return cloned } func isClosedError(err error) bool { return err == nil || errors.Is(err, amqp.ErrClosed) } func unregisterConsumer(target *managedConsumer) { consumerMu.Lock() defer consumerMu.Unlock() for index, consumer := range consumers { if consumer == target { remaining := make([]*managedConsumer, 0, len(consumers)-1) remaining = append(remaining, consumers[:index]...) consumers = append(remaining, consumers[index+1:]...) return } } }