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

package simCardQueue

import (
	"context"
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/config"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/repositories"
	task_core "fwh-toolbox-server/task/core"
	task_pg "fwh-toolbox-server/task/plugin"
	"sync"
	"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"`
	Date string `json:"date,omitempty"` // UTC collection date
}

type simCardPageBatch struct {
	Date  string
	Pages []domain.SimCardRunPage
	Info  []domain.SimCardInfo
}

type SimCardProcessor struct {
	cacheMu      sync.Mutex
	infoCacheMap map[string][]simCardPageBatch
	webhook      string
	secret       string
}

func (processor *SimCardProcessor) putBatches(key string, batches []simCardPageBatch) {
	processor.cacheMu.Lock()
	defer processor.cacheMu.Unlock()
	processor.infoCacheMap[key] = batches
}

func (processor *SimCardProcessor) getBatches(key string) ([]simCardPageBatch, bool) {
	processor.cacheMu.Lock()
	defer processor.cacheMu.Unlock()
	batches, ok := processor.infoCacheMap[key]
	return batches, ok
}

func (processor *SimCardProcessor) deleteBatches(key string) {
	processor.cacheMu.Lock()
	defer processor.cacheMu.Unlock()
	delete(processor.infoCacheMap, key)
}

var TaskClient *asynq.Client

func InitSimCardTaskQueue(tc *task_core.TaskCore, appConfig *config.AppConfig) {
	processor := &SimCardProcessor{
		infoCacheMap: make(map[string][]simCardPageBatch),
		webhook:      appConfig.SimCardDingTalkWebhook,
		secret:       appConfig.SimCardDingTalkSecret,
	}

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

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

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

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

			pageBatches, ok := processor.getBatches(string(payloadStr))
			if !ok {
				zap.L().Error("SimCard aggregation payload missing", zap.String("key", string(payloadStr)))
				continue
			}
			defer processor.deleteBatches(string(payloadStr))

			batches = append(batches, pageBatches...)
		}

		uuidStr := uuid.NewString()
		processor.putBatches(uuidStr, batches)
		zap.L().Debug("SimCardBulk get batches.", zap.Int("len", len(batches)))

		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, date time.Time) 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(date.Day()), // 当天日期值 1-31
			DailyUsage: domain.SimCardDailyUsage{
				UpdateTime: info.UpdateTime,            // UTC时间
				Day:        date.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()
	batches, ok := processor.getBatches(string(payloadStr))
	if !ok {
		return fmt.Errorf("SIM card bulk payload missing: %s", string(payloadStr))
	}
	byDate := make(map[string]*simCardPageBatch)
	for _, batch := range batches {
		combined := byDate[batch.Date]
		if combined == nil {
			combined = &simCardPageBatch{Date: batch.Date}
			byDate[batch.Date] = combined
		}
		combined.Pages = append(combined.Pages, batch.Pages...)
		combined.Info = append(combined.Info, batch.Info...)
	}
	for dateStr, batch := range byDate {
		date, err := time.Parse(time.DateOnly, dateStr)
		if err != nil {
			return err
		}
		if err := processor.saveSimCardInfo(batch.Info, date); err != nil {
			return err
		}
		if err := repositories.MarkSimCardPagesSaved(dateStr, batch.Pages); err != nil {
			return err
		}
		if err := processor.enqueueSimCardReportIfComplete(dateStr); err != nil {
			return err
		}
	}
	processor.deleteBatches(string(payloadStr))
	return nil
}