feat(ipratelimiter): 自 go-hua 迁入按 IP 令牌桶限流器并修正包子句
- 包子句 IPRateLimiter → ipratelimiter(符合 Go 小写命名惯例), 消费方 限定符同步; 新增依赖 golang.org/x/time - 附测试、examples/ipratelimiter 例程与包 README
This commit is contained in:
@@ -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())
|
||||||
|
}
|
||||||
@@ -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