geo/server/internal/tenant/app/image_asset_worker.go

package app

import (
	"context"
	"errors"
	"fmt"
	"os"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	"github.com/geo-platform/tenant-api/internal/shared/response"
)

const (
	defaultImageAssetWorkerRetryInterval  = 5 * time.Second
	defaultImageAssetWorkerProcessTimeout = 2 * time.Minute
)

var errImageAssetConsumerClosed = errors.New("image asset consumer channel closed")

type ImageAssetWorker struct {
	rabbitMQ       *rabbitmq.Client
	service        *ImageService
	logger         *zap.Logger
	retryInterval  time.Duration
	processTimeout time.Duration
	consumerPrefix string
	concurrency    int
}

func NewImageAssetWorker(service *ImageService, rabbitMQClient *rabbitmq.Client, logger *zap.Logger) *ImageAssetWorker {
	return &ImageAssetWorker{
		rabbitMQ:       rabbitMQClient,
		service:        service,
		logger:         logger,
		retryInterval:  defaultImageAssetWorkerRetryInterval,
		processTimeout: defaultImageAssetWorkerProcessTimeout,
		consumerPrefix: fmt.Sprintf("tenant-api-image-asset-%d", os.Getpid()),
		concurrency:    1,
	}
}

func (w *ImageAssetWorker) Start(ctx context.Context) {
	if w == nil {
		return
	}
	for i := 0; i < w.concurrency; i++ {
		consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1)
		go w.run(ctx, consumerName)
	}
}

func (w *ImageAssetWorker) run(ctx context.Context, consumerName string) {
	if w.rabbitMQ == nil || w.service == nil {
		return
	}

	for {
		if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil {
			if w.logger != nil {
				if errors.Is(err, errImageAssetConsumerClosed) {
					w.logger.Info("image asset consumer channel closed, retrying",
						zap.String("consumer", consumerName),
					)
				} else {
					w.logger.Warn("image asset worker stopped unexpectedly",
						zap.Error(err),
						zap.String("consumer", consumerName),
					)
				}
			}
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(w.retryInterval):
		}
	}
}

func (w *ImageAssetWorker) consumeOnce(ctx context.Context, consumerName string) error {
	deliveries, ch, err := w.rabbitMQ.ConsumeImageAssetTask(consumerName)
	if err != nil {
		return err
	}
	defer ch.Close()

	for {
		select {
		case <-ctx.Done():
			return nil
		case delivery, ok := <-deliveries:
			if !ok {
				return errImageAssetConsumerClosed
			}
			w.handleDelivery(ctx, delivery)
		}
	}
}

func (w *ImageAssetWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) {
	envelope, err := decodeImageAssetTaskEnvelope(delivery.Body)
	if err != nil {
		w.rejectDelivery(delivery, false, err, 0, 0, "")
		return
	}

	ctx, cancel := context.WithTimeout(parent, w.processTimeout)
	defer cancel()

	switch envelope.Kind {
	case imageAssetTaskKindFinalizeUpload:
		err = w.service.ProcessFinalizeUpload(ctx, envelope.TenantID, envelope.AssetID)
	case imageAssetTaskKindFinalizeDelete:
		err = w.service.ProcessFinalizeDelete(ctx, envelope.TenantID, envelope.AssetID)
	default:
		err = fmt.Errorf("unsupported image asset task kind %q", envelope.Kind)
	}
	if err != nil {
		requeue := shouldRequeueImageAssetError(err, delivery.Redelivered)
		w.rejectDelivery(delivery, requeue, err, envelope.TenantID, envelope.AssetID, envelope.Kind)
		return
	}

	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("image asset ack failed",
			zap.Error(err),
			zap.Int64("tenant_id", envelope.TenantID),
			zap.Int64("asset_id", envelope.AssetID),
			zap.String("kind", envelope.Kind),
		)
	}
}

func (w *ImageAssetWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, tenantID, assetID int64, kind string) {
	if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil {
		w.logger.Warn("image asset nack failed",
			zap.Error(nackErr),
			zap.Int64("tenant_id", tenantID),
			zap.Int64("asset_id", assetID),
			zap.String("kind", kind),
		)
	}
	if w.logger != nil {
		w.logger.Warn("image asset processing failed",
			zap.Error(err),
			zap.Bool("requeue", requeue),
			zap.Bool("redelivered", delivery.Redelivered),
			zap.Int64("tenant_id", tenantID),
			zap.Int64("asset_id", assetID),
			zap.String("kind", kind),
		)
	}
}

func shouldRequeueImageAssetError(err error, redelivered bool) bool {
	if redelivered {
		return false
	}

	appErr := response.Normalize(err)
	return appErr.HTTPStatus >= 500
}