feat(conn): 自 go-hua 迁入限速统计连接装饰器

- net.Conn 透明叠加 rate 限速与 flow 流量统计
- 内部依赖改指 zogo/flow 与 zogo/rate, 附测试、例程与 README
This commit is contained in:
2026-09-20 12:25:35 +08:00
parent f9f7c7c92f
commit a81e467a31
4 changed files with 326 additions and 0 deletions
+41
View File
@@ -0,0 +1,41 @@
# conn
`net.Conn` 装饰器:在读写路径上透明叠加带宽限速(rate)与流量统计(flow),
可选关联 `sync.WaitGroup` 管理连接生命周期。
> 迁移自 go-hua/conn,内部依赖指向本仓库 flow/rate。
## 用法
```go
import (
"net"
"sync"
"git.zeroonesoft.cn/golib/zogo/conn"
"git.zeroonesoft.cn/golib/zogo/flow"
"git.zeroonesoft.cn/golib/zogo/rate"
)
var wg sync.WaitGroup
wg.Add(1)
f := &flow.Flow{}
r := rate.NewRate(1024 * 512) // 512 KB/s,0 表示不限速
r.Start()
c := conn.NewConn(rawConn, r, f, &wg)
defer c.Close()
// 之后像普通 net.Conn 一样使用:Read/Write 已自动限速并累计流量
c.Write(buf)
n, _ := c.Read(buf)
send, recv := f.Get() // 读写流量统计
```
完整可运行例程:[examples/conn/main.go](../examples/conn/main.go)
## 注意
- `rate` 传 nil 表示不限速;`flow` 传 nil 表示不统计。
- `Close` 幂等,并会 `wg.Done()`(前提是创建时传入了 wg)。
+92
View File
@@ -0,0 +1,92 @@
package conn
import (
"git.zeroonesoft.cn/golib/zogo/flow"
"git.zeroonesoft.cn/golib/zogo/rate"
"net"
"sync"
"time"
)
// 主要就是为了统计 带宽,流量
// Conn 带流量统计与限速的连接包装,实现 net.Conn 接口
type Conn struct {
Conn net.Conn // 多路复用流
Rate *rate.Rate //带宽
Flow *flow.Flow //流量
wg *sync.WaitGroup
once sync.Once
}
// NewConn 包装原始连接
// rate/flow 可为 nil(不统计不限速);wg 非 nil 时 Close 会触发一次 wg.Done
func NewConn(conn net.Conn, rate *rate.Rate, flow *flow.Flow, wg *sync.WaitGroup) *Conn {
c := new(Conn)
c.Conn = conn
c.Rate = rate
c.Flow = flow
c.wg = wg
return c
}
// Read 读取数据,读到的字节数自动计入限速与流量统计
func (c *Conn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
if c.Rate != nil {
c.Rate.Add(n)
}
if c.Flow != nil {
c.Flow.Add(int64(n), 0)
}
return
}
// Write 写入数据,写出的字节数自动计入限速与流量统计
func (c *Conn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
if c.Rate != nil {
c.Rate.Add(n)
}
if c.Flow != nil {
c.Flow.Add(int64(n), 0)
}
return
}
// Close 关闭连接(幂等,只执行一次),并触发一次 wg.Done
func (c *Conn) Close() (err error) {
// 只执行一次
c.once.Do(func() {
if c.wg != nil {
c.wg.Done()
}
err = c.Conn.Close()
})
return
}
// LocalAddr 返回本地网络地址
func (c *Conn) LocalAddr() net.Addr {
return c.Conn.LocalAddr()
}
// RemoteAddr 返回对端网络地址
func (c *Conn) RemoteAddr() net.Addr {
return c.Conn.RemoteAddr()
}
// SetDeadline 设置读写超时时间
func (c *Conn) SetDeadline(t time.Time) error {
return c.Conn.SetDeadline(t)
}
// SetReadDeadline 设置读超时时间
func (c *Conn) SetReadDeadline(t time.Time) error {
return c.Conn.SetReadDeadline(t)
}
// SetWriteDeadline 设置写超时时间
func (c *Conn) SetWriteDeadline(t time.Time) error {
return c.Conn.SetWriteDeadline(t)
}
+132
View File
@@ -0,0 +1,132 @@
package conn
import (
"io"
"net"
"sync"
"testing"
"time"
"git.zeroonesoft.cn/golib/zogo/flow"
"git.zeroonesoft.cn/golib/zogo/rate"
)
func TestConnWriteReadFlow(t *testing.T) {
c1, c2 := net.Pipe()
defer c2.Close()
f := &flow.Flow{}
wg := &sync.WaitGroup{}
wg.Add(1)
c := NewConn(c1, nil, f, wg)
defer c.Close()
// 持续消费对端数据,避免 Write 阻塞
done := make(chan struct{})
go func() {
buf := make([]byte, 1024)
for {
if _, err := c2.Read(buf); err != nil {
close(done)
return
}
}
}()
if n, err := c.Write([]byte("hello")); err != nil || n != 5 {
t.Fatalf("Write = %d, %v, want 5, nil", n, err)
}
// 对端写入,本端读取
go func() {
c2.Write([]byte("abcde"))
}()
buf := make([]byte, 5)
if err := c.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
t.Logf("SetReadDeadline: %v(conn.Read 未实现 Deadline 转发)", err)
}
n, err := c.Read(buf)
if err != nil || n != 5 {
t.Fatalf("Read = %d, %v, want 5, nil", n, err)
}
// conn 包设计:Read/Write 均计入 SendFlow
send, _ := f.Get()
if send != 10 {
t.Errorf("Flow SendFlow = %d, want 10(读写各计 5)", send)
}
// 触发消费者协程退出:关闭本端后其阻塞在 c2.Read 的 goroutine 会返回
c.Close()
<-done
}
func TestConnCloseIdempotent(t *testing.T) {
c1, c2 := net.Pipe()
defer c2.Close()
wg := &sync.WaitGroup{}
wg.Add(1)
c := NewConn(c1, nil, nil, wg)
if err := c.Close(); err != nil {
t.Fatalf("第一次 Close: %v", err)
}
_ = c.Close() // 第二次 Close 不应重复 wg.Done
finished := make(chan struct{})
go func() {
wg.Wait()
close(finished)
}()
select {
case <-finished:
// 正常:wg 只被 Done 一次
case <-time.After(time.Second):
t.Error("Close 幂等性失败:wg 被多次 Done 或未被 Done")
}
}
func TestConnAddr(t *testing.T) {
c1, c2 := net.Pipe()
defer c1.Close()
defer c2.Close()
c := NewConn(c1, nil, nil, nil)
defer c.Close()
if c.LocalAddr() == nil || c.RemoteAddr() == nil {
t.Error("LocalAddr/RemoteAddr 不应为 nil")
}
}
func TestConnWithRateLimit(t *testing.T) {
c1, c2 := net.Pipe()
defer c2.Close()
// 不限速:Add 应立即返回
r := rate.NewRate(0)
wg := &sync.WaitGroup{}
wg.Add(1)
c := NewConn(c1, r, nil, wg)
defer c.Close()
readerDone := make(chan struct{})
go func() {
defer close(readerDone)
buf := make([]byte, 1024)
for {
if _, err := c2.Read(buf); err != nil {
return
}
}
}()
if _, err := io.WriteString(c, "x"); err != nil {
t.Fatalf("Write: %v", err)
}
// 关闭本端使消费者协程退出,避免 goroutine 泄漏
_ = c.Close()
<-readerDone
}