From a81e467a31569fefb756489fa43e2507a8e5311e Mon Sep 17 00:00:00 2001 From: 4566704 <4566704@qq.com> Date: Sun, 20 Sep 2026 12:25:35 +0800 Subject: [PATCH] =?UTF-8?q?feat(conn):=20=E8=87=AA=20go-hua=20=E8=BF=81?= =?UTF-8?q?=E5=85=A5=E9=99=90=E9=80=9F=E7=BB=9F=E8=AE=A1=E8=BF=9E=E6=8E=A5?= =?UTF-8?q?=E8=A3=85=E9=A5=B0=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - net.Conn 透明叠加 rate 限速与 flow 流量统计 - 内部依赖改指 zogo/flow 与 zogo/rate, 附测试、例程与 README --- conn/README.md | 41 +++++++++++++ conn/conn.go | 92 +++++++++++++++++++++++++++++ conn/conn_test.go | 132 ++++++++++++++++++++++++++++++++++++++++++ examples/conn/main.go | 61 +++++++++++++++++++ 4 files changed, 326 insertions(+) create mode 100644 conn/README.md create mode 100644 conn/conn.go create mode 100644 conn/conn_test.go create mode 100644 examples/conn/main.go diff --git a/conn/README.md b/conn/README.md new file mode 100644 index 0000000..39fa8ea --- /dev/null +++ b/conn/README.md @@ -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）。 diff --git a/conn/conn.go b/conn/conn.go new file mode 100644 index 0000000..055c731 --- /dev/null +++ b/conn/conn.go @@ -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) +} diff --git a/conn/conn_test.go b/conn/conn_test.go new file mode 100644 index 0000000..147e50b --- /dev/null +++ b/conn/conn_test.go @@ -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 +} diff --git a/examples/conn/main.go b/examples/conn/main.go new file mode 100644 index 0000000..0fc0c21 --- /dev/null +++ b/examples/conn/main.go @@ -0,0 +1,61 @@ +// conn 包示例：带流量统计/限速的连接包装（读写自动计入 flow/rate，Close 幂等） +package main + +import ( + "fmt" + "io" + "net" + "sync" + + "git.zeroonesoft.cn/golib/zogo/conn" + "git.zeroonesoft.cn/golib/zogo/flow" + "git.zeroonesoft.cn/golib/zogo/rate" +) + +func main() { + // 本地 echo 服务 + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + panic(err) + } + go func() { + for { + c, err := l.Accept() + if err != nil { + return + } + go func() { + io.Copy(c, c) + c.Close() + }() + } + }() + + raw, err := net.Dial("tcp", l.Addr().String()) + if err != nil { + panic(err) + } + + // rate.NewRate(0) 表示不限速；限速时传入 Mbps 并调用 Start() + r := rate.NewRate(0) + f := &flow.Flow{} + wg := &sync.WaitGroup{} + wg.Add(1) + + // 用 conn.Conn 包装原始连接 + c := conn.NewConn(raw, r, f, wg) + + c.Write([]byte("hello go-hua")) + buf := make([]byte, 12) + n, _ := c.Read(buf) + fmt.Println("echo:", string(buf[:n])) + + // 读写流量自动统计 + send, recv := f.Get() + fmt.Printf("连接累计流量: 写入 %d 字节, 读取 %d 字节\n", send, recv) + + // Close 只会真正关闭一次，并触发 wg.Done + c.Close() + wg.Wait() + fmt.Println("连接已关闭") +}