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

package fwhQueue

import (
	"context"
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/cache"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
	"fwh-toolbox-server/logger"
	"slices"
	"time"

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
)

var (
	ctx        = context.Background()
	streamName = "fwhtoolbox:taskQueue:taskStream"
	groupName  = "task_queue"
	consumerID = "worker-taskQueue"

	// retryConsumerID = "retry-taskQueue"
	retryConsumerID = consumerID
	maxRetry        = 3
	idleTime        = 10 * time.Minute

	// 状态管理相关变量
	runningID      = ""                 // 当前执行的任务ID
	failedTaskList = []string{}         // 失败且未处理任务ID列表
	errorTaskList  = []redis.XMessage{} // 超过重试次数仍失败的任务列表

	// 唯一任务key前缀
	uniqueTaskKeyPrefix = "fwhtoolbox:taskQueue:taskUnique:"
)

type TaskStreamInfo struct {
	Ctx        context.Context
	StreamName string
	GroupName  string
	ID         string
}

type TaskPayloadBase struct {
	Type string `json:"type"`
}

type TaskPayload[T any] struct {
	TaskPayloadBase
	Payload T `json:"payload"`
}

func InitFwhTaskQueue() {
	// 确保消费组存在（忽略已存在错误）
	_ = cache.XGroupCreate(ctx, streamName, groupName, "$")

	// 启动后台任务队列（不会阻塞主线程）
	go startTaskQueueWorker()
}

// 任务消费者（单 worker，顺序执行）
func startTaskQueueWorker() {
	streamsMap := make(map[string]string)
	streamsMap[streamName] = ">"

	for {
		// 处理新任务
		res, err := cache.XReadGroup(ctx, groupName, consumerID, streamsMap, 1) // 等待任务
		if err != nil {
			if err == redis.Nil {
				// 未读取到新消息
				// 处理滞留的消息
				retryRes, _, err := cache.XAutoClaim(ctx, streamName, groupName, retryConsumerID, idleTime, 1)
				if err != nil {
					if err == redis.Nil {
						continue
					}
					logger.Std.Println("XAutoClaim error:", err)
					time.Sleep(time.Second)
					continue
				}

				for _, msg := range retryRes {
					pendingRes, err := cache.XPendingExtOne(ctx, streamName, groupName, retryConsumerID, msg.ID)
					if err != nil {
						logger.Std.Println("XPendingExtOne error:", err)
					} else if pendingRes == nil {
						panic("无法读取XPendingExtOne")
					}

					// 判断重试次数
					if pendingRes.RetryCount > int64(maxRetry) {
						if slices.Contains(failedTaskList, msg.ID) {
							failedTaskList = slices.DeleteFunc(failedTaskList, func(s string) bool { return s == msg.ID })
						}
						errorTaskList = append(errorTaskList, msg)
						ackTask(TaskStreamInfo{
							Ctx:        ctx,
							StreamName: streamName,
							GroupName:  groupName,
							ID:         msg.ID,
						}, true, fmt.Errorf("任务超过最大重试次数。"))
						continue
					}

					zap.L().Warn("taskQueue任务重试。", zap.String("ID", msg.ID), zap.Int64("count", pendingRes.RetryCount))

					runningID = msg.ID
					handleTaskQueueMsg(msg)
				}
				continue
			}
			logger.Std.Println("XReadGroup error:", err)
			time.Sleep(time.Second)
			continue
		}

		for _, stream := range res {
			for _, message := range stream.Messages {
				runningID = message.ID
				handleTaskQueueMsg(message)
			}
		}
	}
}

