FranklinWH-Toolbox/service/task/queue/fwhQueue/fwhTaskSiteInfoUpdate.go

package fwhQueue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
	"fwh-toolbox-server/repositories"
	"time"

	"github.com/hibiken/asynq"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.uber.org/zap"
)

// handleSiteInfoUpdate 是站点信息更新任务的 Asynq handler
func handleSiteInfoUpdate(ctx context.Context, task *asynq.Task) error {
	var payload domain.TaskSiteInfoUpdate
	if err := json.Unmarshal(task.Payload(), &payload); err != nil {
		return fmt.Errorf("unmarshal SiteInfoUpdate payload: %w", err)
	}

	taskId := payload.Id.Hex()

	// 更新运行状态
	if err := repositories.UpdateSiteInfoUpdateResult(payload.Id, domain.ProgressRunning, 0); err != nil {
		zap.L().Error("[SiteInfoUpdate] 更新状态失败(Running)", zap.String("id", taskId))
		return err
	}

	// 执行任务
	total, err := executeSiteInfoUpdate(ctx, payload)
	if errors.Is(err, context.Canceled) {
		// CancelSiteInfoUpdate 发出的取消（asynq 超时为 DeadlineExceeded，按失败处理）
		zap.L().Info("[SiteInfoUpdate] 任务已取消", zap.String("id", taskId))
		repositories.UpdateSiteInfoUpdateResult(payload.Id, domain.ProgressCancelled, 0)
		return err
	}
	if err != nil {
		zap.L().Error("[SiteInfoUpdate] 任务执行失败", zap.String("id", taskId), zap.String("err", err.Error()))
		repositories.UpdateSiteInfoUpdateResult(payload.Id, domain.ProgressError, 0)
		return nil // 业务错误不重试
	}

	zap.L().Info("[SiteInfoUpdate] 任务执行成功", zap.String("id", taskId), zap.Uint64("total", total))
	repositories.UpdateSiteInfoUpdateResult(payload.Id, domain.ProgressComplete, total)
	return nil
}

// executeSiteInfoUpdate 执行站点信息更新
func executeSiteInfoUpdate(ctx context.Context, payload domain.TaskSiteInfoUpdate) (uint64, error) {
	var pageSize uint = 50
	var pageNum uint = 1
	var total uint64 = 0
	var realTotal uint64 = 0

	for {
		// 每页检查一次：取消（CancelSiteInfoUpdate）或超时（SiteInfoUpdateTimeout）时停止
		if err := ctx.Err(); err != nil {
			return 0, err
		}
		resp, err := repositories.QuerySiteInfoList(dto.SiteInfoPara{
			PageSize: pageSize,
			PageNum:  pageNum,
		})
		if err != nil {
			return 0, fmt.Errorf("查询站点信息失败(page=%d): %w", pageNum, err)
		}

		for _, info := range *resp.Result {
			if err := repositories.InsertSiteInfoData(*domain.FwhSiteInfoToSiteInfoDomain(&info)); err != nil {
				return 0, fmt.Errorf("保存站点信息失败(page=%d): %w", pageNum, err)
			}
			realTotal++
		}

		pageNum++
		total = resp.Total
		if (int64(pageNum)-1)*int64(pageSize) >= int64(total) {
			break
		}
	}

	return realTotal, nil
}

// SiteInfoUpdateTimeout 站点信息同步任务的执行超时
const SiteInfoUpdateTimeout = 4 * time.Hour

// CancelSiteInfoUpdate 取消一次站点信息同步：
// 正在执行的通知 worker 停止（worker 停下后把记录标记为 cancelled）；还在排队或已不在队列中的直接删除并标记。
// 可在主服务或只入队的独立进程（fwh-agent-mcp）中调用，需要 inspector 已初始化。
func CancelSiteInfoUpdate(id primitive.ObjectID) error {
	if inspector == nil {
		return fmt.Errorf("任务队列未初始化")
	}
	taskId := newTaskID(TaskTypeSiteInfoUpdate, id.Hex())
	info, err := inspector.GetTaskInfo(QueueMain, taskId)
	switch {
	case errors.Is(err, asynq.ErrTaskNotFound) || errors.Is(err, asynq.ErrQueueNotFound):
		// 队列中已没有这个任务（如主服务重启时丢失）：只更新记录
	case err != nil:
		return err
	case info.State == asynq.TaskStateActive:
		return inspector.CancelProcessing(taskId)
	case info.State == asynq.TaskStateCompleted || info.State == asynq.TaskStateArchived:
		// 任务已经结束，worker 已写入最终状态，不覆盖
		return nil
	case info.State == asynq.TaskStatePending || info.State == asynq.TaskStateScheduled || info.State == asynq.TaskStateRetry:
		if err := inspector.DeleteTask(QueueMain, taskId); err != nil {
			return err
		}
	}
	return repositories.UpdateSiteInfoUpdateResult(id, domain.ProgressCancelled, 0)
}