geo/server/internal/worker/generate/template_assist_worker.go

package generate

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

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

	"github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
)

var errTemplateAssistConsumerClosed = errors.New("template assist consumer channel closed")

type TemplateAssistWorker struct {
	rabbitMQ          *rabbitmq.Client
	templateService   *tenantapp.TemplateService
	logger            *zap.Logger
	retryInterval     time.Duration
	processTimeout    time.Duration
	configProvider    config.Provider
	consumerPrefix    string
	workerConcurrency int
}

func NewTemplateAssistWorker(
	rabbitMQClient *rabbitmq.Client,
	templateService *tenantapp.TemplateService,
	logger *zap.Logger,
	workerConcurrency int,
	processTimeout time.Duration,
) *TemplateAssistWorker {
	if workerConcurrency <= 0 {
		workerConcurrency = 1
	}
	if processTimeout <= 0 {
		processTimeout = 5 * time.Minute
	}

	return &TemplateAssistWorker{
		rabbitMQ:          rabbitMQClient,
		templateService:   templateService,
		logger:            logger,
		retryInterval:     5 * time.Second,
		processTimeout:    processTimeout,
		consumerPrefix:    fmt.Sprintf("template-assist-%d", os.Getpid()),
		workerConcurrency: workerConcurrency,
	}
}

func (w *TemplateAssistWorker) WithConfigProvider(provider config.Provider) *TemplateAssistWorker {
	w.configProvider = provider
	return w
}

func (w *TemplateAssistWorker) Start(ctx context.Context) {
	if w == nil || w.rabbitMQ == nil || w.templateService == nil {
		return
	}

	for i := 0; i < w.workerConcurrency; i++ {
		consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1)
		go w.run(ctx, consumerName)
	}
}

func (w *TemplateAssistWorker) run(ctx context.Context, consumerName string) {
	for {
		if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil {
			if errors.Is(err, errTemplateAssistConsumerClosed) {
				w.logger.Info("template assist consumer channel closed, retrying",
					zap.String("consumer", consumerName),
				)
			} else {
				w.logger.Warn("template assist worker stopped unexpectedly",
					zap.Error(err),
					zap.String("consumer", consumerName),
				)
			}
		}

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

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

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

func (w *TemplateAssistWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) {
	job, err := tenantapp.DecodeTemplateAssistTask(delivery.Body)
	if err != nil {
		w.rejectDelivery(delivery, false, err, "", "")
		return
	}

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

	terminal, err := w.templateService.IsAssistTaskTerminal(ctx, *job)
	if err != nil {
		requeue := !delivery.Redelivered
		w.rejectDelivery(delivery, requeue, err, job.TaskID, job.TaskType)
		return
	}
	if terminal {
		if err := delivery.Ack(false); err != nil && w.logger != nil {
			w.logger.Warn("template assist ack failed for terminal task",
				zap.Error(err),
				zap.String("task_id", job.TaskID),
				zap.String("task_type", job.TaskType),
			)
		}
		return
	}

	w.templateService.ProcessQueuedAssistJob(ctx, *job)

	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("template assist ack failed",
			zap.Error(err),
			zap.String("task_id", job.TaskID),
			zap.String("task_type", job.TaskType),
		)
	}
}

func (w *TemplateAssistWorker) runtimeProcessTimeout() time.Duration {
	if w != nil && w.configProvider != nil {
		if cfg := w.configProvider.Current(); cfg != nil {
			return templateAssistProcessTimeout(cfg.Generation)
		}
	}
	if w != nil && w.processTimeout > 0 {
		return w.processTimeout
	}
	return templateAssistProcessTimeout(config.GenerationConfig{})
}

func templateAssistProcessTimeout(cfg config.GenerationConfig) time.Duration {
	if cfg.ArticleTimeout > 0 {
		return cfg.ArticleTimeout
	}
	return 5 * time.Minute
}

func (w *TemplateAssistWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID, taskType string) {
	if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil {
		w.logger.Warn("template assist nack failed",
			zap.Error(nackErr),
			zap.String("task_id", taskID),
			zap.String("task_type", taskType),
		)
	}
	if w.logger != nil {
		w.logger.Warn("template assist processing failed",
			zap.Error(err),
			zap.Bool("requeue", requeue),
			zap.Bool("redelivered", delivery.Redelivered),
			zap.String("task_id", taskID),
			zap.String("task_type", taskType),
		)
	}
}