Fix traffic computation and optimize open ended requests

2024-07-04 03:39:53 +02:00
parent 6deeb45db0
commit 49432dd810
5 changed files with 100 additions and 106 deletions
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -5,7 +5,7 @@ on:
    tags:
      - 'v**'
  schedule:
-    - cron: '50 1 * * *'
+    - cron: '45 1 * * *'

 jobs:
  build:
--- a/internal/handlers/home.go
+++ b/internal/handlers/home.go
@@ -22,29 +22,28 @@ type SponsorResponse struct {
 }

 type RootResponse struct {
-	Version          string            `json:"version"`
-	BuiltAt          string            `json:"built_at"`
-	GitCommit        string            `json:"git_commit"`
-	Html             string            `json:"html"`
-	Dav              string            `json:"dav"`
-	Infuse           string            `json:"infuse"`
-	Logs             string            `json:"logs"`
-	UserInfo         *realdebrid.User  `json:"user_info"`
-	APITraffic       uint64            `json:"traffic_from_api"`
-	ServedMB         uint64            `json:"served_mb"`
-	LibrarySize      int               `json:"library_size"` // Number of torrents in the library
-	TorrentsToRepair string            `json:"repair_queue"` // List of torrents in the repair queue
-	MemAlloc         uint64            `json:"mem_alloc"`    // Memory allocation in MB
-	TotalAlloc       uint64            `json:"total_alloc"`  // Total memory allocated in MB
-	Sys              uint64            `json:"sys"`          // System memory in MB
-	NumGC            uint32            `json:"num_gc"`       // Number of completed GC cycles
-	PID              int               `json:"pid"`          // Process ID
-	Sponsor          SponsorResponse   `json:"sponsor_zurg"` // Sponsorship links
-	Config           config.ZurgConfig `json:"config"`
-	Token            string            `json:"token"`
-	DownloadTokens   []string          `json:"download_tokens"`
-	IDsToDelete      []string          `json:"ids_to_delete"`
-	Hosts            []string          `json:"hosts"`
+	Version             string            `json:"version"`
+	BuiltAt             string            `json:"built_at"`
+	GitCommit           string            `json:"git_commit"`
+	Html                string            `json:"html"`
+	Dav                 string            `json:"dav"`
+	Infuse              string            `json:"infuse"`
+	Logs                string            `json:"logs"`
+	UserInfo            *realdebrid.User  `json:"user_info"`
+	LibrarySize         int               `json:"library_size"` // Number of torrents in the library
+	TorrentsToRepair    string            `json:"repair_queue"` // List of torrents in the repair queue
+	MemAlloc            uint64            `json:"mem_alloc"`    // Memory allocation in MB
+	TotalAlloc          uint64            `json:"total_alloc"`  // Total memory allocated in MB
+	Sys                 uint64            `json:"sys"`          // System memory in MB
+	NumGC               uint32            `json:"num_gc"`       // Number of completed GC cycles
+	PID                 int               `json:"pid"`          // Process ID
+	Sponsor             SponsorResponse   `json:"sponsor_zurg"` // Sponsorship links
+	Config              config.ZurgConfig `json:"config"`
+	Token               string            `json:"token"`
+	DownloadTokens      []string          `json:"download_tokens"`
+	IDsToDelete         []string          `json:"ids_to_delete"`
+	Hosts               []string          `json:"hosts"`
+	TrafficServedPerAPI uint64            `json:"traffic_served_per_api"`
 }

 func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request) (*RootResponse, error) {
@@ -98,23 +97,22 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request
 	}

 	return &RootResponse{
-		Version:          version.GetVersion(),
-		BuiltAt:          version.GetBuiltAt(),
-		GitCommit:        version.GetGitCommit(),
-		Html:             fmt.Sprintf("//%s/http/", req.Host),
-		Dav:              fmt.Sprintf("//%s/dav/", req.Host),
-		Infuse:           fmt.Sprintf("//%s/infuse/", req.Host),
-		Logs:             fmt.Sprintf("//%s/logs/", req.Host),
-		UserInfo:         userInfo,
-		APITraffic:       uint64(trafficFromAPI),
-		ServedMB:         bToMb(zr.downloader.TotalBytes.Load()),
-		LibrarySize:      allTorrents.Count(),
-		TorrentsToRepair: repairQueueStr,
-		MemAlloc:         bToMb(mem.Alloc),
-		TotalAlloc:       bToMb(mem.TotalAlloc),
-		Sys:              bToMb(mem.Sys),
-		NumGC:            mem.NumGC,
-		PID:              os.Getpid(),
+		Version:             version.GetVersion(),
+		BuiltAt:             version.GetBuiltAt(),
+		GitCommit:           version.GetGitCommit(),
+		Html:                fmt.Sprintf("//%s/http/", req.Host),
+		Dav:                 fmt.Sprintf("//%s/dav/", req.Host),
+		Infuse:              fmt.Sprintf("//%s/infuse/", req.Host),
+		Logs:                fmt.Sprintf("//%s/logs/", req.Host),
+		UserInfo:            userInfo,
+		TrafficServedPerAPI: uint64(trafficFromAPI),
+		LibrarySize:         allTorrents.Count(),
+		TorrentsToRepair:    repairQueueStr,
+		MemAlloc:            bToMb(mem.Alloc),
+		TotalAlloc:          bToMb(mem.TotalAlloc),
+		Sys:                 bToMb(mem.Sys),
+		NumGC:               mem.NumGC,
+		PID:                 os.Getpid(),
 		Sponsor: SponsorResponse{
 			Patreon: "https://www.patreon.com/debridmediamanager",
 			Github:  "https://github.com/sponsors/debridmediamanager",
@@ -197,18 +195,6 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
 		response.Logs,
 	)

-	denominator := bToMb(response.APITraffic - zr.trafficOnStartup.Load())
-	if denominator == 0 {
-		denominator = 1
-	}
-	efficiency := response.ServedMB * 100 / denominator
-
-	if zr.trafficOnStartup.Load() > response.APITraffic {
-		// it cannot be bigger than traffic logged
-		// so it must be a reset back to 0
-		zr.trafficOnStartup.Store(0)
-	}
-
 	out += fmt.Sprintf(`
        <tr>
            <td>Library Size</td>
@@ -235,16 +221,12 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
            <td colspan="2">%d</td>
        </tr>
        <tr>
-            <td>Traffic Logged</td>
-            <td colspan="2">%d MB (%d MB added)</td>
+            <td>Traffic Served (main token)</td>
+            <td colspan="2">%d MB (%d MB since startup)</td>
        </tr>
        <tr>
-            <td>Traffic Served</td>
+            <td>Traffic Served (zurg)</td>
            <td colspan="2">%d MB</td>
-        </tr>
-        <tr>
-            <td>Traffic Efficiency</td>
-            <td colspan="2">%d%% (wasted %d MB) disclaimer: this assumes all RD traffic comes from this zurg instance</td>
        </tr>`,
 		response.LibrarySize,
 		response.MemAlloc,
@@ -252,11 +234,9 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
 		response.Sys,
 		response.NumGC,
 		response.PID,
-		bToMb(response.APITraffic),
-		bToMb(response.APITraffic-zr.trafficOnStartup.Load()),
-		response.ServedMB,
-		efficiency,
-		bToMb(response.APITraffic-zr.trafficOnStartup.Load())-response.ServedMB,
+		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*
 	)

 	out += fmt.Sprintf(`
--- a/internal/handlers/router.go
+++ b/internal/handlers/router.go
@@ -6,7 +6,6 @@ import (
 	"net/url"
 	"path/filepath"
 	"strings"
-	"sync/atomic"

 	"github.com/debridmediamanager/zurg/internal/config"
 	"github.com/debridmediamanager/zurg/internal/dav"
@@ -21,14 +20,13 @@ import (
 )

 type Handlers struct {
-	downloader       *universal.Downloader
-	torMgr           *torrent.TorrentManager
-	cfg              config.ConfigInterface
-	rd               *realdebrid.RealDebrid
-	workerPool       *ants.Pool
-	hosts            []string
-	trafficOnStartup atomic.Uint64
-	log              *logutil.Logger
+	downloader *universal.Downloader
+	torMgr     *torrent.TorrentManager
+	cfg        config.ConfigInterface
+	rd         *realdebrid.RealDebrid
+	workerPool *ants.Pool
+	hosts      []string
+	log        *logutil.Logger
 }

 func init() {
@@ -48,16 +46,6 @@ func AttachHandlers(router *chi.Mux, downloader *universal.Downloader, torMgr *t
 		log:        log,
 	}

-	trafficDetails, err := rd.GetTrafficDetails()
-	if err != nil {
-		log.Errorf("Failed to get traffic details: %v", err)
-		trafficDetails = make(map[string]int64)
-	}
-	hs.trafficOnStartup.Store(uint64(0))
-	if _, ok := trafficDetails["real-debrid.com"]; ok {
-		hs.trafficOnStartup.Store(uint64(trafficDetails["real-debrid.com"]))
-	}
-
 	if cfg.GetUsername() != "" {
 		router.Use(hs.basicAuth)
 	}
--- a/internal/universal/downloader.go
+++ b/internal/universal/downloader.go
@@ -2,6 +2,7 @@ package universal

 import (
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"path/filepath"
@@ -18,16 +19,28 @@ import (
 )

 type Downloader struct {
-	rd         *realdebrid.RealDebrid
-	workerPool *ants.Pool
-	TotalBytes atomic.Uint64
+	rd               *realdebrid.RealDebrid
+	workerPool       *ants.Pool
+	TrafficServed    atomic.Uint64
+	TrafficOnStartup atomic.Uint64
 }

 func NewDownloader(rd *realdebrid.RealDebrid, workerPool *ants.Pool) *Downloader {
-	return &Downloader{
+	dl := &Downloader{
 		rd:         rd,
 		workerPool: workerPool,
 	}
+
+	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"]))
+	}
+
+	return dl
 }

 // StartResetBandwidthCountersJob is a permanent job that resets the bandwidth counters at 12AM CET
@@ -48,7 +61,8 @@ func (dl *Downloader) StartResetBandwidthCountersJob() {
 		ticker := time.NewTicker(24 * time.Hour)
 		for {
 			dl.rd.TokenManager.ResetAllTokens()
-			dl.TotalBytes.Store(0)
+			dl.TrafficServed.Store(0)
+			dl.TrafficOnStartup.Store(0)
 			<-ticker.C
 		}
 	})
@@ -162,7 +176,12 @@ func (dl *Downloader) streamFileToResponse(

 	// Add the range header if it exists
 	if req.Header.Get("Range") != "" {
-		dlReq.Header.Add("Range", req.Header.Get("Range"))
+		rangeVal := req.Header.Get("Range")
+		// check if open-ended range request (e.g. "bytes=999-")
+		if strings.HasSuffix(rangeVal, "-") {
+			rangeVal += fmt.Sprintf("%d", unrestrict.Filesize-1)
+		}
+		dlReq.Header.Add("Range", rangeVal)
 	}

 	downloadResp, err := dl.rd.DownloadFile(dlReq)
@@ -227,13 +246,14 @@ func (dl *Downloader) streamFileToResponse(
 		n, _ = io.Copy(resp, downloadResp.Body)
 	}

+	if !strings.HasPrefix(unrestrict.Link, "https://real-debrid.com/d/") {
+		return
+	}
+
 	dl.workerPool.Submit(func() {
-		if n == 0 {
-			return
-		}
-		// Update the download statistics
-		dl.TotalBytes.Add(uint64(n))
-		if cfg.ShouldLogRequests() {
+		dl.TrafficServed.Add(uint64(n))
+
+		if cfg.ShouldLogRequests() && bToMb(uint64(n)) > 0 {
 			log.Debugf("Served %d MB of file %s (range=%s)", bToMb(uint64(n)), unrestrict.Filename, req.Header.Get("Range"))
 		}
 	})
--- a/pkg/realdebrid/api.go
+++ b/pkg/realdebrid/api.go
@@ -70,6 +70,10 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
 		tokenMap, _ := rd.UnrestrictMap.Get(token)
 		if tokenMap.Has(link) {
 			download, _ := tokenMap.Get(link)
+			if !rd.cfg.ShouldServeFromRclone() {
+				return download, nil
+			}
+
 			// check if the link is in the verified links cache
 			if expiry, ok := rd.verifiedLinks.Get(download.ID); ok && expiry > time.Now().Unix() {
 				return download, nil
@@ -98,17 +102,19 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {

 		tokenMap.Set(link, download)

-		err = rd.downloadClient.VerifyLink(download.Download)
-		if utils.IsBytesLimitReached(err) {
-			rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
-			continue
-		} else if utils.IsInvalidDownloadCode(err) {
-			continue
-		} else if err != nil {
-			return nil, err
-		}
+		if rd.cfg.ShouldServeFromRclone() {
+			err = rd.downloadClient.VerifyLink(download.Download)
+			if utils.IsBytesLimitReached(err) {
+				rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
+				continue
+			} else if utils.IsInvalidDownloadCode(err) {
+				continue
+			} else if err != nil {
+				return nil, err
+			}

-		rd.verifiedLinks.Set(download.ID, time.Now().Unix()+DOWNLOAD_LINK_EXPIRY)
+			rd.verifiedLinks.Set(download.ID, time.Now().Unix()+DOWNLOAD_LINK_EXPIRY)
+		}

 		return download, err
 	}