kra-oa/pkg/mq/rabbitmq.go

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 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 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 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, "."))
}