junhong_cmp_fiber/internal/infrastructure/gatewaylimit/groups.go

package gatewaylimit

import (
	"context"
	"sync"
	"time"

	"github.com/bytedance/sonic"
	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// defaultGroupCode 是未被任何启用限流组覆盖时使用的默认组标识（每个账户每个操作独立窗口）。
const defaultGroupCode = "default"

// GroupRule 是一个限流组对某个操作生效的规则快照。
type GroupRule struct {
	GroupCode    string
	MaxQPS       int
	ReservedQPS  int
	HighPriority bool // 该操作在组内是否为高优先级
}

// groupIndex 是启用限流组按"账户 → 操作"建立的索引。
type groupIndex struct {
	byCarrier map[uint]map[string]GroupRule
	global    map[string]GroupRule
}

// GroupLoader 读取全部启用的限流组。
type GroupLoader interface {
	ListEnabled(ctx context.Context) ([]model.CarrierRateLimitGroup, error)
}

// GroupProvider 在进程内缓存启用限流组，超过 TTL 重新加载，后台修改数秒内生效。
type GroupProvider struct {
	loader GroupLoader
	logger *zap.Logger
	mu     sync.Mutex
	index  groupIndex
	loaded time.Time
}

// NewGroupProvider 创建限流组提供者；loader 为 nil 时所有操作使用默认 QPS。
func NewGroupProvider(loader GroupLoader, logger *zap.Logger) *GroupProvider {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &GroupProvider{loader: loader, logger: logger, index: groupIndex{byCarrier: map[uint]map[string]GroupRule{}, global: map[string]GroupRule{}}}
}

// Invalidate 使缓存立即失效，后台保存限流组后调用。
func (p *GroupProvider) Invalidate() {
	if p == nil {
		return
	}
	p.mu.Lock()
	p.loaded = time.Time{}
	p.mu.Unlock()
}

// Resolve 返回某账户某操作的账户组规则（未覆盖时 ok=false）与全局组规则（未配置时 ok=false）。
func (p *GroupProvider) Resolve(ctx context.Context, carrierID uint, op string) (account GroupRule, accountOK bool, global GroupRule, globalOK bool) {
	index := p.current(ctx)
	if rules, ok := index.byCarrier[carrierID]; ok {
		account, accountOK = rules[op]
	}
	global, globalOK = index.global[op]
	return
}

func (p *GroupProvider) current(ctx context.Context) groupIndex {
	if p == nil {
		return groupIndex{}
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.loader == nil || time.Since(p.loaded) < constants.GatewayRateLimitConfigTTL {
		return p.index
	}
	rows, err := p.loader.ListEnabled(ctx)
	p.loaded = time.Now()
	if err != nil {
		p.logger.Warn("加载网关限流组失败，沿用上一次配置", zap.Error(err))
		return p.index
	}
	p.index = buildGroupIndex(rows, p.logger)
	return p.index
}

func buildGroupIndex(rows []model.CarrierRateLimitGroup, logger *zap.Logger) groupIndex {
	index := groupIndex{byCarrier: map[uint]map[string]GroupRule{}, global: map[string]GroupRule{}}
	for _, row := range rows {
		ops, err := DecodeOperations(row.Operations)
		if err != nil {
			logger.Warn("限流组操作列表非法，已跳过该组", zap.Uint("group_id", row.ID), zap.Error(err))
			continue
		}
		highOps, _ := DecodeOperations(row.HighPriorityOperations)
		high := map[string]bool{}
		for _, op := range highOps {
			high[op] = true
		}
		target := index.global
		if row.CarrierID != nil {
			if index.byCarrier[*row.CarrierID] == nil {
				index.byCarrier[*row.CarrierID] = map[string]GroupRule{}
			}
			target = index.byCarrier[*row.CarrierID]
		}
		for _, op := range ops {
			target[op] = GroupRule{GroupCode: row.GroupCode, MaxQPS: row.MaxQPS, ReservedQPS: row.ReservedQPS, HighPriority: high[op]}
		}
	}
	return index
}

// DecodeOperations 解析 JSON 字符串数组形式的操作列表。
func DecodeOperations(raw []byte) ([]string, error) {
	if len(raw) == 0 {
		return nil, nil
	}
	var ops []string
	if err := sonic.Unmarshal(raw, &ops); err != nil {
		return nil, err
	}
	return ops, nil
}