FranklinWH-Toolbox/service/task/queue/simCardQueue/simCardTaskQueue.go

package simCardQueue

import (
	"context"
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/repositories"
	task_core "fwh-toolbox-server/task/core"
	task_pg "fwh-toolbox-server/task/plugin"
	"time"

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

/***********************************************************************
 * 为了控制流量采用单线程队列查询
 * SyncOnePage Task * 1								查询单页数据，并创建数据聚合任务
 * 			|
 *			▼
 * SyncOnePage Task * 2								查询单页数据，并创建数据聚合任务
 * 			...
 *			▼
 * SyncOnePage Task * N								查询单页数据，并创建数据聚合任务
 * 			|
 *			▼
 * Aggregation (Group Aggregator)			聚合数据并创建，批量数据处理任务
 * 			|
 *			▼
 * Bulk Task													将聚合的数据批量写入数据库
 ***********************************************************************/

const (
	TaskTypeSimCardAggregation  = "SimCard:Aggregation"
	GroupNameSimCardAggregation = "SimCard:GroupAggregation"
	TaskTypeSimCardBulk         = "SimCard:Bulk"
)

type TaskCore interface {
	CreateTask() *asynq.Client
}

type SimCardSyncOnePagePayload struct {
	Size uint64 `json:"size"`
	Page uint64 `json:"page"`
}

type SimCardProcessor struct {
	// ... fields for struct
	infoCacheMap map[string][]domain.SimCardInfo
}

var TaskClient *asynq.Client

func InitSimCardTaskQueue(tc *task_core.TaskCore) {
	processor := &SimCardProcessor{
		infoCacheMap: make(map[string][]domain.SimCardInfo),
	}

	// 注册SmartHub收集任务处理函数
	mux := asynq.NewServeMux()
	mux.HandleFunc(TaskTypeSimCardSmartHubSyncOnePage, processor.handleSimCardSmartHubSyncOnePageTask)

	// 注册Acceleronix收集任务处理函数
	mux.HandleFunc(TaskTypeSimCardAcceleronixSyncOnePage, processor.handleSimCardAcceleronixSyncOnePageTask)

	// 注册SimCard批处理任务
	mux.HandleFunc(TaskTypeSimCardBulk, processor.handleSimCardBulkTask)

	// 注册SimCard组聚合器
	aggManager := task_pg.NewAggregatorManager()
	err := aggManager.Register(GroupNameSimCardAggregation, func(tasks []*asynq.Task) *asynq.Task {
		var aggInfoList []domain.SimCardInfo
		for _, t := range tasks {
			payloadStr := t.Payload()

			infoList := processor.infoCacheMap[string(payloadStr)]
			defer delete(processor.infoCacheMap, string(payloadStr))

			aggInfoList = append(aggInfoList, infoList...)
		}

		uuidStr := uuid.NewString()
		processor.infoCacheMap[uuidStr] = aggInfoList
		zap.L().Debug("SimCardBulk get aggInfoList.", zap.Int("len", len(aggInfoList)))

		bulkTask := asynq.NewTask(
			TaskTypeSimCardBulk,
			[]byte(uuidStr),
			asynq.MaxRetry(5),
			asynq.Queue("default"),
			asynq.Timeout(time.Minute),
			asynq.Unique(5*time.Minute),
			asynq.Retention(7*24*time.Hour), // 成功后保留7天
		)

		return bulkTask
	})
	if err != nil {
		zap.L().Error("GroupNameSimCardAggregation register failed.", zap.String("err", err.Error()))
	}

	TaskClient = tc.CreateTask(task_core.TaskConfig{
		Name: "SimCardCollectTask",
		SvrCfg: asynq.Config{
			// Specify how many concurrent workers to use
			Concurrency: 1, // 串行任务，只配置1个worker
			// Optionally specify multiple queues with different priority.
			Queues: map[string]int{
				"critical": 6,
				"default":  3,
				"low":      1,
			},
			// See the godoc for other configuration options
			GroupAggregator:  asynq.GroupAggregatorFunc(aggManager.Aggregate),
			GroupGracePeriod: 30 * time.Second, // 宽限期，30秒没人提交开始聚合
			GroupMaxDelay:    5 * time.Minute,  // 最长等待时间，5分钟必须聚合
			GroupMaxSize:     100,              // 最大数量，数据总量到达100必须聚合
		},
		Mux: mux,
		Agg: aggManager,
	})

	err = CreateSchedulerSimCardCollect(tc)
	if err != nil {
		zap.L().Error("CreateSchedulerSimCardCollect failed.", zap.String("err", err.Error()))
	}
}

func CreateSchedulerSimCardCollect(tc *task_core.TaskCore) error {
	payload := SimCardSyncOnePagePayload{
		Size: SmartHubSimCardInfoPageSize,
		Page: 0, // SmartHub从第0页开始
	}
	payloadStr, err := json.Marshal(payload)
	if err != nil {
		zap.L().Error("SimCardSyncOnePagePayload Marshal failed.", zap.String("err", err.Error()))
		return err
	}

	// 创建SmartHub收集任务
	taskSmartHub := asynq.NewTask(
		TaskTypeSimCardSmartHubSyncOnePage,
		payloadStr,
		asynq.MaxRetry(5),
		asynq.Queue("default"),
		asynq.Timeout(time.Minute),
		asynq.Unique(5*time.Minute),
		asynq.Retention(7*24*time.Hour), // 成功后保留7天
	)
	err = tc.Scheduler.RegisterUtcScheduler("SmartHubSyncOnePageScheduler", "0 23 * * *", taskSmartHub)
	if err != nil {
		zap.L().Error("TaskTypeSimCardSmartHubSyncOnePage RegisterUtcScheduler failed.", zap.String("err", err.Error()), zap.Uint64("size", payload.Size), zap.Uint64("page", payload.Page))
		return err
	}

	payload.Page = 1 // Acceleronix从第1页开始
	payloadStr, err = json.Marshal(payload)
	if err != nil {
		zap.L().Error("SimCardSyncOnePagePayload Marshal failed.", zap.String("err", err.Error()))
		return err
	}

	// 创建Acceleronix收集任务
	taskAcceleronix := asynq.NewTask(
		TaskTypeSimCardAcceleronixSyncOnePage,
		payloadStr,
		asynq.MaxRetry(5),
		asynq.Queue("default"),
		asynq.Timeout(30*time.Minute),
		asynq.Unique(5*time.Minute),
		asynq.Retention(7*24*time.Hour), // 成功后保留7天
	)
	err = tc.Scheduler.RegisterUtcScheduler("AcceleronixSyncOnePageScheduler", "0 23 * * *", taskAcceleronix)
	if err != nil {
		zap.L().Error("TaskTypeSimCardAcceleronixSyncOnePage RegisterUtcScheduler failed.", zap.String("err", err.Error()), zap.Uint64("size", payload.Size), zap.Uint64("page", payload.Page))
		return err
	}

	return nil
}

// 创建SimCard信息收集任务
func CreateTaskSimCardCollect() error {
	if TaskClient == nil {
		return fmt.Errorf("*asynq.Client Uninitialized.")
	}

	payload := SimCardSyncOnePagePayload{
		Size: SmartHubSimCardInfoPageSize,
		Page: 0, // SmartHub从第0页开始
	}
	payloadStr, err := json.Marshal(payload)
	if err != nil {
		zap.L().Error("SimCardSyncOnePagePayload Marshal failed.", zap.String("err", err.Error()))
		return err
	}

	// 创建SmartHub收集任务
	taskSmartHub := asynq.NewTask(
		TaskTypeSimCardSmartHubSyncOnePage,
		payloadStr,
		asynq.MaxRetry(5),
		asynq.Queue("default"),
		asynq.Timeout(time.Minute),
		asynq.Unique(5*time.Minute),
		asynq.Retention(7*24*time.Hour), // 成功后保留7天
	)
	_, err = TaskClient.Enqueue(taskSmartHub)
	if err != nil {
		zap.L().Error("TaskTypeSimCardSmartHubSyncOnePage Enqueue failed.", zap.String("err", err.Error()), zap.Uint64("size", payload.Size), zap.Uint64("page", payload.Page))
		return err
	}

	payload.Page = 1 // Acceleronix从第1页开始
	payloadStr, err = json.Marshal(payload)
	if err != nil {
		zap.L().Error("SimCardSyncOnePagePayload Marshal failed.", zap.String("err", err.Error()))
		return err
	}

	// 创建Acceleronix收集任务
	taskAcceleronix := asynq.NewTask(
		TaskTypeSimCardAcceleronixSyncOnePage,
		payloadStr,
		asynq.MaxRetry(5),
		asynq.Queue("default"),
		asynq.Timeout(30*time.Minute),
		asynq.Unique(5*time.Minute),
		asynq.Retention(7*24*time.Hour), // 成功后保留7天
	)
	_, err = TaskClient.Enqueue(taskAcceleronix)
	if err != nil {
		zap.L().Error("TaskTypeSimCardAcceleronixSyncOnePage Enqueue failed.", zap.String("err", err.Error()), zap.Uint64("size", payload.Size), zap.Uint64("page", payload.Page))
		return err
	}

	return nil
}

// 保存SimCard信息到数据库
func (processor *SimCardProcessor) saveSimCardInfo(infoList []domain.SimCardInfo) error {
	zap.L().Info("saveSimCardInfo", zap.Int("Num", len(infoList)))

	// 保存基本信息
	_, err := repositories.UpdateSimCardBatchInfo(infoList)
	if err != nil {
		zap.L().Error("UpdateSimCardInfo failed.", zap.String("err", err.Error()), zap.Int("infoNum", len(infoList)))
		return err
	}

	// 提取日数据
	dayDataList := []domain.SimCardDayData{}
	for _, info := range infoList {
		dayData := domain.SimCardDayData{
			SimId:        info.GICCID,                   // GICCID
			UpdateTime:   info.UpdateTime,               // UTC时间
			Month:        info.Month,                    // 统计月份 格式：yyyy-MM
			AllDataUsage: info.AllDataUsage,             // 月数据用量(KB)
			Day:          uint8(time.Now().UTC().Day()), // 当天日期值 1-31
			DailyUsage: domain.SimCardDailyUsage{
				UpdateTime: info.UpdateTime,                        // UTC时间
				Day:        time.Now().UTC().Format(time.DateOnly), // 统计日期 格式：yyyy-MM-DD
				DataUsage:  info.AllDataUsage,                      // 数据用量(KB)
			}, // 日用量数据信息
		}
		dayDataList = append(dayDataList, dayData)
	}

	// 保存月数据信息
	_, err = repositories.UpdateSimCardBatchDayData(dayDataList)
	if err != nil {
		zap.L().Error("UpdateSimCardBatchMonthlyData failed.", zap.String("err", err.Error()), zap.Int("infoNum", len(infoList)))
		return err
	}

	return nil
}

// 处理SimCard批处理任务
func (processor *SimCardProcessor) handleSimCardBulkTask(ctx context.Context, task *asynq.Task) error {
	payloadStr := task.Payload()
	originalList := processor.infoCacheMap[string(payloadStr)]
	defer delete(processor.infoCacheMap, string(payloadStr))

	// infoList := make([]domain.SimCardInfo, len(originalList))
	// copy(infoList, originalList) // 这里使用拷贝，防止引用被删除
	infoList := originalList

	return processor.saveSimCardInfo(infoList)
}