feat(ipratelimiter): 自 go-hua 迁入按 IP 令牌桶限流器并修正包子句
- 包子句 IPRateLimiter → ipratelimiter(符合 Go 小写命名惯例), 消费方 限定符同步; 新增依赖 golang.org/x/time - 附测试、examples/ipratelimiter 例程与包 README
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
# ipratelimiter
|
||||
|
||||
按 IP 维度管理的令牌桶限流器(基于 `golang.org/x/time/rate`):为每个 IP 维护
|
||||
独立 `rate.Limiter`,典型用于 Gin 中间件按客户端 IP 限流。
|
||||
|
||||
> 迁移自 go-hua/ipratelimiter,并将包子句 `IPRateLimiter` 修正为 Go 惯例的
|
||||
> 小写 `ipratelimiter`(原包名大写不符合 Go 命名规范)。
|
||||
|
||||
## 用法
|
||||
|
||||
```go
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"git.zeroonesoft.cn/golib/zogo/ipratelimiter"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// 每 IP 每秒 2 个请求,桶容量 3
|
||||
limiterSet := ipratelimiter.NewIPRateLimiter(rate.Limit(2), 3)
|
||||
|
||||
func handler(w http.ResponseWriter, r *http.Request) {
|
||||
lim := limiterSet.GetLimiter(r.RemoteAddr) // 不存在则自动 AddIP
|
||||
if !lim.Allow() {
|
||||
http.Error(w, "too many requests", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
// ...业务
|
||||
}
|
||||
```
|
||||
|
||||
完整可运行例程:[examples/ipratelimiter/main.go](../examples/ipratelimiter/main.go)
|
||||
|
||||
## 注意
|
||||
|
||||
- 内部 map 只增不减,海量公网 IP 场景需自行加过期清理(或前置 CDN/网关限流)。
|
||||
- 限流参数与 `golang.org/x/time/rate` 语义一致:`rate.Limit` 为每秒令牌数,`b` 为桶容量。
|
||||
@@ -0,0 +1,53 @@
|
||||
// Package ipratelimiter 提供按 IP 维度管理的令牌桶限流器。
|
||||
package ipratelimiter
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// IPRateLimiter 按 IP 维度管理的令牌桶限流器
|
||||
type IPRateLimiter struct {
|
||||
ips map[string]*rate.Limiter
|
||||
mu *sync.RWMutex
|
||||
r rate.Limit
|
||||
b int
|
||||
}
|
||||
|
||||
// NewIPRateLimiter 创建限流器
|
||||
// r 为每秒令牌补充速率,b 为桶容量
|
||||
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
|
||||
i := &IPRateLimiter{
|
||||
ips: make(map[string]*rate.Limiter),
|
||||
mu: &sync.RWMutex{},
|
||||
r: r,
|
||||
b: b,
|
||||
}
|
||||
|
||||
return i
|
||||
}
|
||||
|
||||
// AddIP 为指定 IP 创建新的限流器并加入缓存
|
||||
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
limiter := rate.NewLimiter(i.r, i.b)
|
||||
i.ips[ip] = limiter
|
||||
return limiter
|
||||
}
|
||||
|
||||
// GetLimiter 返回指定 IP 的限流器,不存在时自动创建
|
||||
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
|
||||
i.mu.Lock()
|
||||
limiter, exists := i.ips[ip]
|
||||
|
||||
if !exists {
|
||||
i.mu.Unlock()
|
||||
return i.AddIP(ip)
|
||||
}
|
||||
|
||||
i.mu.Unlock()
|
||||
|
||||
return limiter
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package ipratelimiter
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
func TestNewIPRateLimiter(t *testing.T) {
|
||||
l := NewIPRateLimiter(rate.Limit(100), 100)
|
||||
if l == nil {
|
||||
t.Fatal("NewIPRateLimiter 返回 nil")
|
||||
}
|
||||
lim := l.GetLimiter("10.1.1.1")
|
||||
if lim == nil {
|
||||
t.Fatal("GetLimiter 不应返回 nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLimiterSameInstance(t *testing.T) {
|
||||
l := NewIPRateLimiter(rate.Limit(100), 100)
|
||||
|
||||
lim1 := l.GetLimiter("192.168.1.1")
|
||||
lim2 := l.GetLimiter("192.168.1.1")
|
||||
if lim1 != lim2 {
|
||||
t.Error("同一 IP 的 GetLimiter 应返回相同实例")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLimiterDifferentIP(t *testing.T) {
|
||||
l := NewIPRateLimiter(rate.Limit(100), 100)
|
||||
|
||||
lim1 := l.GetLimiter("192.168.1.1")
|
||||
lim2 := l.GetLimiter("10.0.0.1")
|
||||
if lim1 == lim2 {
|
||||
t.Error("不同 IP 的 GetLimiter 应返回不同实例")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddIP(t *testing.T) {
|
||||
l := NewIPRateLimiter(rate.Limit(100), 100)
|
||||
|
||||
lim1 := l.AddIP("172.16.0.1")
|
||||
lim2 := l.GetLimiter("172.16.0.1")
|
||||
if lim1 != lim2 {
|
||||
t.Error("AddIP 后 GetLimiter 应返回同一实例")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllow(t *testing.T) {
|
||||
// 大速率:短时间内必然放行
|
||||
l := NewIPRateLimiter(rate.Limit(1000), 1000)
|
||||
lim := l.GetLimiter("192.168.1.100")
|
||||
for i := 0; i < 10; i++ {
|
||||
if !lim.Allow() {
|
||||
t.Fatalf("第 %d 次请求被拒绝", i+1)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user