Files
EdgeNode/internal/utils/conns/conn_no_stat.go
2023-04-21 15:08:44 +08:00

73 lines
1.6 KiB
Go

// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package connutils
import (
"github.com/TeaOSLab/EdgeNode/internal/zero"
"net"
"sync"
"time"
)
// 记录不需要带宽统计的连接
// 比如本地的清理和预热
var noStatAddrMap = map[string]zero.Zero{} // addr => Zero
var noStatLocker = &sync.RWMutex{}
// IsNoStatConn 检查是否为不统计连接
func IsNoStatConn(addr string) bool {
noStatLocker.RLock()
_, ok := noStatAddrMap[addr]
noStatLocker.RUnlock()
return ok
}
type NoStatConn struct {
rawConn net.Conn
}
func NewNoStat(rawConn net.Conn) net.Conn {
noStatLocker.Lock()
noStatAddrMap[rawConn.LocalAddr().String()] = zero.New()
noStatLocker.Unlock()
return &NoStatConn{rawConn: rawConn}
}
func (this *NoStatConn) Read(b []byte) (n int, err error) {
return this.rawConn.Read(b)
}
func (this *NoStatConn) Write(b []byte) (n int, err error) {
return this.rawConn.Write(b)
}
func (this *NoStatConn) Close() error {
err := this.rawConn.Close()
noStatLocker.Lock()
delete(noStatAddrMap, this.rawConn.LocalAddr().String())
noStatLocker.Unlock()
return err
}
func (this *NoStatConn) LocalAddr() net.Addr {
return this.rawConn.LocalAddr()
}
func (this *NoStatConn) RemoteAddr() net.Addr {
return this.rawConn.RemoteAddr()
}
func (this *NoStatConn) SetDeadline(t time.Time) error {
return this.rawConn.SetDeadline(t)
}
func (this *NoStatConn) SetReadDeadline(t time.Time) error {
return this.rawConn.SetReadDeadline(t)
}
func (this *NoStatConn) SetWriteDeadline(t time.Time) error {
return this.rawConn.SetWriteDeadline(t)
}