github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/conns/conn_no_stat.go (about)

     1  // Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package connutils
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeNode/internal/zero"
     7  	"net"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  // 记录不需要带宽统计的连接
    13  // 比如本地的清理和预热
    14  var noStatAddrMap = map[string]zero.Zero{} // addr => Zero
    15  var noStatLocker = &sync.RWMutex{}
    16  
    17  // IsNoStatConn 检查是否为不统计连接
    18  func IsNoStatConn(addr string) bool {
    19  	noStatLocker.RLock()
    20  	_, ok := noStatAddrMap[addr]
    21  	noStatLocker.RUnlock()
    22  	return ok
    23  }
    24  
    25  type NoStatConn struct {
    26  	rawConn net.Conn
    27  }
    28  
    29  func NewNoStat(rawConn net.Conn) net.Conn {
    30  	noStatLocker.Lock()
    31  	noStatAddrMap[rawConn.LocalAddr().String()] = zero.New()
    32  	noStatLocker.Unlock()
    33  	return &NoStatConn{rawConn: rawConn}
    34  }
    35  
    36  func (this *NoStatConn) Read(b []byte) (n int, err error) {
    37  	return this.rawConn.Read(b)
    38  }
    39  
    40  func (this *NoStatConn) Write(b []byte) (n int, err error) {
    41  	return this.rawConn.Write(b)
    42  }
    43  
    44  func (this *NoStatConn) Close() error {
    45  	err := this.rawConn.Close()
    46  
    47  	noStatLocker.Lock()
    48  	delete(noStatAddrMap, this.rawConn.LocalAddr().String())
    49  	noStatLocker.Unlock()
    50  
    51  	return err
    52  }
    53  
    54  func (this *NoStatConn) LocalAddr() net.Addr {
    55  	return this.rawConn.LocalAddr()
    56  }
    57  
    58  func (this *NoStatConn) RemoteAddr() net.Addr {
    59  	return this.rawConn.RemoteAddr()
    60  }
    61  
    62  func (this *NoStatConn) SetDeadline(t time.Time) error {
    63  	return this.rawConn.SetDeadline(t)
    64  }
    65  
    66  func (this *NoStatConn) SetReadDeadline(t time.Time) error {
    67  	return this.rawConn.SetReadDeadline(t)
    68  }
    69  
    70  func (this *NoStatConn) SetWriteDeadline(t time.Time) error {
    71  	return this.rawConn.SetWriteDeadline(t)
    72  }