FranklinWH-Toolbox/service/taskQueue/taskQueue.go

package taskQueue

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

	"github.com/redis/go-redis/v9"
)

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

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 InitTaskQueue() {
	// 确保消费组存在（忽略已存在错误）
	_ = 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 {
				continue
			}
			logger.Std.Println("XReadGroup error:", err)
			time.Sleep(time.Second)
			continue
		}

		for _, stream := range res {
			for _, message := range stream.Messages {
				id := message.ID
				rawPayload := message.Values["payload"].(string)

				// 先反序列化成 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()))
					continue
				}

				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()))
						continue
					}
					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()))
						continue
					}
				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()))
						continue
					}
					logger.Std.Printf("Processing task %s: %+v\n", id, payload)
					if err := handleTaskAlarmStatistics(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()))
						continue
					}
				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.XAck(streamInfo.Ctx, streamInfo.StreamName, streamInfo.GroupName, streamInfo.ID); err != nil {
			logger.Std.Println("XAck error:", err)
		}
	} else {
		// 任务失败
		logger.Std.Println("Task queue execution failed, error:", err)
	}
}

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
}