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

package generate

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"go.uber.org/zap"

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

const (
	defaultGenerationTaskRecoveryInterval = time.Minute
	defaultGenerationTaskRecoveryTimeout  = 30 * time.Second
)

type generationTaskRecoveryConfig struct {
	LeaseTTL         time.Duration
	RecoveryInterval time.Duration
	RecoveryTimeout  time.Duration
	BatchSize        int
	QueuedStaleAfter time.Duration
	MaxAttempts      int
}

type claimedGenerationTaskLease struct {
	task  generationTaskRecoveryRecord
	token string
}

type generationTaskRecoveryRecord struct {
	ID                 int64
	TenantID           int64
	OperatorID         int64
	ArticleID          int64
	QuotaReservationID int64
	TaskType           string
	InputParams        map[string]interface{}
	AttemptCount       int
}

func (r generationTaskRecoveryRecord) toClaimedTask() tenantapp.ClaimedGenerationTask {
	return tenantapp.ClaimedGenerationTask{
		ID:                 r.ID,
		TenantID:           r.TenantID,
		OperatorID:         r.OperatorID,
		ArticleID:          r.ArticleID,
		QuotaReservationID: r.QuotaReservationID,
		TaskType:           r.TaskType,
		InputParams:        r.InputParams,
	}
}

type recoveredGenerationTask struct {
	ID           int64
	TenantID     int64
	ArticleID    int64
	TaskType     string
	AttemptCount int
}

func (w *ArticleGenerationWorker) startRecovery(ctx context.Context) {
	if w == nil || w.pool == nil {
		return
	}
	go w.runRecovery(ctx)
}

func (w *ArticleGenerationWorker) runRecovery(ctx context.Context) {
	w.runRecoveryOnce(context.Background())

	for {
		interval := w.runtimeRecoveryConfig().RecoveryInterval
		timer := time.NewTimer(interval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
			w.runRecoveryOnce(context.Background())
		}
	}
}

func (w *ArticleGenerationWorker) runRecoveryOnce(parent context.Context) {
	cfg := w.runtimeRecoveryConfig()
	ctx, cancel := context.WithTimeout(parent, cfg.RecoveryTimeout)
	defer cancel()

	queuedRecovered, requeueErr := w.requeueStaleGenerationTasks(ctx, cfg)
	if requeueErr != nil && w.logger != nil {
		w.logger.Warn("article generation queued task recovery failed", zap.Error(requeueErr))
	}

	expired, expireErr := w.expireStaleRunningGenerationTasks(ctx, cfg)
	if expireErr != nil && w.logger != nil {
		w.logger.Warn("article generation running task recovery failed", zap.Error(expireErr))
	}

	if w.logger != nil && (queuedRecovered > 0 || expired > 0) {
		w.logger.Info("article generation task recovery completed",
			zap.Int("queued_recovered_task_count", queuedRecovered),
			zap.Int("running_recovered_task_count", expired),
		)
	}
}

func (w *ArticleGenerationWorker) runtimeRecoveryConfig() generationTaskRecoveryConfig {
	if w != nil && w.configProvider != nil {
		if cfg := w.configProvider.Current(); cfg != nil {
			return normalizeGenerationTaskRecoveryConfig(cfg.Generation)
		}
	}
	if w != nil {
		return normalizeGenerationTaskRecoveryConfig(w.generationCfg)
	}
	return normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{})
}

func normalizeGenerationTaskRecoveryConfig(cfg config.GenerationConfig) generationTaskRecoveryConfig {
	normalized := cfg
	config.NormalizeGenerationConfig(&normalized)
	return generationTaskRecoveryConfig{
		LeaseTTL:         normalized.TaskLeaseTTL,
		RecoveryInterval: normalized.TaskRecoveryInterval,
		RecoveryTimeout:  normalized.TaskRecoveryTimeout,
		BatchSize:        normalized.TaskRecoveryBatchSize,
		QueuedStaleAfter: normalized.TaskQueuedStaleAfter,
		MaxAttempts:      normalized.TaskMaxAttempts,
	}
}

