zogo/conn/conn.go

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)
}