feat(rate): 自 go-hua 迁入秒级带宽限速器
- 按秒配额补充令牌, Add 阻塞限速, 配合 conn 使用; 附测试、例程与 README
This commit is contained in:
@@ -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())
|
||||||
|
}
|
||||||
@@ -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` 中叠加使用。
|
||||||
+109
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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 阻塞(属预期)")
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user