diff --git a/examples/ipratelimiter/main.go b/examples/ipratelimiter/main.go new file mode 100644 index 0000000..584b886 --- /dev/null +++ b/examples/ipratelimiter/main.go @@ -0,0 +1,34 @@ +// ipratelimiter 包示例:按 IP 维度的令牌桶限流器 +package main + +import ( + "fmt" + "time" + + "git.zeroonesoft.cn/golib/zogo/ipratelimiter" + + "golang.org/x/time/rate" +) + +func main() { + // 每秒补充 2 个令牌,桶容量 3 + limiter := ipratelimiter.NewIPRateLimiter(rate.Limit(2), 3) + + ip := "192.168.1.100" + allowed := 0 + for i := 0; i < 10; i++ { + // GetLimiter:该 IP 不存在时自动创建 + l := limiter.GetLimiter(ip) + if l.Allow() { + allowed++ + fmt.Printf("第 %2d 次请求: 放行\n", i+1) + } else { + fmt.Printf("第 %2d 次请求: 拒绝\n", i+1) + } + time.Sleep(50 * time.Millisecond) + } + fmt.Printf("共放行 %d 次\n", allowed) + + // 不同 IP 相互独立限流 + fmt.Println("另一 IP 首次请求放行:", limiter.GetLimiter("10.0.0.1").Allow()) +} diff --git a/ipratelimiter/README.md b/ipratelimiter/README.md new file mode 100644 index 0000000..2bff1b7 --- /dev/null +++ b/ipratelimiter/README.md @@ -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` 为桶容量。 diff --git a/ipratelimiter/ipRatelimiter.go b/ipratelimiter/ipRatelimiter.go new file mode 100644 index 0000000..5f388a2 --- /dev/null +++ b/ipratelimiter/ipRatelimiter.go @@ -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 +} diff --git a/ipratelimiter/ipRatelimiter_test.go b/ipratelimiter/ipRatelimiter_test.go new file mode 100644 index 0000000..fbaa10c --- /dev/null +++ b/ipratelimiter/ipRatelimiter_test.go @@ -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) + } + } +}