func (w *ArticleGenerationWorker) claimGenerationTask(ctx context.Context, taskID int64, owner string, leaseTTL time.Duration) (*claimedGenerationTaskLease, bool, error) {
	if leaseTTL <= 0 {
		leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL
	}
	token, err := newGenerationTaskLeaseToken()
	if err != nil {
		return nil, false, err
	}
	leaseExpiresAt := time.Now().UTC().Add(leaseTTL)

	var (
		row             generationTaskRecoveryRecord
		operatorID      sql.NullInt64
		articleID       sql.NullInt64
		reservationID   sql.NullInt64
		inputParamsJSON []byte
	)
	err = w.pool.QueryRow(ctx, `
		UPDATE generation_tasks
		SET status = 'running',
		    started_at = COALESCE(started_at, NOW()),
		    lease_token = $2,
		    lease_owner = $3,
		    lease_expires_at = $4,
		    last_heartbeat_at = NOW(),
		    attempt_count = attempt_count + 1,
		    updated_at = NOW()
		WHERE id = $1
		  AND status = 'queued'
		RETURNING id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count
	`, taskID, token, owner, leaseExpiresAt).Scan(
		&row.ID,
		&row.TenantID,
		&operatorID,
		&articleID,
		&reservationID,
		&row.TaskType,
		&inputParamsJSON,
		&row.AttemptCount,
	)
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil, true, nil
		}
		return nil, false, err
	}

	row.OperatorID = operatorID.Int64
	row.ArticleID = articleID.Int64
	row.QuotaReservationID = reservationID.Int64
	params, err := parseJSONMap(inputParamsJSON)
	if err != nil {
		row.InputParams = map[string]interface{}{}
		return &claimedGenerationTaskLease{task: row, token: token}, false, fmt.Errorf("decode generation input params: %w", err)
	}
	row.InputParams = params

	return &claimedGenerationTaskLease{task: row, token: token}, false, nil
}

func (w *ArticleGenerationWorker) renewGenerationTaskLease(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) (bool, error) {
	if taskID <= 0 || token == "" {
		return false, nil
	}
	if leaseTTL <= 0 {
		leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL
	}
	tag, err := w.pool.Exec(ctx, `
		UPDATE generation_tasks
		SET lease_expires_at = $3,
		    last_heartbeat_at = NOW(),
		    updated_at = NOW()
		WHERE id = $1
		  AND lease_token = $2
		  AND status = 'running'
	`, taskID, token, time.Now().UTC().Add(leaseTTL))
	if err != nil {
		return false, err
	}
	return tag.RowsAffected() > 0, nil
}

func (w *ArticleGenerationWorker) releaseGenerationTaskLease(ctx context.Context, taskID int64, token string) {
	if w == nil || w.pool == nil || taskID <= 0 || token == "" {
		return
	}
	releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if _, err := w.pool.Exec(releaseCtx, `
		UPDATE generation_tasks
		SET lease_token = NULL,
		    lease_owner = NULL,
		    lease_expires_at = NULL,
		    updated_at = NOW()
		WHERE id = $1
		  AND lease_token = $2
		  AND status IN ('completed', 'failed')
	`, taskID, token); err != nil && w.logger != nil {
		w.logger.Warn("article generation task lease release failed",
			zap.Error(err),
			zap.Int64("task_id", taskID),
		)
	}
}

func (w *ArticleGenerationWorker) runGenerationLeaseHeartbeat(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) {
	if taskID <= 0 || token == "" || leaseTTL <= 0 {
		return
	}

	interval := leaseTTL / 3
	if interval < 10*time.Second {
		interval = 10 * time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			heartbeatCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			ok, err := w.renewGenerationTaskLease(heartbeatCtx, taskID, token, leaseTTL)
			cancel()
			if err != nil {
				if w.logger != nil {
					w.logger.Warn("article generation task lease heartbeat failed",
						zap.Error(err),
						zap.Int64("task_id", taskID),
					)
				}
				continue
			}
			if !ok {
				if w.logger != nil {
					w.logger.Warn("article generation task lease lost",
						zap.Int64("task_id", taskID),
					)
				}
				return
			}
		}
	}
}

