xinmi/internal/container/recover_pending_wiki_tasks.go

package container

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	apprepo "github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/application/service"
	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/hibiken/asynq"
	"gorm.io/gorm"
)

// wikiRestartClaimGrace bounds how long a wiki claim may outlive the worker that
// stamped it. It must exceed the longest plausible batch turn (the per-turn model
// call budget is four attempts bounded by wikiLLMCallTimeout) and stay well below
// wikiClaimStaleAfter so a restart resumes faster than the steady-state rule.
const wikiRestartClaimGrace = 15 * time.Minute

// sweepOrphanedWikiClaims releases wiki queue claims stamped before now-grace, which
// no live worker of this deployment can still own.
func sweepOrphanedWikiClaims(ctx context.Context, db *gorm.DB, now time.Time) (int64, error) {
	if db == nil || !db.Migrator().HasTable("task_pending_ops") {
		return 0, nil
	}
	result := db.WithContext(ctx).
		Table("task_pending_ops").
		Where("task_type IN ?", []string{types.TypeWikiIngest, types.TypeWikiFinalize}).
		Where("claimed_at IS NOT NULL AND claimed_at < ?", now.UTC().Add(-wikiRestartClaimGrace)).
		Update("claimed_at", nil)
	return result.RowsAffected, result.Error
}

// pendingWikiScope is the minimum routing information needed to recreate an
// ephemeral wiki trigger from the durable task_pending_ops queue.
type pendingWikiScope struct {
	TenantID uint64 `gorm:"column:tenant_id"`
	TaskType string `gorm:"column:task_type"`
	ScopeID  string `gorm:"column:scope_id"`
}

// recoverPendingWikiTasks recreates one trigger per pending wiki queue lane.
//
// task_pending_ops is deliberately durable, but the trigger that wakes its
// consumer is not durable in Lite mode (SyncTaskExecutor is process-local) and
// may also be absent after an interrupted Redis enqueue. Running this after all
// handlers are registered closes that gap. Duplicate triggers are harmless:
// ingest claims/peeks disjoint rows and finalize coalesces its pending lane.
func recoverPendingWikiTasks(db *gorm.DB, task interfaces.TaskEnqueuer) {
	if db == nil || task == nil {
		return
	}
	ctx := context.Background()
	const activeKnowledgeBase = `EXISTS (
		SELECT 1 FROM knowledge_bases kb
		WHERE kb.id = task_pending_ops.scope_id
			AND kb.tenant_id = task_pending_ops.tenant_id
			AND kb.deleted_at IS NULL
	)`
	wikiTaskTypes := []string{types.TypeWikiIngest, types.TypeWikiFinalize}
	if db.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
		var stopped []types.WikiRecoveryBudget
		if err := db.WithContext(ctx).Where("status = ? OR (status = ? AND deadline_at <= ?)", "stopped_budget", "active", time.Now().UTC()).Find(&stopped).Error; err != nil {
			logger.Errorf(ctx, "[WikiRecovery] cannot verify recovery budgets: %v", err)
			return
		}
		repo := apprepo.WikiRecoveryBudgetRepository{DB: db}
		for _, budget := range stopped {
			reason := budget.StopReason
			if reason == "" {
				reason = "deadline"
			}
			if err := repo.StopAndArchive(ctx, budget, reason, "startup"); err != nil {
				logger.Errorf(ctx, "[WikiRecovery] cannot settle stopped recovery %d: %v", budget.PendingOpID, err)
				return
			}
		}
	}

	// Durable rows for a deleted/missing KB must not recreate ephemeral
	// triggers at startup. Fail closed if this cleanup cannot be verified.
	cleanup := db.WithContext(ctx).
		Where("scope = ? AND task_type IN ?", types.TaskScopeKnowledgeBase, wikiTaskTypes).
		Where("NOT " + activeKnowledgeBase).
		Delete(&types.TaskPendingOp{})
	if cleanup.Error != nil {
		logger.Warnf(ctx, "[WikiRecovery] failed to clear deleted KB queues: %v", cleanup.Error)
		return
	}
	if cleanup.RowsAffected > 0 {
		logger.Infof(ctx, "[WikiRecovery] removed %d pending row(s) for deleted knowledge bases", cleanup.RowsAffected)
	}

	// Claims are not renewed, so a claim older than the longest plausible batch turn
	// can only belong to a worker that is gone. This runs at startup precisely because
	// a restart means the previous incarnation of this process died: releasing those
	// claims lets the queue resume at once instead of waiting out the steady-state
	// stale threshold (wikiClaimStaleAfter, 90 minutes).
	released, sweepErr := sweepOrphanedWikiClaims(ctx, db, time.Now())
	if sweepErr != nil {
		logger.Warnf(ctx, "[WikiRecovery] failed to sweep stale wiki claims: %v", sweepErr)
	} else if released > 0 {
		logger.Infof(ctx, "[WikiRecovery] released %d wiki claim(s) left by a stopped worker", released)
	}

	var scopes []pendingWikiScope
	if err := db.WithContext(ctx).
		Model(&types.TaskPendingOp{}).
		Distinct("tenant_id", "task_type", "scope_id").
		Where("scope = ? AND task_type IN ?", types.TaskScopeKnowledgeBase, wikiTaskTypes).
		Where(activeKnowledgeBase).
		Find(&scopes).Error; err != nil {
		logger.Warnf(ctx, "[WikiRecovery] failed to list pending queues: %v", err)
		return
	}

	recovered := 0
	for _, scope := range scopes {
		if scope.ScopeID == "" {
			continue
		}
		payload, err := json.Marshal(service.WikiIngestPayload{
			TenantID:        scope.TenantID,
			KnowledgeBaseID: scope.ScopeID,
		})
		if err != nil {
			logger.Warnf(ctx, "[WikiRecovery] marshal trigger for KB %s failed: %v", scope.ScopeID, err)
			continue
		}
		opts := []asynq.Option{
			asynq.Queue(types.QueueWiki),
			asynq.MaxRetry(10), // keep aligned with the wiki ingest retry policy
			asynq.Timeout(60 * time.Minute),
		}
		if scope.TaskType == types.TypeWikiFinalize {
			opts[2] = asynq.Timeout(30 * time.Minute)
			// Match scheduleFinalize so simultaneous replica startups collapse
			// into the same per-KB finalize trigger.
			opts = append(opts, asynq.TaskID("wiki-finalize-"+scope.ScopeID))
		}
		trigger := asynq.NewTask(scope.TaskType, payload, opts...)
		if _, err := task.Enqueue(trigger); err != nil {
			if errors.Is(err, asynq.ErrTaskIDConflict) || errors.Is(err, asynq.ErrDuplicateTask) {
				recovered++
				continue
			}
			logger.Warnf(ctx, "[WikiRecovery] enqueue %s trigger for KB %s failed: %v",
				scope.TaskType, scope.ScopeID, err)
			continue
		}
		recovered++
	}
	if recovered > 0 {
		logger.Infof(ctx, "[WikiRecovery] recreated %d trigger(s) from durable pending queues", recovered)
	}
}