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 }