feat(rate): 自 go-hua 迁入秒级带宽限速器

- 按秒配额补充令牌, Add 阻塞限速, 配合 conn 使用; 附测试、例程与 README
2026-09-20 12:25:36 +08:00
parent a0133f5881
commit 45864a61dd
4 changed files with 250 additions and 0 deletions
@@ -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` 中叠加使用。
@@ -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 阻塞（属预期）")
+	}
+}