package service import ( "context" "fmt" "time" "cid/consts" "github.com/gogf/gf/v2/frame/g" ) type rateLimit struct{} // RateLimit 限流服务 var RateLimit = new(rateLimit) // TenantRateLimitConfig 租户限流配置 type TenantRateLimitConfig struct { TenantID int64 // 租户ID RequestsPerSecond float64 // 每秒请求数 Burst int // 突发请求数 Window time.Duration // 时间窗口 } // CheckTenantRequestLimit 检查租户请求次数限制 func (s *rateLimit) CheckTenantRequestLimit(ctx context.Context, tenantID int64, config *TenantRateLimitConfig) (bool, error) { if config == nil { // 使用默认配置 config = s.getDefaultTenantRateLimitConfig(tenantID) } // 构建Redis键 - 使用当前小时的键,确保按小时计数 now := time.Now() hourKey := fmt.Sprintf("%s%d:%d", consts.AdRequestLimitKeyPrefix, tenantID, now.Hour()) // 获取当前计数 currentCountVar, err := g.Redis().Get(ctx, hourKey) if err != nil && err.Error() != "redis: nil" { return false, err } currentCount := currentCountVar.Int64() // 如果是第一次请求,设置计数和过期时间(到下一个小时) if currentCount == 0 { // 设置过期时间为到下一个小时的剩余时间 nextHour := now.Truncate(time.Hour).Add(time.Hour) ttl := nextHour.Sub(now) // 使用SetEX一次性设置值和过期时间 err = g.Redis().SetEX(ctx, hourKey, 1, int64(ttl.Seconds())) if err != nil { return false, err } return true, nil } // 检查是否超过限制 maxRequests := int64(config.RequestsPerSecond * config.Window.Seconds()) if currentCount >= maxRequests { return false, nil } // 增加计数 _, err = g.Redis().Incr(ctx, hourKey) if err != nil { return false, err } return true, nil } // GetDefaultTenantRateLimitConfig 获取默认的租户限流配置 func (s *rateLimit) getDefaultTenantRateLimitConfig(tenantID int64) *TenantRateLimitConfig { // 从配置文件中读取限流参数 ctx := context.Background() // 检查是否启用租户限流 enabled := g.Cfg().MustGet(ctx, "tenantRateLimit.enabled", false).Bool() if !enabled { // 如果未启用,返回一个很大的限制值,相当于不限制 return &TenantRateLimitConfig{ TenantID: tenantID, RequestsPerSecond: 10000, // 每秒10000个请求,相当于不限制 Burst: 20000, // 突发20000个请求 Window: time.Hour, } } // 从配置文件中获取限流参数 requestsPerHour := g.Cfg().MustGet(ctx, "tenantRateLimit.requestsPerHour", 3600).Int64() windowSeconds := g.Cfg().MustGet(ctx, "tenantRateLimit.window", 3600).Int64() burst := g.Cfg().MustGet(ctx, "tenantRateLimit.burst", 100).Int() // 转换为每秒请求数 requestsPerSecond := float64(requestsPerHour) / float64(windowSeconds) return &TenantRateLimitConfig{ TenantID: tenantID, RequestsPerSecond: requestsPerSecond, Burst: burst, Window: time.Duration(windowSeconds) * time.Second, } } // SetTenantRateLimitConfig 设置租户限流配置 func (s *rateLimit) SetTenantRateLimitConfig(ctx context.Context, config *TenantRateLimitConfig) error { // 注意:实际使用的是config.yml中的全局配置,此方法仅用于兼容旧API // 实际限流参数请修改config.yml中的tenantRateLimit部分 return nil } // GetTenantCurrentUsage 获取租户当前请求使用情况 func (s *rateLimit) GetTenantCurrentUsage(ctx context.Context, tenantID int64, config *TenantRateLimitConfig) (current int64, max int64, err error) { if config == nil { config = s.getDefaultTenantRateLimitConfig(tenantID) } // 构建当前小时的Redis键 now := time.Now() hourKey := fmt.Sprintf("%s%d:%d", consts.AdRequestLimitKeyPrefix, tenantID, now.Hour()) // 获取当前计数 currentVar, err := g.Redis().Get(ctx, hourKey) if err != nil && err.Error() == "redis: nil" { current = 0 err = nil } else if err != nil { return 0, 0, err } else { current = currentVar.Int64() } max = int64(config.RequestsPerSecond * config.Window.Seconds()) return current, max, nil }