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
	publishPool chan *publishChannel
}

type publishChannel struct {
	conn *amqp.Connection
	ch   *amqp.Channel
}

type QueueStats struct {
	Name      string
	Messages  int
	Consumers int
}

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,
		publishPool: make(chan *publishChannel, cfg.PublishChannelPoolSize),
	}
	if _, err := client.ensureConnection(); err != nil {
		return nil, err
	}

	return client, nil
}

// Config returns a copy of the effective rabbitmq configuration. Services that
// need to derive routing keys (e.g. desktop dispatch) consult this instead of
// hard-coding string constants.
func (c *Client) Config() config.RabbitMQConfig {
	if c == nil {
		return config.RabbitMQConfig{}
	}
	return c.cfg
}

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

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

	for _, ch := range pooledChannels {
		_ = ch.Close()
	}

	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) InspectGenerationQueue(ctx context.Context) (QueueStats, error) {
	return c.inspectQueue(ctx, c.cfg.GenerationQueue)
}

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) PublishGenerationTask(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.GenerationExchange, c.cfg.GenerationRoutingKey, body)
}

func (c *Client) PublishGenerationEvent(ctx context.Context, body []byte) error {
	return c.publishFanout(ctx, c.cfg.GenerationEventExchange, body)
}

func (c *Client) PublishDesktopTaskEvent(ctx context.Context, body []byte) error {
	return c.publishFanout(ctx, c.cfg.DesktopTaskEventExchange, body)
}

func (c *Client) PublishDesktopDispatch(ctx context.Context, routingKey string, body []byte) error {
	return c.publish(ctx, c.cfg.DesktopDispatchExchange, routingKey, body)
}

// ConsumeDesktopDispatch declares a transient per-process queue bound to the
// desktop dispatch topic exchange using bindingKeys, and begins consuming from
// it. Every tenant-api instance gets its own queue so every instance sees every
// dispatch message; the hub inside the instance then filters by target client
// presence. Messages are auto-ack (transient delivery is OK: if the instance
// dies, the client reconnects and the server re-drives the task via the
// existing lease/recovery path).
func (c *Client) ConsumeDesktopDispatch(consumerName string, bindingKeys []string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	if len(bindingKeys) == 0 {
		return nil, nil, fmt.Errorf("consume rabbitmq desktop dispatch: bindingKeys is required")
	}

	for attempt := 0; attempt < 2; attempt++ {
		conn, ch, err := c.openChannel()
		if err != nil {
			return nil, nil, err
		}

		queue, err := ch.QueueDeclare(
			"",
			false,
			true,
			true,
			false,
			nil,
		)
		if err != nil {
			_ = ch.Close()
			c.invalidateConnection(conn)
			if attempt == 1 {
				return nil, nil, fmt.Errorf("declare rabbitmq desktop dispatch queue: %w", err)
			}
			continue
		}

		bindFailed := false
		for _, key := range bindingKeys {
			if err := ch.QueueBind(
				queue.Name,
				key,
				c.cfg.DesktopDispatchExchange,
				false,
				nil,
			); err != nil {
				_ = ch.Close()
				c.invalidateConnection(conn)
				if attempt == 1 {
					return nil, nil, fmt.Errorf("bind rabbitmq desktop dispatch queue: %w", err)
				}
				bindFailed = true
				break
			}
		}
		if bindFailed {
			continue
		}

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

		return deliveries, ch, nil
	}

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

func (c *Client) PublishTemplateAssistTask(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.TemplateAssistExchange, c.cfg.TemplateAssistRoutingKey, body)
}

func (c *Client) PublishKolAssistTask(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.KolAssistExchange, c.cfg.KolAssistRoutingKey, body)
}

