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) } func (c *Client) PublishDesktopAccountHealth(ctx context.Context, body []byte) error { return c.publish(ctx, c.cfg.DesktopAccountHealthExchange, c.cfg.DesktopAccountHealthRoutingKey, 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) ConsumeDesktopAccountHealth(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { return c.consume(consumerName, c.cfg.DesktopAccountHealthQueue) } 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.DesktopAccountHealthExchange, c.cfg.DesktopAccountHealthRoutingKey, c.cfg.DesktopAccountHealthQueue, c.cfg.DesktopAccountHealthDLX, c.cfg.DesktopAccountHealthDLQ, c.cfg.DesktopAccountHealthDLQRouteKey, ); err != nil { return 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 }