Explorar o código

statsConn: avoid repetition of close function

fatedier %!s(int64=7) %!d(string=hai) anos
pai
achega
1d0865ca49
Modificáronse 1 ficheiros con 8 adicións e 3 borrados
  1. 8 3
      utils/net/conn.go

+ 8 - 3
utils/net/conn.go

@@ -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) {
-	err = statsConn.Conn.Close()
-	if statsConn.statsFunc != nil {
-		statsConn.statsFunc(statsConn.totalRead, statsConn.totalWrite)
+	old := atomic.SwapInt64(&statsConn.closed, 1)
+	if old != 1 {
+		err = statsConn.Conn.Close()
+		if statsConn.statsFunc != nil {
+			statsConn.statsFunc(statsConn.totalRead, statsConn.totalWrite)
+		}
 	}
 	return
 }