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/go-redis/redis/v8"
)

var (
	ctx        = context.Background()
	rdb        = cache.C
	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() {
	// 刷新redis指针
	rdb = cache.C

	// 确保消费组存在（忽略已存在错误）
	_ = 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 {
					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
					}
				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 := rdb.XAck(streamInfo.Ctx, streamInfo.StreamName, streamInfo.GroupName, streamInfo.ID).Err(); 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
	}

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

	return nil
}