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

tunnel/manager.go

@@ -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)
 	}
 }