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

package generate

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"

	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"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"
)

const defaultArticleGenerationWorkerRetryInterval = 5 * time.Second

var errArticleGenerationConsumerClosed = errors.New("article generation consumer channel closed")

type ArticleGenerationWorker struct {
	pool              *pgxpool.Pool
	rabbitMQ          *rabbitmq.Client
	templateService   *tenantapp.TemplateService
	promptRuleService *tenantapp.PromptRuleGenerationService
	imitationService  *tenantapp.ArticleImitationService
	kolWorker         *KolGenerationWorker
	cache             sharedcache.Cache
	logger            *zap.Logger
	retryInterval     time.Duration
	processTimeout    time.Duration
	generationCfg     config.GenerationConfig
	configProvider    config.Provider
	consumerPrefix    string
	workerConcurrency int
}

func NewArticleGenerationWorker(
	pool *pgxpool.Pool,
	rabbitMQClient *rabbitmq.Client,
	templateService *tenantapp.TemplateService,
	promptRuleService *tenantapp.PromptRuleGenerationService,
	imitationService *tenantapp.ArticleImitationService,
	kolWorker *KolGenerationWorker,
	logger *zap.Logger,
	cfg config.GenerationConfig,
) *ArticleGenerationWorker {
	workerConcurrency := cfg.WorkerConcurrency
	if workerConcurrency <= 0 {
		workerConcurrency = 1
	}

	return &ArticleGenerationWorker{
		pool:              pool,
		rabbitMQ:          rabbitMQClient,
		templateService:   templateService,
		promptRuleService: promptRuleService,
		imitationService:  imitationService,
		kolWorker:         kolWorker,
		logger:            logger,
		retryInterval:     defaultArticleGenerationWorkerRetryInterval,
		processTimeout:    articleGenerationProcessTimeout(cfg),
		generationCfg:     cfg,
		consumerPrefix:    fmt.Sprintf("worker-generate-%d", os.Getpid()),
		workerConcurrency: workerConcurrency,
	}
}

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

func (w *ArticleGenerationWorker) WithCache(c sharedcache.Cache) *ArticleGenerationWorker {
	w.cache = c
	return w
}

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

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

func (w *ArticleGenerationWorker) 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, errArticleGenerationConsumerClosed) {
				w.logger.Info("article generation consumer channel closed, retrying",
					zap.String("consumer", consumerName),
				)
			} else {
				w.logger.Warn("article generation worker stopped unexpectedly",
					zap.Error(err),
					zap.String("consumer", consumerName),
				)
			}
		}

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

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

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

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

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

	lease, terminal, err := w.claimGenerationTask(ctx, envelope.TaskID, generationWorkerOwner(delivery.ConsumerTag), recoveryCfg.LeaseTTL)
	if err != nil {
		if lease != nil {
			task := lease.task.toClaimedTask()
			w.handleTaskFailure(ctx, task, "task_load", err)
			w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token)
			if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil {
				w.logger.Warn("article generation ack failed after task claim failure",
					zap.Error(ackErr),
					zap.Int64("task_id", envelope.TaskID),
				)
			}
			return
		}
		requeue := !delivery.Redelivered && !isPermanentGenerationTaskDecodeError(err)
		w.rejectDelivery(delivery, requeue, err, envelope.TaskID)
		return
	}
	if terminal {
		if err := delivery.Ack(false); err != nil && w.logger != nil {
			w.logger.Warn("article generation ack failed for terminal task",
				zap.Error(err),
				zap.Int64("task_id", envelope.TaskID),
			)
		}
		return
	}
	if lease == nil {
		w.rejectDelivery(delivery, !delivery.Redelivered, fmt.Errorf("generation task claim returned empty lease"), envelope.TaskID)
		return
	}

	task := lease.task.toClaimedTask()
	heartbeatCtx, stopHeartbeat := context.WithCancel(parent)
	go w.runGenerationLeaseHeartbeat(heartbeatCtx, task.ID, lease.token, recoveryCfg.LeaseTTL)
	defer func() {
		stopHeartbeat()
		w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token)
	}()

	task.TaskType = strings.TrimSpace(task.TaskType)
	switch {
	case task.ArticleID <= 0 || task.QuotaReservationID <= 0:
		w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("generation task references are incomplete"))
	case task.TaskType == "template":
		w.templateService.ProcessQueuedTask(ctx, task)
	case task.TaskType == "custom_generation":
		w.promptRuleService.ProcessQueuedTask(ctx, task)
	case task.TaskType == "imitation":
		if w.imitationService != nil {
			w.imitationService.ProcessQueuedTask(ctx, task)
		} else {
			w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("imitation generation worker is not configured"))
		}
	case task.TaskType == kolGenerationTaskType:
		if w.kolWorker != nil {
			w.kolWorker.Process(ctx, task)
		} else {
			w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("kol generation worker is not configured"))
		}
	default:
		w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("unsupported generation task type: %s", task.TaskType))
	}

	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("article generation ack failed",
			zap.Error(err),
			zap.Int64("task_id", task.ID),
			zap.Int64("article_id", task.ArticleID),
		)
	}
}

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

func articleGenerationProcessTimeout(cfg config.GenerationConfig) time.Duration {
	config.NormalizeGenerationConfig(&cfg)
	processTimeout := cfg.ArticleTimeout + 2*time.Minute
	if processTimeout <= 2*time.Minute {
		return 10 * time.Minute
	}
	return processTimeout
}

func (w *ArticleGenerationWorker) handleTaskFailure(ctx context.Context, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) {
	task.TaskType = strings.TrimSpace(task.TaskType)
	var err error
	switch {
	case task.TaskType == kolGenerationTaskType && w.kolWorker != nil:
		w.kolWorker.HandleTaskFailure(ctx, task, stage, taskErr)
	case task.TaskType == "imitation" && w.imitationService != nil:
		w.imitationService.HandleTaskFailure(ctx, task, stage, taskErr)
	default:
		err = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, taskErr)
	}
	if err != nil && w.logger != nil {
		w.logger.Warn("article generation task failure handling failed",
			zap.Error(err),
			zap.Int64("task_id", task.ID),
			zap.Int64("article_id", task.ArticleID),
			zap.String("task_type", task.TaskType),
			zap.String("stage", stage),
		)
	}
}

func isPermanentGenerationTaskDecodeError(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "decode generation input params")
}

func (w *ArticleGenerationWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID int64) {
	if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil {
		w.logger.Warn("article generation nack failed",
			zap.Error(nackErr),
			zap.Int64("task_id", taskID),
		)
	}
	if w.logger != nil {
		w.logger.Warn("article generation processing failed",
			zap.Error(err),
			zap.Bool("requeue", requeue),
			zap.Bool("redelivered", delivery.Redelivered),
			zap.Int64("task_id", taskID),
		)
	}
}

func parseJSONMap(raw []byte) (map[string]interface{}, error) {
	if len(raw) == 0 {
		return map[string]interface{}{}, nil
	}

	payload := make(map[string]interface{})
	if err := json.Unmarshal(raw, &payload); err != nil {
		return nil, err
	}
	return payload, nil
}