guoswork
/
frp
mirror of https://gitee.com/IrisVega/frp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

statsConn: avoid repetition of close function

This commit is contained in:
fatedier 2018-02-01 11:15:35 +08:00
parent 5c9909aeef
commit 1d0865ca49
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
utils/net/conn.go
View File

@ -21,6 +21,7 @@ import (
"io"
"net"
"sync"
"sync/atomic"
"time"
"github.com/fatedier/frp/utils/log"
@ -178,6 +179,7 @@ func (sc *SharedConn) WriteBuff(buffer []byte) (err error) {
type StatsConn struct {
Conn
closed int64 // 1 means closed
totalRead int64
totalWrite int64
statsFunc func(totalRead, totalWrite int64)
@ -203,9 +205,12 @@ func (statsConn *StatsConn) Write(p []byte) (n int, err error) {
}
func (statsConn *StatsConn) Close() (err error) {
old := atomic.SwapInt64(&statsConn.closed, 1)
if old != 1 {
err = statsConn.Conn.Close()
if statsConn.statsFunc != nil {
statsConn.statsFunc(statsConn.totalRead, statsConn.totalWrite)
}
}
return
}