geo/server/internal/shared/messaging/rabbitmq/client.go

package rabbitmq

import (
	"context"
	"fmt"
	"strings"
	"sync"

	amqp "github.com/rabbitmq/amqp091-go"

	"github.com/geo-platform/tenant-api/internal/shared/config"
)

type Client struct {
	cfg    config.RabbitMQConfig
	mu     sync.Mutex
	conn   *amqp.Connection
	closed bool
}

func New(_ context.Context, cfg config.RabbitMQConfig) (*Client, error) {
	if strings.TrimSpace(cfg.URL) == "" {
		return nil, fmt.Errorf("rabbitmq url is required")
	}

	client := &Client{cfg: cfg}
	if _, err := client.ensureConnection(); err != nil {
		return nil, err
	}

	return client, nil
}

func (c *Client) Close() error {
	if c == nil {
		return nil
	}

	c.mu.Lock()
	c.closed = true
	conn := c.conn
	c.conn = nil
	c.mu.Unlock()

	if conn == nil || conn.IsClosed() {
		return nil
	}
	return conn.Close()
}

func (c *Client) Ping(_ context.Context) error {
	_, ch, err := c.openChannel()
	if err != nil {
		return err
	}
	defer ch.Close()

	return nil
}

func (c *Client) PublishMonitorResultIngest(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.MonitorResultExchange, c.cfg.MonitorResultRoutingKey, body)
}

func (c *Client) PublishMonitorProjectionRebuild(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.MonitorProjectionExchange, c.cfg.MonitorProjectionRoutingKey, body)
}

func (c *Client) publish(ctx context.Context, exchange, routingKey string, body []byte) error {
	for attempt := 0; attempt < 2; attempt++ {
		conn, ch, err := c.openChannel()
		if err != nil {
			return err
		}

		publishErr := ch.PublishWithContext(ctx, exchange, routingKey, false, false, amqp.Publishing{
			ContentType:  "application/json",
			DeliveryMode: amqp.Persistent,
			Body:         body,
		})
		_ = ch.Close()
		if publishErr == nil {
			return nil
		}

		c.invalidateConnection(conn)
		if attempt == 1 {
			return fmt.Errorf("publish rabbitmq message: %w", publishErr)
		}
	}

	return fmt.Errorf("publish rabbitmq message: retry exhausted")
}

func (c *Client) ConsumeMonitorResultIngest(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	return c.consume(consumerName, c.cfg.MonitorResultQueue)
}

func (c *Client) ConsumeMonitorProjectionRebuild(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	return c.consume(consumerName, c.cfg.MonitorProjectionQueue)
}

func (c *Client) consume(consumerName, queue string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	for attempt := 0; attempt < 2; attempt++ {
		conn, ch, err := c.openChannel()
		if err != nil {
			return nil, nil, err
		}

		if c.cfg.ConsumerPrefetch > 0 {
			if err := ch.Qos(c.cfg.ConsumerPrefetch, 0, false); err != nil {
				_ = ch.Close()
				c.invalidateConnection(conn)
				if attempt == 1 {
					return nil, nil, fmt.Errorf("configure rabbitmq qos: %w", err)
				}
				continue
			}
		}

		deliveries, err := ch.Consume(
			queue,
			consumerName,
			false,
			false,
			false,
			false,
			nil,
		)
		if err != nil {
			_ = ch.Close()
			c.invalidateConnection(conn)
			if attempt == 1 {
				return nil, nil, fmt.Errorf("consume rabbitmq queue: %w", err)
			}
			continue
		}

		return deliveries, ch, nil
	}

	return nil, nil, fmt.Errorf("consume rabbitmq queue: retry exhausted")
}

func (c *Client) ensureConnection() (*amqp.Connection, error) {
	if c == nil {
		return nil, fmt.Errorf("rabbitmq client is nil")
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return nil, fmt.Errorf("rabbitmq client is closed")
	}
	if c.conn != nil && !c.conn.IsClosed() {
		return c.conn, nil
	}
	if c.conn != nil {
		_ = c.conn.Close()
		c.conn = nil
	}

	conn, err := amqp.Dial(c.cfg.URL)
	if err != nil {
		return nil, fmt.Errorf("dial rabbitmq: %w", err)
	}
	if err := c.ensureTopology(conn); err != nil {
		_ = conn.Close()
		return nil, err
	}

	c.conn = conn
	return conn, nil
}

func (c *Client) openChannel() (*amqp.Connection, *amqp.Channel, error) {
	for attempt := 0; attempt < 2; attempt++ {
		conn, err := c.ensureConnection()
		if err != nil {
			return nil, nil, err
		}

		ch, err := conn.Channel()
		if err == nil {
			return conn, ch, nil
		}

		c.invalidateConnection(conn)
		if attempt == 1 {
			return nil, nil, fmt.Errorf("open rabbitmq channel: %w", err)
		}
	}

	return nil, nil, fmt.Errorf("open rabbitmq channel: retry exhausted")
}

func (c *Client) invalidateConnection(conn *amqp.Connection) {
	if c == nil || conn == nil {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != conn {
		return
	}
	c.conn = nil
	if !conn.IsClosed() {
		_ = conn.Close()
	}
}

func (c *Client) ensureTopology(conn *amqp.Connection) error {
	ch, err := conn.Channel()
	if err != nil {
		return fmt.Errorf("open rabbitmq channel: %w", err)
	}
	defer ch.Close()

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.MonitorResultExchange,
		c.cfg.MonitorResultRoutingKey,
		c.cfg.MonitorResultQueue,
		c.cfg.MonitorResultDLX,
		c.cfg.MonitorResultDLQ,
		c.cfg.MonitorResultDLQRouteKey,
	); err != nil {
		return err
	}

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.MonitorProjectionExchange,
		c.cfg.MonitorProjectionRoutingKey,
		c.cfg.MonitorProjectionQueue,
		c.cfg.MonitorProjectionDLX,
		c.cfg.MonitorProjectionDLQ,
		c.cfg.MonitorProjectionDLQRouteKey,
	); err != nil {
		return err
	}

	return nil
}

func (c *Client) ensureWorkQueue(ch *amqp.Channel, exchange, routingKey, queue, dlx, dlq, dlqRoutingKey string) error {
	if err := ch.ExchangeDeclare(
		exchange,
		"direct",
		true,
		false,
		false,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("declare rabbitmq exchange %s: %w", exchange, err)
	}

	if err := ch.ExchangeDeclare(
		dlx,
		"direct",
		true,
		false,
		false,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("declare rabbitmq dlx %s: %w", dlx, err)
	}

	queueArgs := amqp.Table{
		"x-queue-type":              "quorum",
		"x-dead-letter-exchange":    dlx,
		"x-dead-letter-routing-key": dlqRoutingKey,
	}
	if _, err := ch.QueueDeclare(
		queue,
		true,
		false,
		false,
		false,
		queueArgs,
	); err != nil {
		return fmt.Errorf("declare rabbitmq queue %s: %w", queue, err)
	}

	if err := ch.QueueBind(
		queue,
		routingKey,
		exchange,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("bind rabbitmq queue %s: %w", queue, err)
	}

	dlqArgs := amqp.Table{
		"x-queue-type": "quorum",
	}
	if _, err := ch.QueueDeclare(
		dlq,
		true,
		false,
		false,
		false,
		dlqArgs,
	); err != nil {
		return fmt.Errorf("declare rabbitmq dlq %s: %w", dlq, err)
	}

	if err := ch.QueueBind(
		dlq,
		dlqRoutingKey,
		dlx,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("bind rabbitmq dlq %s: %w", dlq, err)
	}

	return nil
}