aiaa-notification-server/internal/subscriber/subscriber.go

package subscriber

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"aiaa-notification-service/internal/config"
	"aiaa-notification-service/internal/subscriber/tradesignal"

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

type Subscriber struct {
	cfg     config.SubscriptionConfig
	conv    *tradesignal.Converter
	lookup  SourceLookup
	process ProcessFunc
	deduper Deduper
}

func New(cfg config.SubscriptionConfig, lookup SourceLookup, process ProcessFunc, deduper Deduper) (*Subscriber, error) {
	if cfg.Formatter != "trade_signal" {
		return nil, fmt.Errorf("unknown formatter %q", cfg.Formatter)
	}
	return &Subscriber{
		cfg:     cfg,
		conv:    tradesignal.NewConverter(cfg.StrategyOverrides),
		lookup:  lookup,
		process: process,
		deduper: deduper,
	}, nil
}

func (s *Subscriber) Run(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		if err := s.consumeOnce(ctx); err != nil {
			slog.Error("subscriber error, reconnecting", "name", s.cfg.Name, "error", err)
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(5 * time.Second):
			}
		}
	}
}

func (s *Subscriber) consumeOnce(ctx context.Context) error {
	conn, err := amqp.Dial(s.cfg.URL)
	if err != nil {
		return fmt.Errorf("dial rabbitmq: %w", err)
	}
	defer conn.Close()

	ch, err := conn.Channel()
	if err != nil {
		return fmt.Errorf("open channel: %w", err)
	}
	defer ch.Close()

	if err := ch.Qos(1, 0, false); err != nil {
		return fmt.Errorf("set qos: %w", err)
	}
	if err := s.ensureQueue(ch); err != nil {
		return err
	}

	deliveries, err := ch.Consume(s.cfg.Queue, s.cfg.Name, false, false, false, false, nil)
	if err != nil {
		return fmt.Errorf("consume queue %s: %w", s.cfg.Queue, err)
	}
	slog.Info("listening on queue", "name", s.cfg.Name, "queue", s.cfg.Queue)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case d, ok := <-deliveries:
			if !ok {
				return fmt.Errorf("delivery channel closed")
			}
			s.handleDelivery(ch, d)
		}
	}
}

func (s *Subscriber) ensureQueue(ch *amqp.Channel) error {
	if s.cfg.Exchange != "" {
		if err := ch.ExchangeDeclare(s.cfg.Exchange, s.cfg.ExchangeType, true, false, false, false, nil); err != nil {
			return fmt.Errorf("declare exchange %q: %w", s.cfg.Exchange, err)
		}
	}

	if _, err := ch.QueueDeclare(s.cfg.Queue, true, false, false, false, nil); err != nil {
		return fmt.Errorf("declare queue %q: %w", s.cfg.Queue, err)
	}

	if s.cfg.Exchange != "" {
		if err := ch.QueueBind(s.cfg.Queue, s.cfg.RoutingKey, s.cfg.Exchange, false, nil); err != nil {
			return fmt.Errorf("bind queue %q to exchange %q: %w", s.cfg.Queue, s.cfg.Exchange, err)
		}
		slog.Info("queue bound", "queue", s.cfg.Queue, "exchange", s.cfg.Exchange, "type", s.cfg.ExchangeType, "routing_key", s.cfg.RoutingKey)
	}

	if s.cfg.DeadLetterQueue != "" {
		if _, err := ch.QueueDeclare(s.cfg.DeadLetterQueue, true, false, false, false, nil); err != nil {
			slog.Warn("declare dead letter queue failed", "queue", s.cfg.DeadLetterQueue, "error", err)
		}
	}
	return nil
}

func (s *Subscriber) handleDelivery(ch *amqp.Channel, d amqp.Delivery) {
	disp := HandleMessage(context.Background(), HandleInput{
		Body:       d.Body,
		Headers:    map[string]any(d.Headers),
		Name:       s.cfg.Name,
		Queue:      s.cfg.Queue,
		SourceName: s.cfg.Source,
		MaxRetry:   s.cfg.MaxRetry,
		Deduper:    s.deduper,
	}, s.conv, s.lookup, s.process)

	switch disp {
	case DispositionRetry:
		s.republish(ch, d, s.cfg.Queue)
	case DispositionDLQ:
		if s.cfg.DeadLetterQueue == "" {
			slog.Warn("max retry reached, discarded", "name", s.cfg.Name)
			_ = d.Ack(false)
			return
		}
		s.republish(ch, d, s.cfg.DeadLetterQueue)
	default:
		_ = d.Ack(false)
	}
}

func (s *Subscriber) republish(ch *amqp.Channel, d amqp.Delivery, queue string) {
	headers := copyAMQPHeaders(d.Headers)
	headers[retryHeader] = RetryCount(map[string]any(d.Headers)) + 1
	if err := publishToQueue(ch, queue, d.Body, headers); err != nil {
		slog.Error("requeue failed", "queue", queue, "error", err)
		_ = d.Nack(false, true)
		return
	}
	_ = d.Ack(false)
	if queue == s.cfg.DeadLetterQueue {
		slog.Warn("message moved to dlq", "name", s.cfg.Name, "retries", s.cfg.MaxRetry)
		return
	}
	slog.Info("message requeued", "name", s.cfg.Name, "retry", headers[retryHeader], "max", s.cfg.MaxRetry)
}

func copyAMQPHeaders(headers amqp.Table) amqp.Table {
	out := amqp.Table{}
	for k, v := range headers {
		out[k] = v
	}
	return out
}

func publishToQueue(ch *amqp.Channel, queue string, body []byte, headers amqp.Table) error {
	return ch.Publish(
		"",
		queue,
		false,
		false,
		amqp.Publishing{
			ContentType:  "application/json",
			DeliveryMode: amqp.Persistent,
			Headers:      headers,
			Body:         body,
			Timestamp:    time.Now(),
		},
	)
}