FranklinWH-Toolbox/service/taskQueue/taskQueue.go

package taskqueue

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

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

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

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

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

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

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

// 任务消费者（单 worker，顺序执行）
func startTaskQueueWorker() {
	for {
		res, err := rdb.XReadGroup(ctx, &redis.XReadGroupArgs{
			Group:    groupName,
			Consumer: consumerID,
			Streams:  []string{streamName, ">"},
			Count:    1,
			Block:    5 * time.Second, // 等待任务
		}).Result()

		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 {
					logger.Std.Println("unmarshal base error:", err)
					continue
				}

				switch base.Type {
				case TaskTypeSiteAnalysis:
					var payload TaskPayload[TaskSiteAnalysis]
					if err := json.Unmarshal([]byte(rawPayload), &payload); err != nil {
						log.Println("unmarshal error:", err)
						continue
					}
					logger.Std.Printf("Processing task %s: %+v\n", id, payload)
					handleTaskSiteAnalysis(payload.Payload)
				default:
					logger.Std.Println("Unknown task type:", base.Type)
				}

				// 确认任务完成
				if err := rdb.XAck(ctx, streamName, groupName, id).Err(); err != nil {
					logger.Std.Println("XAck error:", err)
				}
			}
		}
	}
}

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

	if _, err := rdb.XAdd(ctx, &redis.XAddArgs{
		Stream: streamName,
		Values: map[string]interface{}{"payload": string(data)},
	}).Result(); err != nil {
		return err
	}

	return nil
}