Report traffic for all download tokens

internal/handlers/home.go

@@ -43,7 +43,7 @@ type RootResponse struct {
 	DownloadTokens      []string          `json:"download_tokens"`
 	IDsToDelete         []string          `json:"ids_to_delete"`
 	Hosts               []string          `json:"hosts"`
-	TrafficServedPerAPI uint64            `json:"traffic_served_per_api"`
+	TrafficServedPerAPI map[string]uint64 `json:"traffic_served_per_api"`
 }
 
 func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request) (*RootResponse, error) {
@@ -77,15 +77,16 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request
 	sortedIDs := zr.torMgr.OnceDoneBin.ToSlice()
 	sort.Strings(sortedIDs)
 
-	trafficDetails, err := zr.rd.GetTrafficDetails()
-	if err != nil {
-		http.Error(resp, err.Error(), http.StatusInternalServerError)
-		return nil, err
-	}
-	var trafficFromAPI int64
-	trafficFromAPI = 0
-	if _, ok := trafficDetails["real-debrid.com"]; ok {
-		trafficFromAPI = trafficDetails["real-debrid.com"]
+	trafficFromAPI := make(map[string]uint64)
+	for _, token := range zr.rd.TokenManager.GetAllTokens() {
+		trafficDetails, err := zr.rd.GetTrafficDetails(token)
+		if err != nil {
+			trafficDetails = make(map[string]uint64)
+		}
+		trafficFromAPI[token] = 0
+		if _, ok := trafficDetails["real-debrid.com"]; ok {
+			trafficFromAPI[token] = trafficDetails["real-debrid.com"]
+		}
 	}
 
 	userInfo.Premium = userInfo.Premium / 86400
@@ -105,7 +106,7 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request
 		Infuse:              fmt.Sprintf("//%s/infuse/", req.Host),
 		Logs:                fmt.Sprintf("//%s/logs/", req.Host),
 		UserInfo:            userInfo,
-		TrafficServedPerAPI: uint64(trafficFromAPI),
+		TrafficServedPerAPI: trafficFromAPI,
 		LibrarySize:         allTorrents.Count(),
 		TorrentsToRepair:    repairQueueStr,
 		MemAlloc:            bToMb(mem.Alloc),
@@ -220,10 +221,6 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
             <td>Process ID</td>
             <td colspan="2">%d</td>
         </tr>
-        <tr>
-            <td>Traffic Served (main token)</td>
-            <td colspan="2">%d MB (%d MB since startup)</td>
-        </tr>
         <tr>
             <td>Traffic Served (zurg)</td>
             <td colspan="2">%d MB</td>
@@ -234,11 +231,22 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
 		response.Sys,
 		response.NumGC,
 		response.PID,
-		bToMb(response.TrafficServedPerAPI), // traffic served *api*
-		bToMb(response.TrafficServedPerAPI-zr.downloader.TrafficOnStartup.Load()), // traffic served *api* since startup
-		bToMb(zr.downloader.TrafficServed.Load()),                                 // traffic served *zurg*
+		bToMb(zr.downloader.TrafficServed.Load()), // traffic served *zurg*
 	)
 
+	for token, traffic := range response.TrafficServedPerAPI {
+		trafficOnStartup, _ := zr.downloader.TrafficOnStartup.Get(token)
+		out += fmt.Sprintf(`
+        <tr>
+            <td>Traffic Served (%s)</td>
+            <td colspan="2">%d MB (%d MB since startup)</td>
+        </tr>`,
+			utils.MaskToken(token)[40:],
+			bToMb(traffic),                         // traffic served *api*
+			bToMb(traffic-trafficOnStartup.Load()), // traffic served *api* since startup
+		)
+	}
+
 	out += fmt.Sprintf(`
         <tr>
             <td rowspan="3">Sponsor Zurg</td>

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