51 lines
1.2 KiB
Go
51 lines
1.2 KiB
Go
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
|
||
}
|