func (w *ArticleGenerationWorker) requeueStaleGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) {
	if w == nil || w.pool == nil || w.rabbitMQ == nil || cfg.BatchSize <= 0 {
		return 0, nil
	}
	staleBefore := time.Now().UTC().Add(-cfg.QueuedStaleAfter)

	rows, err := w.pool.Query(ctx, `
		WITH candidates AS (
			SELECT id, tenant_id, article_id, task_type, attempt_count
			FROM generation_tasks
			WHERE status = 'queued'
			  AND article_id IS NOT NULL
			  AND updated_at < $1
			ORDER BY updated_at ASC, id ASC
			FOR UPDATE SKIP LOCKED
			LIMIT $2
		),
		touched AS (
			UPDATE generation_tasks gt
			SET updated_at = NOW()
			FROM candidates c
			WHERE gt.id = c.id
			RETURNING c.id, c.tenant_id, c.article_id, c.task_type, c.attempt_count
		)
		SELECT id, tenant_id, article_id, task_type, attempt_count
		FROM touched
	`, staleBefore, cfg.BatchSize)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	items := make([]recoveredGenerationTask, 0, cfg.BatchSize)
	for rows.Next() {
		var item recoveredGenerationTask
		if err := rows.Scan(&item.ID, &item.TenantID, &item.ArticleID, &item.TaskType, &item.AttemptCount); err != nil {
			return 0, err
		}
		items = append(items, item)
	}
	if err := rows.Err(); err != nil {
		return 0, err
	}

	requeued := 0
	for _, item := range items {
		if item.AttemptCount >= cfg.MaxAttempts {
			ok, err := w.finalizeExpiredGenerationTask(ctx, generationTaskRecoveryRecord{
				ID:           item.ID,
				TenantID:     item.TenantID,
				ArticleID:    item.ArticleID,
				TaskType:     item.TaskType,
				AttemptCount: item.AttemptCount,
			})
			if err != nil && w.logger != nil {
				w.logger.Warn("article generation stale queued task finalization failed",
					zap.Error(err),
					zap.Int64("task_id", item.ID),
					zap.Int64("article_id", item.ArticleID),
				)
			}
			if ok {
				requeued++
			}
			continue
		}
		if err := w.publishGenerationTask(ctx, item.ID, item.TaskType, item.TenantID, item.ArticleID); err != nil {
			if w.logger != nil {
				w.logger.Warn("article generation stale queued task republish failed",
					zap.Error(err),
					zap.Int64("task_id", item.ID),
					zap.Int64("article_id", item.ArticleID),
				)
			}
			continue
		}
		requeued++
	}

	return requeued, nil
}

func (w *ArticleGenerationWorker) expireStaleRunningGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) {
	if w == nil || w.pool == nil || cfg.BatchSize <= 0 {
		return 0, nil
	}
	now := time.Now().UTC()

	rows, err := w.pool.Query(ctx, `
		SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count
		FROM generation_tasks
		WHERE status = 'running'
		  AND article_id IS NOT NULL
		  AND (
			lease_expires_at IS NULL
			OR lease_expires_at < $1
		  )
		ORDER BY COALESCE(lease_expires_at, started_at, updated_at, created_at) ASC, id ASC
		LIMIT $2
	`, now, cfg.BatchSize)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	recovered := make([]generationTaskRecoveryRecord, 0, cfg.BatchSize)
	for rows.Next() {
		var (
			item            generationTaskRecoveryRecord
			operatorID      sql.NullInt64
			articleID       sql.NullInt64
			reservationID   sql.NullInt64
			inputParamsJSON []byte
		)
		if err := rows.Scan(
			&item.ID,
			&item.TenantID,
			&operatorID,
			&articleID,
			&reservationID,
			&item.TaskType,
			&inputParamsJSON,
			&item.AttemptCount,
		); err != nil {
			return 0, err
		}
		item.OperatorID = operatorID.Int64
		item.ArticleID = articleID.Int64
		item.QuotaReservationID = reservationID.Int64
		params, err := parseJSONMap(inputParamsJSON)
		if err != nil {
			params = map[string]interface{}{}
		}
		item.InputParams = params
		recovered = append(recovered, item)
	}
	if err := rows.Err(); err != nil {
		return 0, err
	}

	handled := 0
	for _, item := range recovered {
		if item.AttemptCount < cfg.MaxAttempts {
			ok, err := w.requeueExpiredRunningGenerationTask(ctx, item)
			if err != nil && w.logger != nil {
				w.logger.Warn("article generation expired task requeue failed",
					zap.Error(err),
					zap.Int64("task_id", item.ID),
					zap.Int64("article_id", item.ArticleID),
				)
			}
			if ok {
				handled++
			}
			continue
		}

		ok, err := w.finalizeExpiredGenerationTask(ctx, item)
		if err != nil && w.logger != nil {
			w.logger.Warn("article generation expired task finalization failed",
				zap.Error(err),
				zap.Int64("task_id", item.ID),
				zap.Int64("article_id", item.ArticleID),
			)
		}
		if ok {
			handled++
		}
	}

	return handled, nil
}

