FranklinWH-Toolbox/service/task/plugin/aggregator.go

package task_pg

import (
	"fmt"
	"sync"

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

type AggregatorFunc func(tasks []*asynq.Task) *asynq.Task

type AggregatorManager struct {
	mu          sync.RWMutex
	aggregators map[string]AggregatorFunc // 聚合器映射表
}

func NewAggregatorManager() *AggregatorManager {
	return &AggregatorManager{
		aggregators: make(map[string]AggregatorFunc),
	}
}

func (manager *AggregatorManager) Aggregate(group string, tasks []*asynq.Task) *asynq.Task {
	manager.mu.RLock()
	fn, ok := manager.aggregators[group]
	manager.mu.RUnlock()

	if !ok {
		zap.L().Error("aggregator: group not found", zap.String("group", group))
		return nil // 或记录错误
	}

	return fn(tasks)
}

func (manager *AggregatorManager) Register(group string, fn AggregatorFunc) error {
	manager.mu.Lock()
	defer manager.mu.Unlock()

	if group == "" {
		return fmt.Errorf("AggregatorManager: group is empty")
	}

	if fn == nil {
		return fmt.Errorf("AggregatorManager: nil AggregatorFunc")
	}

	if _, ok := manager.aggregators[group]; ok {
		panic(fmt.Sprintf("aggregator: group [%s] already registered", group))
	}

	manager.aggregators[group] = fn

	return nil
}

func (manager *AggregatorManager) Unregister(group string) {
	manager.mu.Lock()
	defer manager.mu.Unlock()

	delete(manager.aggregators, group)
}