From 45864a61dddb64fb16c2ea3d1de837e93cd6fcc4 Mon Sep 17 00:00:00 2001 From: 4566704 <4566704@qq.com> Date: Sun, 20 Sep 2026 12:25:36 +0800 Subject: [PATCH] =?UTF-8?q?feat(rate):=20=E8=87=AA=20go-hua=20=E8=BF=81?= =?UTF-8?q?=E5=85=A5=E7=A7=92=E7=BA=A7=E5=B8=A6=E5=AE=BD=E9=99=90=E9=80=9F?= =?UTF-8?q?=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 按秒配额补充令牌, Add 阻塞限速, 配合 conn 使用; 附测试、例程与 README --- examples/rate/main.go | 34 +++++++++++++ rate/README.md | 28 +++++++++++ rate/rate.go | 109 ++++++++++++++++++++++++++++++++++++++++++ rate/rate_test.go | 79 ++++++++++++++++++++++++++++++ 4 files changed, 250 insertions(+) create mode 100644 examples/rate/main.go create mode 100644 rate/README.md create mode 100644 rate/rate.go create mode 100644 rate/rate_test.go diff --git a/examples/rate/main.go b/examples/rate/main.go new file mode 100644 index 0000000..8129f27 --- /dev/null +++ b/examples/rate/main.go @@ -0,0 +1,34 @@ +// rate 包示例：带宽限速（按秒配额，超限时 Add 阻塞等待），常配合 conn 包做连接限速 +package main + +import ( + "fmt" + "time" + + "git.zeroonesoft.cn/golib/zogo/rate" +) + +func main() { + // 10Mbps 带宽（内部换算为 1310720 字节/秒的配额） + r := rate.NewRate(10) + r.Start() // 启动每秒配额重置 + + // 模拟写入流量：累计超过每秒配额后 Add 会阻塞，直到下一个配额周期 + total := 0 + for i := 0; i < 5; i++ { + size := 512 * 1024 // 512KB + start := time.Now() + r.Add(size) + total += size + fmt.Printf("写入 %dKB，本次等待 %v，当前秒已用 %d 字节\n", + size/1024, time.Since(start).Round(time.Millisecond), r.GetNow()) + } + fmt.Println("累计写入:", total, "字节") + + // 运行期间可动态调整限速（0 表示不限制） + r.SetLimit(0) + fmt.Println("已取消限速") + + // 取出并清零当前秒用量 + fmt.Println("ResetMax:", r.ResetMax()) +} diff --git a/rate/README.md b/rate/README.md new file mode 100644 index 0000000..e65cd1d --- /dev/null +++ b/rate/README.md @@ -0,0 +1,28 @@ +# rate + +按秒配额的带宽限速器：为读写字节流提供令牌补充式限速（字节数/秒）， +常配合 `conn` 包使用。 + +> 迁移自 go-hua/rate，代码保持原样。 + +## 用法 + +```go +import "git.zeroonesoft.cn/golib/zogo/rate" + +r := rate.NewRate(512 * 1024) // 512 KB/s；0 表示不限速 +r.Start() // 启动每秒配额补充 +defer r.Stop() + +r.Add(n) // 每次读/写后上报本秒已消耗字节数，超配额时 Add 阻塞到下个周期 + +r.SetLimit(1024 * 1024) // 动态调整限速 +used := r.GetNow() // 本秒已消耗 +``` + +完整可运行例程：[examples/rate/main.go](../examples/rate/main.go) + +## 注意 + +- `NewRate(0)` 为不限速模式，`Add` 直接放行。 +- 限速粒度为秒级令牌桶，短时突发精度有限；逐连接限速在 `conn` 中叠加使用。 diff --git a/rate/rate.go b/rate/rate.go new file mode 100644 index 0000000..59d83f4 --- /dev/null +++ b/rate/rate.go @@ -0,0 +1,109 @@ +// Package rate 提供按秒配额的带宽限速，常配合 conn 包使用。 +package rate + +import ( + "sync/atomic" + "time" +) + +// Rate 带宽限速器：每秒重置配额，配额用尽时 Add 阻塞等待 +type Rate struct { + Limit int32 `json:"limit"` // 限制速度 字节 0为不限制 + Surplus int32 `json:"surplus"` // 剩余流量 字节 + Now int32 `json:"now"` // 接收(下载/下行) 当前流量 字节 + Max int32 `json:"max"` // 接收(下载/下行) 最大流量 字节 + IsStop chan bool `json:"-"` +} + +// NewRate 创建限速器，bandwidth 为带宽上限（单位 Mbps） +func NewRate(bandwidth int) *Rate { + // 带宽 应该是Mbps 要换算成Mbyte + limit := bandwidth * 1024 * 1024 / 8 + r := new(Rate) + r.Limit = int32(limit) + r.Now = 0 + r.Surplus = int32(limit) + r.IsStop = make(chan bool) // 初始化停止通道，否则 Stop() 会 panic + return r +} + +// SetLimit 动态调整带宽上限，bandwidth 为 Mbps（0 表示不限速） +func (r *Rate) SetLimit(bandwidth int) { + // 带宽 应该是Mbps 要换算成Mbyte + limit := 0 + if bandwidth > 0 { + limit = bandwidth * 1024 * 1024 / 8 + } else { + limit = 0 + } + + atomic.StoreInt32(&r.Limit, int32(limit)) +} + +// Start 启动每秒配额重置协程 +func (r *Rate) Start() { + go r.proc() +} + +// Stop 停止配额重置协程 +func (r *Rate) Stop() { + r.IsStop <- true +} + +// proc 配额重置循环 +func (r *Rate) proc() { + ticker := time.NewTicker(time.Second * 1) + + for { + select { + case <-ticker.C: + r.reset() + case <-r.IsStop: + ticker.Stop() + return + } + } +} + +// reset 统计上一秒用量并重置配额 +func (r *Rate) reset() { + n := r.Limit - atomic.LoadInt32(&r.Surplus) + atomic.StoreInt32(&r.Now, n) + atomic.StoreInt32(&r.Surplus, r.Limit) + + if n > 0 { + now := atomic.LoadInt32(&r.Now) + if now > 0 && now > atomic.LoadInt32(&r.Max) { + atomic.StoreInt32(&r.Max, now) + } + } + //fmt.Printf("now:%d limit:%d Surplus:%d \n", n, r.Limit, r.Limit) +} + +// GetNow 获取当前秒已用量（字节） +func (r *Rate) GetNow() int { + n := atomic.LoadInt32(&r.Now) + return int(n) +} + +// ResetMax 取出并清零当前秒用量 +func (r *Rate) ResetMax() int { + n := atomic.SwapInt32(&r.Now, 0) + return int(n) +} + +// Add 计入本秒用量；配额用尽时阻塞等待下一秒配额释放 +func (r *Rate) Add(size int) { + if atomic.LoadInt32(&r.Surplus) > 0 || atomic.LoadInt32(&r.Limit) == 0 { + atomic.AddInt32(&r.Surplus, -int32(size)) + return + } + for { + //fmt.Println("等待") + time.Sleep(time.Millisecond * 10) + if atomic.LoadInt32(&r.Surplus) > 0 { + atomic.AddInt32(&r.Surplus, -int32(size)) + return + } + } +} diff --git a/rate/rate_test.go b/rate/rate_test.go new file mode 100644 index 0000000..611aa90 --- /dev/null +++ b/rate/rate_test.go @@ -0,0 +1,79 @@ +package rate + +import ( + "sync/atomic" + "testing" + "time" +) + +func TestNewRateLimitConvert(t *testing.T) { + // 10 Mbps = 10*1024*1024/8 = 1310720 字节/秒 + r := NewRate(10) + if r.Limit != 1310720 { + t.Errorf("Limit = %d, want 1310720", r.Limit) + } + if r.Surplus != 1310720 { + t.Errorf("初始 Surplus = %d, want 1310720", r.Surplus) + } +} + +func TestSetLimit(t *testing.T) { + r := NewRate(10) + r.SetLimit(0) + if r.Limit != 0 { + t.Errorf("SetLimit(0) 后 Limit = %d, want 0", r.Limit) + } + r.SetLimit(1) + if r.Limit != 1*1024*1024/8 { + t.Errorf("SetLimit(1) 后 Limit = %d, want %d", r.Limit, 1*1024*1024/8) + } +} + +func TestAddWithoutLimit(t *testing.T) { + r := NewRate(0) // 不限速 + done := make(chan struct{}) + go func() { + r.Add(1 << 20) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Error("不限速时 Add 阻塞了") + } +} + +func TestAddConsumeQuota(t *testing.T) { + r := NewRate(1) // 1 Mbps = 131072 字节/秒 + r.Add(1000) + if got := atomic.LoadInt32(&r.Surplus); got != 131072-1000 { + t.Errorf("Add 后 Surplus = %d, want %d", got, 131072-1000) + } +} + +func TestStartStopReset(t *testing.T) { + r := NewRate(1) + r.Start() + + r.Add(500) + // 等待配额重置（每秒一次） + time.Sleep(1100 * time.Millisecond) + if got := atomic.LoadInt32(&r.Surplus); got != 131072 { + t.Errorf("重置后 Surplus = %d, want 131072", got) + } + + // Stop 应正常返回（IsStop 通道已在 NewRate 中初始化） + r.Stop() + + // 重复 Stop 不应 panic：第二次发送时 proc 协程已退出，会阻塞，故放入独立协程 + done := make(chan struct{}) + go func() { + defer close(done) + r.Stop() + }() + select { + case <-done: + case <-time.After(time.Second): + t.Log("proc 协程已退出，第二次 Stop 阻塞（属预期）") + } +}