4566704
a81e467a31
- net.Conn 透明叠加 rate 限速与 flow 流量统计 - 内部依赖改指 zogo/flow 与 zogo/rate, 附测试、例程与 README
133 lines
2.6 KiB
Go
133 lines
2.6 KiB
Go
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
|
||
}
|