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), ) } }