Report traffic for all download tokens

internal/universal/downloader.go

@@ -15,6 +15,7 @@ import (
 	"github.com/debridmediamanager/zurg/pkg/logutil"
 	"github.com/debridmediamanager/zurg/pkg/realdebrid"
 	"github.com/debridmediamanager/zurg/pkg/utils"
+	cmap "github.com/orcaman/concurrent-map/v2"
 	"github.com/panjf2000/ants/v2"
 )
 
@@ -22,22 +23,27 @@ type Downloader struct {
 	rd               *realdebrid.RealDebrid
 	workerPool       *ants.Pool
 	TrafficServed    atomic.Uint64
-	TrafficOnStartup atomic.Uint64
+	TrafficOnStartup cmap.ConcurrentMap[string, *atomic.Uint64]
 }
 
 func NewDownloader(rd *realdebrid.RealDebrid, workerPool *ants.Pool) *Downloader {
 	dl := &Downloader{
-		rd:         rd,
-		workerPool: workerPool,
+		rd:               rd,
+		workerPool:       workerPool,
+		TrafficOnStartup: cmap.New[*atomic.Uint64](),
 	}
 
-	trafficDetails, err := dl.rd.GetTrafficDetails()
-	if err != nil {
-		trafficDetails = make(map[string]int64)
-	}
-	dl.TrafficOnStartup.Store(uint64(0))
-	if _, ok := trafficDetails["real-debrid.com"]; ok {
-		dl.TrafficOnStartup.Store(uint64(trafficDetails["real-debrid.com"]))
+	for _, token := range rd.TokenManager.GetAllTokens() {
+		trafficDetails, err := dl.rd.GetTrafficDetails(token)
+		if err != nil {
+			trafficDetails = make(map[string]uint64)
+		}
+		var trafficOnStartup atomic.Uint64
+		trafficOnStartup.Store(0)
+		if _, ok := trafficDetails["real-debrid.com"]; ok {
+			trafficOnStartup.Store(uint64(trafficDetails["real-debrid.com"]))
+		}
+		dl.TrafficOnStartup.Set(token, &trafficOnStartup)
 	}
 
 	return dl
@@ -62,7 +68,9 @@ func (dl *Downloader) StartResetBandwidthCountersJob() {
 		for {
 			dl.rd.TokenManager.ResetAllTokens()
 			dl.TrafficServed.Store(0)
-			dl.TrafficOnStartup.Store(0)
+			dl.TrafficOnStartup.IterCb(func(token string, traffic *atomic.Uint64) {
+				traffic.Store(0)
+			})
 			<-ticker.C
 		}
 	})