package mq import ( "context" "errors" "fmt" "strings" "sync" "sync/atomic" "time" amqp "github.com/rabbitmq/amqp091-go" ) type RabbitMQConfig struct { Enabled bool Host string Port int Username string Password string VHost string Exchange string ExchangeType string Queue string RoutingKey string Durable bool AutoDelete bool PrefetchCount int Heartbeat time.Duration ConnectTimeout time.Duration ReconnectInterval time.Duration TLS bool } type rabbitSubscription struct { qos byte handler Handler } // RabbitMQ adapts AMQP exchanges and routing keys to the shared topic-based // Client contract. All subscriptions share the configured queue and a single // consumer; deliveries are dispatched to matching handlers locally. type RabbitMQ struct { mu sync.RWMutex opMu sync.Mutex publishMu sync.Mutex consumeMu sync.Mutex connection *amqp.Connection publishChannel *amqp.Channel consumeChannel *amqp.Channel config RabbitMQConfig subscriptions map[string]rabbitSubscription consumerTag string consuming bool stop chan struct{} closed bool reconnecting atomic.Bool } func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) { client := &RabbitMQ{subscriptions: make(map[string]rabbitSubscription), stop: make(chan struct{})} if !config.Enabled { return client, nil } config = defaultRabbitMQConfig(config) if err := validateRabbitMQConfig(config); err != nil { return nil, err } scheme := "amqp" if config.TLS { scheme = "amqps" } address := amqp.URI{ Scheme: scheme, Host: config.Host, Port: config.Port, Username: config.Username, Password: config.Password, Vhost: config.VHost, ConnectionTimeout: int(config.ConnectTimeout.Milliseconds()), }.String() connection, err := amqp.DialConfig(address, amqp.Config{ Heartbeat: config.Heartbeat, Recovery: &amqp.Recovery{ ReconnectionConfig: &amqp.ReconnectionConfig{MaxRetryCount: 5, RetryInterval: config.ReconnectInterval}, OnTopologyEntityError: func(_ *amqp.Connection, entity amqp.TopologyRecoveryEntity) bool { amqp.Logger.Printf("rabbitmq topology recovery failed: %s", entity.Error()) return false }, }, }) if err != nil { return nil, fmt.Errorf("connect rabbitmq: %w", err) } publishChannel, err := connection.Channel() if err != nil { _ = connection.Close() return nil, fmt.Errorf("open rabbitmq publish channel: %w", err) } consumeChannel, err := connection.Channel() if err != nil { _ = publishChannel.Close() _ = connection.Close() return nil, fmt.Errorf("open rabbitmq consume channel: %w", err) } closeOnError := func() { _ = consumeChannel.Close() _ = publishChannel.Close() _ = connection.Close() } if err = consumeChannel.ExchangeDeclare(config.Exchange, config.ExchangeType, config.Durable, config.AutoDelete, false, false, nil); err != nil { closeOnError() return nil, fmt.Errorf("declare rabbitmq exchange: %w", err) } if _, err = consumeChannel.QueueDeclare(config.Queue, config.Durable, config.AutoDelete, false, false, nil); err != nil { closeOnError() return nil, fmt.Errorf("declare rabbitmq queue: %w", err) } if config.PrefetchCount > 0 { if err = consumeChannel.Qos(config.PrefetchCount, 0, false); err != nil { closeOnError() return nil, fmt.Errorf("configure rabbitmq qos: %w", err) } } client.connection = connection client.publishChannel = publishChannel client.consumeChannel = consumeChannel client.config = config client.consumerTag = fmt.Sprintf("kra-%d", time.Now().UnixNano()) stateChanges := make(chan *amqp.StateChanged, 16) connection.NotifyStateChange(stateChanges) go client.watchState(stateChanges) return client, nil } func defaultRabbitMQConfig(config RabbitMQConfig) RabbitMQConfig { if config.Port <= 0 { if config.TLS { config.Port = 5671 } else { config.Port = 5672 } } if config.Username == "" { config.Username = "guest" } if config.Password == "" { config.Password = "guest" } if config.VHost == "" { config.VHost = "/" } if config.ExchangeType == "" { config.ExchangeType = "topic" } if config.RoutingKey == "" { config.RoutingKey = "#" } if config.Heartbeat <= 0 { config.Heartbeat = 10 * time.Second } if config.ConnectTimeout <= 0 { config.ConnectTimeout = 10 * time.Second } if config.ReconnectInterval <= 0 { config.ReconnectInterval = 5 * time.Second } return config } func validateRabbitMQConfig(config RabbitMQConfig) error { if strings.TrimSpace(config.Host) == "" { return errors.New("rabbitmq host is empty") } if config.Port < 1 || config.Port > 65535 { return fmt.Errorf("invalid rabbitmq port %d", config.Port) } if strings.TrimSpace(config.Exchange) == "" { return errors.New("rabbitmq exchange is empty") } if strings.TrimSpace(config.Queue) == "" { return errors.New("rabbitmq queue is empty") } switch strings.ToLower(config.ExchangeType) { case "direct", "fanout", "topic": return nil default: return fmt.Errorf("invalid rabbitmq exchange type %q", config.ExchangeType) } } func (c *RabbitMQ) Publish(ctx context.Context, topic string, payload []byte, qos byte, _ bool) error { if c == nil { return ErrUnavailable } if ctx == nil { ctx = context.Background() } if c != nil && strings.TrimSpace(topic) == "" { topic = c.config.RoutingKey } if strings.TrimSpace(topic) == "" { return errors.New("rabbitmq routing key is empty") } if qos > AtLeastOnce { return fmt.Errorf("rabbitmq supports qos 0 or 1, got %d", qos) } select { case <-ctx.Done(): return ctx.Err() default: } c.publishMu.Lock() defer c.publishMu.Unlock() c.mu.RLock() if c.closed || c.publishChannel == nil || c.publishChannel.IsClosed() { c.mu.RUnlock() return ErrUnavailable } channel := c.publishChannel exchange := c.config.Exchange c.mu.RUnlock() deliveryMode := amqp.Transient if qos >= AtLeastOnce { deliveryMode = amqp.Persistent } return channel.PublishWithContext(ctx, exchange, topic, false, false, amqp.Publishing{ ContentType: "application/octet-stream", DeliveryMode: deliveryMode, Timestamp: time.Now(), Body: append([]byte(nil), payload...), }) } func (c *RabbitMQ) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error { if c == nil { return ErrUnavailable } if ctx == nil { ctx = context.Background() } if c != nil && strings.TrimSpace(topic) == "" { topic = c.config.RoutingKey } if strings.TrimSpace(topic) == "" { return errors.New("rabbitmq routing key is empty") } if qos > AtLeastOnce { return fmt.Errorf("rabbitmq supports qos 0 or 1, got %d", qos) } if handler == nil { return errors.New("rabbitmq handler is nil") } select { case <-ctx.Done(): return ctx.Err() default: } c.opMu.Lock() defer c.opMu.Unlock() c.consumeMu.Lock() defer c.consumeMu.Unlock() c.mu.RLock() if c.closed || c.consumeChannel == nil || c.consumeChannel.IsClosed() { c.mu.RUnlock() return ErrUnavailable } channel := c.consumeChannel config := c.config _, exists := c.subscriptions[topic] c.mu.RUnlock() if !exists { if err := channel.QueueBind(config.Queue, topic, config.Exchange, false, nil); err != nil { return fmt.Errorf("bind rabbitmq queue: %w", err) } } c.mu.Lock() c.subscriptions[topic] = rabbitSubscription{qos: qos, handler: handler} shouldStart := !c.consuming c.mu.Unlock() if !shouldStart { return nil } deliveries, err := channel.Consume(config.Queue, c.consumerTag, false, false, false, false, nil) if err != nil { c.mu.Lock() delete(c.subscriptions, topic) c.mu.Unlock() if !exists { _ = channel.QueueUnbind(config.Queue, topic, config.Exchange, nil) } return fmt.Errorf("consume rabbitmq queue: %w", err) } c.mu.Lock() c.consuming = true c.mu.Unlock() go c.consume(deliveries) return nil } func (c *RabbitMQ) consume(deliveries <-chan amqp.Delivery) { for { select { case <-c.stop: return case delivery, ok := <-deliveries: if !ok { c.mu.Lock() c.consuming = false c.mu.Unlock() return } subscriptions := c.subscriptionSnapshot() for pattern, item := range subscriptions { if rabbitRoutingKeyMatches(c.config.ExchangeType, pattern, delivery.RoutingKey) { item.handler(context.Background(), Message{Topic: delivery.RoutingKey, Payload: append([]byte(nil), delivery.Body...), QoS: item.qos}) } } c.consumeMu.Lock() _ = delivery.Ack(false) c.consumeMu.Unlock() } } } func (c *RabbitMQ) subscriptionSnapshot() map[string]rabbitSubscription { c.mu.RLock() defer c.mu.RUnlock() result := make(map[string]rabbitSubscription, len(c.subscriptions)) for topic, item := range c.subscriptions { result[topic] = item } return result } func (c *RabbitMQ) Unsubscribe(ctx context.Context, topics ...string) error { if c == nil { return ErrUnavailable } if ctx == nil { ctx = context.Background() } if len(topics) == 0 { return errors.New("rabbitmq routing keys are empty") } select { case <-ctx.Done(): return ctx.Err() default: } c.opMu.Lock() defer c.opMu.Unlock() c.consumeMu.Lock() defer c.consumeMu.Unlock() c.mu.RLock() if c.closed || c.consumeChannel == nil || c.consumeChannel.IsClosed() { c.mu.RUnlock() return ErrUnavailable } channel := c.consumeChannel config := c.config c.mu.RUnlock() for _, topic := range topics { c.mu.RLock() _, exists := c.subscriptions[topic] c.mu.RUnlock() if !exists { continue } if err := channel.QueueUnbind(config.Queue, topic, config.Exchange, nil); err != nil { return fmt.Errorf("unbind rabbitmq queue: %w", err) } c.mu.Lock() delete(c.subscriptions, topic) c.mu.Unlock() } c.mu.RLock() empty := len(c.subscriptions) == 0 consuming := c.consuming c.mu.RUnlock() if empty && consuming { if err := channel.Cancel(c.consumerTag, false); err != nil { return fmt.Errorf("cancel rabbitmq consumer: %w", err) } c.mu.Lock() c.consuming = false c.mu.Unlock() } return nil } func (c *RabbitMQ) Connected() bool { if c == nil { return false } c.mu.RLock() defer c.mu.RUnlock() return !c.closed && !c.reconnecting.Load() && c.connection != nil && !c.connection.IsClosed() } // Reconnecting reports whether the AMQP driver is actively recovering the // current connection. The outer runtime waits for this state to settle before // deciding whether a fresh client must be built. func (c *RabbitMQ) Reconnecting() bool { return c != nil && !c.closed && c.reconnecting.Load() } func (c *RabbitMQ) watchState(states <-chan *amqp.StateChanged) { for state := range states { if state == nil { continue } switch state.To { case amqp.StateReconnecting: c.reconnecting.Store(true) case amqp.StateOpen, amqp.StateClosed: c.reconnecting.Store(false) } } } func (c *RabbitMQ) Close() error { if c == nil { return nil } c.opMu.Lock() defer c.opMu.Unlock() c.publishMu.Lock() defer c.publishMu.Unlock() c.consumeMu.Lock() defer c.consumeMu.Unlock() c.mu.Lock() if c.closed { c.mu.Unlock() return nil } c.closed = true close(c.stop) publishChannel := c.publishChannel consumeChannel := c.consumeChannel connection := c.connection c.publishChannel = nil c.consumeChannel = nil c.connection = nil c.mu.Unlock() var result error if consumeChannel != nil { result = errors.Join(result, consumeChannel.Close()) } if publishChannel != nil { result = errors.Join(result, publishChannel.Close()) } if connection != nil { result = errors.Join(result, connection.Close()) } return result } func rabbitRoutingKeyMatches(exchangeType, pattern, routingKey string) bool { switch strings.ToLower(exchangeType) { case "fanout": return true case "direct": return pattern == routingKey } patternParts := strings.Split(pattern, ".") routingParts := strings.Split(routingKey, ".") for len(patternParts) > 0 { head := patternParts[0] patternParts = patternParts[1:] if head == "#" { if len(patternParts) == 0 { return true } for index := 0; index <= len(routingParts); index++ { if rabbitTopicPartsMatch(patternParts, routingParts[index:]) { return true } } return false } if len(routingParts) == 0 || (head != "*" && head != routingParts[0]) { return false } routingParts = routingParts[1:] } return len(routingParts) == 0 } func rabbitTopicPartsMatch(patternParts, routingParts []string) bool { return rabbitRoutingKeyMatches("topic", strings.Join(patternParts, "."), strings.Join(routingParts, ".")) }