Fix: tunnel manager & tracker race condition (#1048)

This commit is contained in:
Jason Lyu
2020-10-29 17:51:14 +08:00
committed by GitHub
parent b98e9ea202
commit 87e4d94290
7 changed files with 82 additions and 68 deletions

View File

@ -2,26 +2,34 @@ package tunnel
import (
"sync"
"sync/atomic"
"time"
"go.uber.org/atomic"
)
var DefaultManager *Manager
func init() {
DefaultManager = &Manager{}
DefaultManager = &Manager{
uploadTemp: atomic.NewInt64(0),
downloadTemp: atomic.NewInt64(0),
uploadBlip: atomic.NewInt64(0),
downloadBlip: atomic.NewInt64(0),
uploadTotal: atomic.NewInt64(0),
downloadTotal: atomic.NewInt64(0),
}
go DefaultManager.handle()
}
type Manager struct {
connections sync.Map
uploadTemp int64
downloadTemp int64
uploadBlip int64
downloadBlip int64
uploadTotal int64
downloadTotal int64
uploadTemp *atomic.Int64
downloadTemp *atomic.Int64
uploadBlip *atomic.Int64
downloadBlip *atomic.Int64
uploadTotal *atomic.Int64
downloadTotal *atomic.Int64
}
func (m *Manager) Join(c tracker) {
@ -33,17 +41,17 @@ func (m *Manager) Leave(c tracker) {
}
func (m *Manager) PushUploaded(size int64) {
atomic.AddInt64(&m.uploadTemp, size)
atomic.AddInt64(&m.uploadTotal, size)
m.uploadTemp.Add(size)
m.uploadTotal.Add(size)
}
func (m *Manager) PushDownloaded(size int64) {
atomic.AddInt64(&m.downloadTemp, size)
atomic.AddInt64(&m.downloadTotal, size)
m.downloadTemp.Add(size)
m.downloadTotal.Add(size)
}
func (m *Manager) Now() (up int64, down int64) {
return atomic.LoadInt64(&m.uploadBlip), atomic.LoadInt64(&m.downloadBlip)
return m.uploadBlip.Load(), m.downloadBlip.Load()
}
func (m *Manager) Snapshot() *Snapshot {
@ -54,29 +62,29 @@ func (m *Manager) Snapshot() *Snapshot {
})
return &Snapshot{
UploadTotal: atomic.LoadInt64(&m.uploadTotal),
DownloadTotal: atomic.LoadInt64(&m.downloadTotal),
UploadTotal: m.uploadTotal.Load(),
DownloadTotal: m.downloadTotal.Load(),
Connections: connections,
}
}
func (m *Manager) ResetStatistic() {
m.uploadTemp = 0
m.uploadBlip = 0
m.uploadTotal = 0
m.downloadTemp = 0
m.downloadBlip = 0
m.downloadTotal = 0
m.uploadTemp.Store(0)
m.uploadBlip.Store(0)
m.uploadTotal.Store(0)
m.downloadTemp.Store(0)
m.downloadBlip.Store(0)
m.downloadTotal.Store(0)
}
func (m *Manager) handle() {
ticker := time.NewTicker(time.Second)
for range ticker.C {
atomic.StoreInt64(&m.uploadBlip, atomic.LoadInt64(&m.uploadTemp))
atomic.StoreInt64(&m.uploadTemp, 0)
atomic.StoreInt64(&m.downloadBlip, atomic.LoadInt64(&m.downloadTemp))
atomic.StoreInt64(&m.downloadTemp, 0)
m.uploadBlip.Store(m.uploadTemp.Load())
m.uploadTemp.Store(0)
m.downloadBlip.Store(m.downloadTemp.Load())
m.downloadTemp.Store(0)
}
}

View File

@ -2,11 +2,12 @@ package tunnel
import (
"net"
"sync/atomic"
"time"
C "github.com/Dreamacro/clash/constant"
"github.com/gofrs/uuid"
"go.uber.org/atomic"
)
type tracker interface {
@ -15,14 +16,14 @@ type tracker interface {
}
type trackerInfo struct {
UUID uuid.UUID `json:"id"`
Metadata *C.Metadata `json:"metadata"`
UploadTotal int64 `json:"upload"`
DownloadTotal int64 `json:"download"`
Start time.Time `json:"start"`
Chain C.Chain `json:"chains"`
Rule string `json:"rule"`
RulePayload string `json:"rulePayload"`
UUID uuid.UUID `json:"id"`
Metadata *C.Metadata `json:"metadata"`
UploadTotal *atomic.Int64 `json:"upload"`
DownloadTotal *atomic.Int64 `json:"download"`
Start time.Time `json:"start"`
Chain C.Chain `json:"chains"`
Rule string `json:"rule"`
RulePayload string `json:"rulePayload"`
}
type tcpTracker struct {
@ -39,7 +40,7 @@ func (tt *tcpTracker) Read(b []byte) (int, error) {
n, err := tt.Conn.Read(b)
download := int64(n)
tt.manager.PushDownloaded(download)
atomic.AddInt64(&tt.DownloadTotal, download)
tt.DownloadTotal.Add(download)
return n, err
}
@ -47,7 +48,7 @@ func (tt *tcpTracker) Write(b []byte) (int, error) {
n, err := tt.Conn.Write(b)
upload := int64(n)
tt.manager.PushUploaded(upload)
atomic.AddInt64(&tt.UploadTotal, upload)
tt.UploadTotal.Add(upload)
return n, err
}
@ -63,11 +64,13 @@ func newTCPTracker(conn C.Conn, manager *Manager, metadata *C.Metadata, rule C.R
Conn: conn,
manager: manager,
trackerInfo: &trackerInfo{
UUID: uuid,
Start: time.Now(),
Metadata: metadata,
Chain: conn.Chains(),
Rule: "",
UUID: uuid,
Start: time.Now(),
Metadata: metadata,
Chain: conn.Chains(),
Rule: "",
UploadTotal: atomic.NewInt64(0),
DownloadTotal: atomic.NewInt64(0),
},
}
@ -94,7 +97,7 @@ func (ut *udpTracker) ReadFrom(b []byte) (int, net.Addr, error) {
n, addr, err := ut.PacketConn.ReadFrom(b)
download := int64(n)
ut.manager.PushDownloaded(download)
atomic.AddInt64(&ut.DownloadTotal, download)
ut.DownloadTotal.Add(download)
return n, addr, err
}
@ -102,7 +105,7 @@ func (ut *udpTracker) WriteTo(b []byte, addr net.Addr) (int, error) {
n, err := ut.PacketConn.WriteTo(b, addr)
upload := int64(n)
ut.manager.PushUploaded(upload)
atomic.AddInt64(&ut.UploadTotal, upload)
ut.UploadTotal.Add(upload)
return n, err
}
@ -118,11 +121,13 @@ func newUDPTracker(conn C.PacketConn, manager *Manager, metadata *C.Metadata, ru
PacketConn: conn,
manager: manager,
trackerInfo: &trackerInfo{
UUID: uuid,
Start: time.Now(),
Metadata: metadata,
Chain: conn.Chains(),
Rule: "",
UUID: uuid,
Start: time.Now(),
Metadata: metadata,
Chain: conn.Chains(),
Rule: "",
UploadTotal: atomic.NewInt64(0),
DownloadTotal: atomic.NewInt64(0),
},
}