func (c *Client) PublishImageAssetTask(ctx context.Context, body []byte) error {
	return c.publish(ctx, c.cfg.ImageAssetExchange, c.cfg.ImageAssetRoutingKey, body)
}

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

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

		c.discardPublishChannel(ch)
		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) publishFanout(ctx context.Context, exchange string, body []byte) error {
	for attempt := 0; attempt < 2; attempt++ {
		conn, ch, pooled, err := c.acquirePublishChannel()
		if err != nil {
			return err
		}

		publishErr := ch.PublishWithContext(ctx, exchange, "", false, false, amqp.Publishing{
			ContentType:  "application/json",
			DeliveryMode: amqp.Transient,
			Body:         body,
		})
		if publishErr == nil {
			c.releasePublishChannel(conn, ch, pooled)
			return nil
		}

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

	return fmt.Errorf("publish rabbitmq fanout 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) ConsumeGenerationTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	return c.consume(consumerName, c.cfg.GenerationQueue)
}

func (c *Client) ConsumeGenerationEvents(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	return c.consumeFanout(consumerName, c.cfg.GenerationEventExchange)
}

func (c *Client) ConsumeDesktopTaskEvents(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
	return c.consumeFanout(consumerName, c.cfg.DesktopTaskEventExchange)
}

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

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

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

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) consumeFanout(consumerName, exchange 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
		}

		queue, err := ch.QueueDeclare(
			"",
			false,
			true,
			true,
			false,
			nil,
		)
		if err != nil {
			_ = ch.Close()
			c.invalidateConnection(conn)
			if attempt == 1 {
				return nil, nil, fmt.Errorf("declare rabbitmq fanout queue: %w", err)
			}
			continue
		}

		if err := ch.QueueBind(
			queue.Name,
			"",
			exchange,
			false,
			nil,
		); err != nil {
			_ = ch.Close()
			c.invalidateConnection(conn)
			if attempt == 1 {
				return nil, nil, fmt.Errorf("bind rabbitmq fanout queue: %w", err)
			}
			continue
		}

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

		return deliveries, ch, nil
	}

	return nil, nil, fmt.Errorf("consume rabbitmq fanout 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 {
		for _, ch := range c.drainPublishChannelsLocked() {
			_ = ch.Close()
		}
		_ = 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) acquirePublishChannel() (*amqp.Connection, *amqp.Channel, bool, error) {
	for attempt := 0; attempt < 2; attempt++ {
		conn, err := c.ensureConnection()
		if err != nil {
			return nil, nil, false, err
		}

		if ch := c.borrowPublishChannel(conn); ch != nil {
			return conn, ch, true, nil
		}

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

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

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

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

	c.mu.Lock()
	if c.conn != conn {
		c.mu.Unlock()
		return
	}
	c.conn = nil
	pooledChannels := c.drainPublishChannelsLocked()
	c.mu.Unlock()

	for _, ch := range pooledChannels {
		_ = ch.Close()
	}
	if !conn.IsClosed() {
		_ = conn.Close()
	}
}

func (c *Client) inspectQueue(ctx context.Context, queue string) (QueueStats, error) {
	if err := ctx.Err(); err != nil {
		return QueueStats{}, err
	}

	for attempt := 0; attempt < 2; attempt++ {
		conn, ch, err := c.openChannel()
		if err != nil {
			return QueueStats{}, err
		}

		info, inspectErr := ch.QueueInspect(queue)
		_ = ch.Close()
		if inspectErr == nil {
			return QueueStats{
				Name:      info.Name,
				Messages:  info.Messages,
				Consumers: info.Consumers,
			}, nil
		}

		c.invalidateConnection(conn)
		if attempt == 1 {
			return QueueStats{}, fmt.Errorf("inspect rabbitmq queue %s: %w", queue, inspectErr)
		}
	}

	return QueueStats{}, fmt.Errorf("inspect rabbitmq queue %s: retry exhausted", queue)
}

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

	for {
		select {
		case pooled := <-c.publishPool:
			if pooled == nil || pooled.ch == nil {
				continue
			}
			if pooled.conn != conn {
				_ = pooled.ch.Close()
				continue
			}
			return pooled.ch
		default:
			return nil
		}
	}
}

func (c *Client) releasePublishChannel(conn *amqp.Connection, ch *amqp.Channel, reusable bool) {
	if ch == nil {
		return
	}
	if !reusable || c == nil || c.publishPool == nil || conn == nil || conn.IsClosed() {
		_ = ch.Close()
		return
	}

	c.mu.Lock()
	currentConn := c.conn
	closed := c.closed
	c.mu.Unlock()

	if closed || currentConn != conn {
		_ = ch.Close()
		return
	}

	select {
	case c.publishPool <- &publishChannel{conn: conn, ch: ch}:
	default:
		_ = ch.Close()
	}
}

func (c *Client) discardPublishChannel(ch *amqp.Channel) {
	if ch == nil {
		return
	}
	_ = ch.Close()
}

func (c *Client) drainPublishChannelsLocked() []*amqp.Channel {
	if c == nil || c.publishPool == nil {
		return nil
	}

	drained := make([]*amqp.Channel, 0, len(c.publishPool))
	for {
		select {
		case pooled := <-c.publishPool:
			if pooled != nil && pooled.ch != nil {
				drained = append(drained, pooled.ch)
			}
		default:
			return drained
		}
	}
}

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
	}

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.GenerationExchange,
		c.cfg.GenerationRoutingKey,
		c.cfg.GenerationQueue,
		c.cfg.GenerationDLX,
		c.cfg.GenerationDLQ,
		c.cfg.GenerationDLQRouteKey,
	); err != nil {
		return err
	}

	if err := ch.ExchangeDeclare(
		c.cfg.GenerationEventExchange,
		"fanout",
		false,
		false,
		false,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.GenerationEventExchange, err)
	}

	if err := ch.ExchangeDeclare(
		c.cfg.DesktopTaskEventExchange,
		"fanout",
		false,
		false,
		false,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.DesktopTaskEventExchange, err)
	}

	if err := ch.ExchangeDeclare(
		c.cfg.DesktopDispatchExchange,
		"topic",
		true,
		false,
		false,
		false,
		nil,
	); err != nil {
		return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.DesktopDispatchExchange, err)
	}

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.TemplateAssistExchange,
		c.cfg.TemplateAssistRoutingKey,
		c.cfg.TemplateAssistQueue,
		c.cfg.TemplateAssistDLX,
		c.cfg.TemplateAssistDLQ,
		c.cfg.TemplateAssistDLQRouteKey,
	); err != nil {
		return err
	}

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.KolAssistExchange,
		c.cfg.KolAssistRoutingKey,
		c.cfg.KolAssistQueue,
		c.cfg.KolAssistDLX,
		c.cfg.KolAssistDLQ,
		c.cfg.KolAssistDLQRouteKey,
	); err != nil {
		return err
	}

	if err := c.ensureWorkQueue(
		ch,
		c.cfg.ImageAssetExchange,
		c.cfg.ImageAssetRoutingKey,
		c.cfg.ImageAssetQueue,
		c.cfg.ImageAssetDLX,
		c.cfg.ImageAssetDLQ,
		c.cfg.ImageAssetDLQRouteKey,
	); 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
}