package limit import ( "fmt" "golang.org/x/time/rate" "sync" "time" ) type RateLimiter struct { // RateLimiterMap 用来存储每个 key 对应的限流器 rateLimiters map[string]*rate.Limiter mu sync.Mutex } type Parameters struct { Key string // 限流器的标识 Burst int // 每分钟允许的最大请求数 Interval time.Duration // 请求数量的时间窗口内 } func NewRateLimiter() *RateLimiter { return &RateLimiter{ rateLimiters: make(map[string]*rate.Limiter), } } // RateLimit 创建一个函数来检查某个 ID 的请求是否频繁 func (r *RateLimiter) RateLimit(p *Parameters) error { // 锁定访问 rateLimiters,确保线程安全 r.mu.Lock() defer r.mu.Unlock() // 如果该 id 没有对应的限流器,则创建一个新的 if _, exists := r.rateLimiters[p.Key]; !exists { // 每分钟允许 1 次请求 r.rateLimiters[p.Key] = rate.NewLimiter(rate.Every(p.Interval), p.Burst) } // 获取该 id 对应的限流器 limiter := r.rateLimiters[p.Key] // 判断请求是否被限流 if !limiter.Allow() { return fmt.Errorf("请求频率过快,请稍后再试") } // 请求通过限制,可以继续执行 return nil }