func (w *ArticleGenerationWorker) requeueExpiredRunningGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) {
	if w == nil || w.rabbitMQ == nil {
		return false, nil
	}
	tx, err := w.pool.Begin(ctx)
	if err != nil {
		return false, err
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	var (
		tenantID  int64
		articleID int64
		taskType  string
	)
	if err := tx.QueryRow(ctx, `
		UPDATE generation_tasks
		SET status = 'queued',
		    error_message = NULL,
		    lease_token = NULL,
		    lease_owner = NULL,
		    lease_expires_at = NULL,
		    updated_at = NOW()
		WHERE id = $1
		  AND status = 'running'
		  AND (
			lease_expires_at IS NULL
			OR lease_expires_at < NOW()
		  )
		RETURNING tenant_id, article_id, task_type
	`, item.ID).Scan(&tenantID, &articleID, &taskType); err != nil {
		if err == pgx.ErrNoRows {
			return false, nil
		}
		return false, err
	}

	if err := tx.Commit(ctx); err != nil {
		return false, err
	}
	if err := w.publishGenerationTask(ctx, item.ID, taskType, tenantID, articleID); err != nil {
		return false, err
	}
	return true, nil
}

func (w *ArticleGenerationWorker) finalizeExpiredGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) {
	errMsg := formatGenerationTaskRecoveryFailure()

	tx, err := w.pool.Begin(ctx)
	if err != nil {
		return false, err
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	var (
		claimedID          int64
		operatorID         sql.NullInt64
		articleID          sql.NullInt64
		quotaReservationID sql.NullInt64
		taskType           string
	)
	if err := tx.QueryRow(ctx, `
		UPDATE generation_tasks
		SET status = 'failed',
		    error_message = $2,
		    completed_at = COALESCE(completed_at, NOW()),
		    lease_token = NULL,
		    lease_owner = NULL,
		    lease_expires_at = NULL,
		    updated_at = NOW()
		WHERE id = $1
		  AND status IN ('running', 'queued')
		  AND (
			status = 'queued'
			OR lease_expires_at IS NULL
			OR lease_expires_at < NOW()
		  )
		RETURNING id, operator_id, article_id, quota_reservation_id, task_type
	`, item.ID, errMsg).Scan(&claimedID, &operatorID, &articleID, &quotaReservationID, &taskType); err != nil {
		if err == pgx.ErrNoRows {
			return false, nil
		}
		return false, err
	}
	_ = claimedID
	item.OperatorID = operatorID.Int64
	item.ArticleID = articleID.Int64
	item.QuotaReservationID = quotaReservationID.Int64
	item.TaskType = strings.TrimSpace(taskType)

	if item.ArticleID > 0 {
		if _, err := tx.Exec(ctx, `
			UPDATE articles
			SET generate_status = 'failed',
			    updated_at = NOW()
			WHERE id = $1
			  AND tenant_id = $2
			  AND generate_status = 'generating'
			  AND deleted_at IS NULL
		`, item.ArticleID, item.TenantID); err != nil {
			return false, err
		}
	}

	if item.TaskType == kolGenerationTaskType {
		if _, err := tx.Exec(ctx, `
			UPDATE task_records
			SET status = 'failed',
			    error_message = $1,
			    updated_at = NOW()
			WHERE tenant_id = $2
			  AND task_type = $3
			  AND resource_type = $4
			  AND resource_id = $5
		`, errMsg, item.TenantID, "article_generation", "generation_task", item.ID); err != nil {
			return false, err
		}
		if _, err := tx.Exec(ctx, `
			UPDATE kol_usage_logs
			SET status = 'failed',
			    error_message = $1,
			    finished_at = COALESCE(finished_at, NOW())
			WHERE generation_task_id = $2
			  AND tenant_id = $3
			  AND status IN ('pending', 'running')
		`, errMsg, item.ID, item.TenantID); err != nil {
			return false, err
		}
	}

	if item.QuotaReservationID > 0 {
		var reservationStatus string
		if err := tx.QueryRow(ctx, `
			WITH refunded AS (
				UPDATE quota_reservations
				SET status = 'refunded',
				    refunded_amount = reserved_amount,
				    updated_at = NOW()
				WHERE id = $1
				  AND tenant_id = $2
				  AND status = 'pending'
				RETURNING status
			)
			SELECT status FROM refunded
		`, item.QuotaReservationID, item.TenantID).Scan(&reservationStatus); err != nil && err != pgx.ErrNoRows {
			return false, err
		}
		if reservationStatus == "refunded" {
			if err := insertGenerationRecoveryRefundLedger(ctx, tx, item); err != nil {
				return false, err
			}
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return false, err
	}

	cacheCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if item.ArticleID > 0 {
		tenantapp.InvalidateArticleCaches(cacheCtx, w.cache, item.TenantID, &item.ArticleID)
	}
	return true, nil
}

func insertGenerationRecoveryRefundLedger(ctx context.Context, tx pgx.Tx, item generationTaskRecoveryRecord) error {
	reason := "generation_refund"
	referenceType := "generation_task"
	var operatorID interface{}
	if item.OperatorID > 0 {
		operatorID = item.OperatorID
	}
	_, err := tx.Exec(ctx, `
		INSERT INTO tenant_quota_ledgers (
			tenant_id,
			operator_id,
			quota_type,
			delta,
			balance_after,
			reason,
			reference_type,
			reference_id
		)
		SELECT
			$1,
			$2,
			'article_generation',
			1,
			COALESCE((
				SELECT balance_after
				FROM tenant_quota_ledgers
				WHERE tenant_id = $1
				  AND quota_type = 'article_generation'
				ORDER BY created_at DESC, id DESC
				LIMIT 1
			), 0) + 1,
			$3,
			$4,
			$5
		WHERE NOT EXISTS (
			SELECT 1
			FROM tenant_quota_ledgers
			WHERE tenant_id = $1
			  AND quota_type = 'article_generation'
				AND reason = $3
			  AND reference_type = $4
			  AND reference_id = $5
		)
	`, item.TenantID, operatorID, reason, referenceType, item.ID)
	return err
}

func formatGenerationTaskRecoveryFailure() string {
	return "文章生成任务执行超时，worker 重启恢复后自动标记失败 [worker_recovery]"
}

func (w *ArticleGenerationWorker) publishGenerationTask(ctx context.Context, taskID int64, taskType string, tenantID, articleID int64) error {
	if w == nil || w.rabbitMQ == nil {
		return nil
	}
	payload, err := json.Marshal(tenantapp.ArticleGenerationTaskEnvelope{
		TaskID:    taskID,
		TaskType:  taskType,
		TenantID:  tenantID,
		ArticleID: articleID,
	})
	if err != nil {
		return fmt.Errorf("marshal generation task envelope: %w", err)
	}
	if err := w.rabbitMQ.PublishGenerationTask(ctx, payload); err != nil {
		return fmt.Errorf("publish generation task: %w", err)
	}
	return nil
}

func newGenerationTaskLeaseToken() (string, error) {
	var raw [16]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", fmt.Errorf("create generation task lease token: %w", err)
	}
	return hex.EncodeToString(raw[:]), nil
}

func generationWorkerOwner(consumerName string) string {
	host, _ := os.Hostname()
	owner := strings.TrimSpace(host + "/" + consumerName)
	if owner == "/" {
		return "worker-generate"
	}
	if len(owner) > 128 {
		return owner[:128]
	}
	return owner
}