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

pkg/http/client.go

@@ -129,7 +129,7 @@ func NewHTTPClient(
 }
 
 func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
-	if r.token != "" {
+	if r.token != "" && req.Header.Get("Authorization") == "" {
 		req.Header.Set("Authorization", "Bearer "+r.token)
 	}
 

pkg/realdebrid/api.go

@@ -338,12 +338,12 @@ func (rd *RealDebrid) GetUserInformation() (*User, error) {
 
 // TrafficDetails represents the structure of the traffic details response
 type TrafficDetails map[string]struct {
-	Host  map[string]int64 `json:"host"`
-	Bytes int64            `json:"bytes"`
+	Host  map[string]uint64 `json:"host"`
+	Bytes int64             `json:"bytes"`
 }
 
 // GetTrafficDetails gets the traffic details from the Real-Debrid API
-func (rd *RealDebrid) GetTrafficDetails() (map[string]int64, error) {
+func (rd *RealDebrid) GetTrafficDetails(token string) (map[string]uint64, error) {
 	// Construct request URL
 	reqURL := "https://api.real-debrid.com/rest/1.0/traffic/details"
 	req, err := http.NewRequest(http.MethodGet, reqURL, nil)
@@ -352,6 +352,8 @@ func (rd *RealDebrid) GetTrafficDetails() (map[string]int64, error) {
 		return nil, err
 	}
 
+	req.Header.Set("Authorization", "Bearer "+token)
+
 	// Send the request
 	resp, err := rd.apiClient.Do(req)
 	if err != nil {

pkg/realdebrid/token_manager.go

@@ -102,3 +102,11 @@ func (dtm *DownloadTokenManager) GetExpiredTokens() []string {
 	}
 	return tokens
 }
+
+func (dtm *DownloadTokenManager) GetAllTokens() []string {
+	var tokens []string
+	for _, t := range dtm.tokens {
+		tokens = append(tokens, t.value)
+	}
+	return tokens
+}