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

package fwhQueue

import (
	"context"
	"fmt"
	"time"

	"fwh-toolbox-server/dto"
	task_core "fwh-toolbox-server/task/core"

	"github.com/hibiken/asynq"
	"go.uber.org/zap"
)

// 任务类型常量
const (
	TaskTypeBatchDownload  = "fwh:BatchDownload"
	TaskTypeCallFile       = "fwh:CallFile"
	TaskTypeSiteAnalysis   = "fwh:SiteAnalysis"
	TaskTypeShutdownStats  = "fwh:ShutdownStatistics"
	TaskTypeSiteInfoUpdate = "fwh:SiteInfoUpdate"
)

// 队列名称
const (
	QueueMain     = "fwh:main"     // 主任务队列（顺序执行）
	QueueDownload = "fwh:download" // 文件下载队列（可并发）
)

// 模块级变量
var (
	TaskClient *asynq.Client
	inspector  *asynq.Inspector
)

// InitEnqueueClient 只创建入队客户端，不启动 worker、不注册处理函数。
// 供 fwh-agent-mcp 等独立进程把任务投递给主服务执行；主服务使用 InitFwhTaskQueue。
func InitEnqueueClient(redisOpt asynq.RedisClientOpt) {
	TaskClient = asynq.NewClient(redisOpt)
}

// InitFwhTaskQueue 初始化新的 FWH 任务队列
// 通过 TaskCore 创建 Asynq Server（worker）、Client（入队）、Inspector（状态查询）
func InitFwhTaskQueue(tc *task_core.TaskCore) {
	// 获取 Inspector（用于主任务轮询子任务状态）
	inspector = tc.Inspector.GetInspector()

	// 注册处理函数
	mux := asynq.NewServeMux()
	mux.HandleFunc(TaskTypeBatchDownload, handleBatchDownload)
	mux.HandleFunc(TaskTypeCallFile, handleCallFile)
	mux.HandleFunc(TaskTypeSiteAnalysis, handleSiteAnalysis)
	mux.HandleFunc(TaskTypeShutdownStats, handleShutdownStatistics)
	mux.HandleFunc(TaskTypeSiteInfoUpdate, handleSiteInfoUpdate)

	// 通过 TaskCore 创建任务 Server 和 Client
	TaskClient = tc.CreateTask(task_core.TaskConfig{
		Name: "FwhTask",
		SvrCfg: asynq.Config{
			Concurrency: 6, // 总并发数
			Queues: map[string]int{
				QueueDownload: 4, // 文件下载队列：4 个并发 worker
				QueueMain:     2, // 主任务队列：2 个并发 worker
			},
			// 错误处理
			ErrorHandler: asynq.ErrorHandlerFunc(func(ctx context.Context, task *asynq.Task, err error) {
				zap.L().Error("[fwhQueue] task failed",
					zap.String("type", task.Type()),
					zap.String("err", err.Error()),
				)
			}),
			// 重试延迟策略（指数退避）
			RetryDelayFunc: func(n int, e error, t *asynq.Task) time.Duration {
				return time.Duration(n*n) * 30 * time.Second // 30s, 2min, 4.5min
			},
		},
		Mux: mux,
	})

	zap.L().Info("[fwhQueue] initialized successfully")
}

// GetClient 获取 Asynq Client（供内部使用）
func GetClient() *asynq.Client {
	return TaskClient
}

// GetInspector 获取 Asynq Inspector（供内部使用）
func GetInspector() *asynq.Inspector {
	return inspector
}

// GetTaskInfoList 获取所有 FWH 队列中的任务列表（供外部查询）
func GetTaskInfoList() ([]dto.BackgroundTaskItem, error) {
	if inspector == nil {
		return nil, fmt.Errorf("inspector not initialized")
	}

	var result []dto.BackgroundTaskItem
	queues := []string{QueueMain, QueueDownload}

	for _, queue := range queues {
		// 查询活跃任务（running）
		activeTasks, err := inspector.ListActiveTasks(queue, asynq.PageSize(100))
		if err == nil {
			for _, t := range activeTasks {
				result = append(result, dto.BackgroundTaskItem{
					Id:      t.ID,
					Status:  "running",
					Retry:   int64(t.Retried),
					Payload: string(t.Payload),
				})
			}
		}

		// 查询等待中的任务（pending）
		pendingTasks, err := inspector.ListPendingTasks(queue, asynq.PageSize(100))
		if err == nil {
			for _, t := range pendingTasks {
				result = append(result, dto.BackgroundTaskItem{
					Id:      t.ID,
					Status:  "pending",
					Retry:   int64(t.Retried),
					Payload: string(t.Payload),
				})
			}
		}

		// 查询重试中的任务（retry）
		retryTasks, err := inspector.ListRetryTasks(queue, asynq.PageSize(100))
		if err == nil {
			for _, t := range retryTasks {
				result = append(result, dto.BackgroundTaskItem{
					Id:      t.ID,
					Status:  "retry",
					Retry:   int64(t.Retried),
					Payload: string(t.Payload),
				})
			}
		}

		// 查询已完成的任务（completed）
		completedTasks, err := inspector.ListCompletedTasks(queue, asynq.PageSize(100))
		if err == nil {
			for _, t := range completedTasks {
				result = append(result, dto.BackgroundTaskItem{
					Id:      t.ID,
					Status:  "completed",
					Retry:   int64(t.Retried),
					Payload: string(t.Payload),
				})
			}
		}

		// 查询失败的任务（archived/failed）
		archivedTasks, err := inspector.ListArchivedTasks(queue, asynq.PageSize(100))
		if err == nil {
			for _, t := range archivedTasks {
				result = append(result, dto.BackgroundTaskItem{
					Id:      t.ID,
					Status:  "failed",
					Retry:   int64(t.Retried),
					Payload: string(t.Payload),
				})
			}
		}
	}

	return result, nil
}

// newTaskID 生成固定格式的 TaskID，方便追踪
func newTaskID(taskType string, id string) string {
	return fmt.Sprintf("%s:%s", taskType, id)
}