// 处理单条消息
func handleTaskQueueMsg(message redis.XMessage) {
	id := message.ID
	rawPayload := message.Values["payload"].(string)

	var rawLockKey *string
	var rawLockToken *string

	if val, ok := message.Values["uniqueLockKey"]; ok {
		if str, ok := val.(string); ok {
			rawLockKey = &str
		}
	}

	if val, ok := message.Values["uniqueLockToken"]; ok {
		if str, ok := val.(string); ok {
			rawLockToken = &str
		}
	}

	defer func() {
		if rawLockKey != nil && rawLockToken != nil {
			// 释放唯一任务锁
			FinishUniqueTask(*rawLockKey, *rawLockToken)
		}
	}()

	// 先反序列化成 map 确定 type
	var base TaskPayloadBase
	if err := json.Unmarshal([]byte(rawPayload), &base); err != nil {
		ackTask(TaskStreamInfo{
			Ctx:        ctx,
			StreamName: streamName,
			GroupName:  groupName,
			ID:         id,
		}, false, fmt.Errorf("unmarshal base error - %s", err.Error()))
		return
	}

	switch base.Type {
	case TaskTypeSiteAnalysis: // 站点分析任务
		var payload TaskPayload[domain.TaskSiteAnalysis]
		if err := json.Unmarshal([]byte(rawPayload), &payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("unmarshal payload error - %s", err.Error()))
			return
		}
		logger.Std.Printf("Processing task %s: %+v\n", id, payload)
		if err := handleTaskSiteAnalysis(ctx, payload.Payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("handle task SiteAnalysis error - %s", err.Error()))
			return
		}
	case TaskTypeShutdownStatistics: // Shutdown统计任务
		var payload TaskPayload[domain.TaskShutdownStatistics]
		if err := json.Unmarshal([]byte(rawPayload), &payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("unmarshal payload error - %s", err.Error()))
			return
		}
		logger.Std.Printf("Processing task %s: %+v\n", id, payload)
		if err := handleTaskShutdownStatistics(ctx, payload.Payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("handle task SiteAnalysis error - %s", err.Error()))
			return
		}
	case TaskTypeBatchDownload: // 批量数据下载任务
		var payload TaskPayload[domain.TaskBatchDownload]
		if err := json.Unmarshal([]byte(rawPayload), &payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("unmarshal payload error - %s", err.Error()))
			return
		}
		logger.Std.Printf("Processing task %s: %+v\n", id, payload)
		if err := handleTaskBatchDownload(ctx, payload.Payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("handle task BatchDownload error - %s", err.Error()))
			return
		}
	case TaskTypeSiteInfoUpdate:
		var payload TaskPayload[domain.TaskSiteInfoUpdate]
		if err := json.Unmarshal([]byte(rawPayload), &payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("unmarshal payload error - %s", err.Error()))
			return
		}
		logger.Std.Printf("Processing task %s: %+v\n", id, payload)
		if err := handleTaskSiteInfoUpdate(ctx, payload.Payload); err != nil {
			ackTask(TaskStreamInfo{
				Ctx:        ctx,
				StreamName: streamName,
				GroupName:  groupName,
				ID:         id,
			}, false, fmt.Errorf("handle task SiteInfoUpdate error - %s", err.Error()))
			return
		}
	default:
		logger.Std.Println("Unknown task type:", base.Type)
	}

	// 任务完成
	ackTask(TaskStreamInfo{
		Ctx:        ctx,
		StreamName: streamName,
		GroupName:  groupName,
		ID:         id,
	}, true, nil)
}

func ackTask(streamInfo TaskStreamInfo, success bool, err error) {
	if success {
		// 确认任务完成
		if err := cache.XAckDel(streamInfo.Ctx, streamInfo.StreamName, streamInfo.GroupName, streamInfo.ID); err != nil {
			logger.Std.Println("XAck error:", err)
		}
	} else {
		if !slices.Contains(failedTaskList, streamInfo.ID) {
			failedTaskList = append(failedTaskList, streamInfo.ID)
		}
		// 任务失败
		logger.Std.Println("Task queue execution failed, error:", err)
	}

	if err != nil {
		zap.L().Error("taskQueue ackTask error.", zap.String("err", err.Error()))
	}
}

func EnqueueTask[T any](payload TaskPayload[T]) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	if _, err := cache.XAdd(ctx, streamName, map[string]interface{}{"payload": string(data)}); err != nil {
		return err
	}

	return nil
}

func EnqueueUniqueTask[T any](payload TaskPayload[T]) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	key := uniqueTaskKeyPrefix + payload.Type

	lock, ok, err := cache.TryLock(ctx, key, 2*time.Hour)
	if err != nil {
		return err
	}

	if !ok {
		return fmt.Errorf("task type already exists: %s", payload.Type)
	}

	if _, err := cache.XAdd(ctx, streamName, map[string]interface{}{"payload": string(data), "uniqueLockKey": lock.GetKey(), "uniqueLockToken": lock.GetToken()}); err != nil {
		lock.Release(ctx)
		return err
	}

	return nil
}

func FinishUniqueTask(key string, token string) {
	cache.LockRelease(ctx, key, token)
}

func GetTaskInfoList() ([]dto.BackgroundTaskItem, error) {
	c, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()

	allMsg, err := cache.XRangeN(c, streamName, "-", "+", 1000) // 查询所有信息，限制1000条
	if err != nil {
		zap.L().Error("XRangeN查询失败", zap.String("err", err.Error()))
		return nil, err
	}

	res := []dto.BackgroundTaskItem{}
	for _, msg := range append(errorTaskList, allMsg...) {
		// 判断任务状态
		taskStatus := "pending" // 默认为pending
		if msg.ID == runningID {
			taskStatus = "running"
		} else if slices.Contains(failedTaskList, msg.ID) {
			taskStatus = "failed"
		}

		// 获取任务重试次数
		var retryCount int64
		pendingRes, err := cache.XPendingExtOne(ctx, streamName, groupName, retryConsumerID, msg.ID)
		if err != nil {
			zap.L().Error("XPendingExtOne查询失败", zap.String("err", err.Error()))
			return nil, err
		} else if pendingRes == nil {
			retryCount = 0
		} else {
			retryCount = pendingRes.RetryCount
		}

		// 获取任务参数
		payload, ok := msg.Values["payload"].(string)
		if !ok {
			payload = fmt.Sprint(msg.Values["taskId"])
		}
		item := dto.BackgroundTaskItem{
			Id:      msg.ID,
			Status:  taskStatus,
			Retry:   retryCount,
			Payload: payload,
		}

		res = append(res, item)
	}
	return res, nil
}