Add bandwidth tracking

internal/universal/downloader.go

@@ -2,26 +2,55 @@ package universal
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"path/filepath"
+	"strconv"
 	"strings"
+	"sync/atomic"
+	"time"
 
 	"github.com/debridmediamanager/zurg/internal/config"
 	intTor "github.com/debridmediamanager/zurg/internal/torrent"
 	zurghttp "github.com/debridmediamanager/zurg/pkg/http"
 	"github.com/debridmediamanager/zurg/pkg/logutil"
 	"github.com/debridmediamanager/zurg/pkg/realdebrid"
+	"github.com/panjf2000/ants/v2"
 )
 
 type Downloader struct {
-	client *zurghttp.HTTPClient
+	client         *zurghttp.HTTPClient
+	RequestedBytes atomic.Uint64
+	TotalBytes     atomic.Uint64
 }
 
-func NewDownloader(client *zurghttp.HTTPClient) *Downloader {
-	return &Downloader{
+func NewDownloader(client *zurghttp.HTTPClient, workerPool *ants.Pool) *Downloader {
+	dl := &Downloader{
 		client: client,
 	}
+
+	// track bandwidth usage and reset at 12AM CET
+	now := time.Now()
+	tomorrow := now.AddDate(0, 0, 1)
+	cetTZ, err := time.LoadLocation("CET")
+	if err != nil {
+		cetTZ = time.FixedZone("CET", 1*60*60)
+	}
+	nextMidnightInCET := time.Date(tomorrow.Year(), tomorrow.Month(), tomorrow.Day(), 0, 0, 0, 0, cetTZ)
+	duration := nextMidnightInCET.Sub(now)
+	timer := time.NewTimer(duration)
+	workerPool.Submit(func() {
+		<-timer.C
+		ticker := time.NewTicker(24 * time.Hour)
+		for {
+			dl.RequestedBytes.Store(0)
+			dl.TotalBytes.Store(0)
+			<-ticker.C
+		}
+	})
+
+	return dl
 }
 
 // DownloadFile handles a GET request for files in torrents
@@ -146,7 +175,7 @@ func (dl *Downloader) streamFileToResponse(
 	// Add the range header if it exists
 	if req.Header.Get("Range") != "" {
 		dlReq.Header.Add("Range", req.Header.Get("Range"))
-		// log.Debugf("Range request for file %s: %s", unrestrict.Download, req.Header.Get("Range"))
+		// log.Debugf("Serving file %s: %s", unrestrict.Filename, req.Header.Get("Range"))
 	}
 
 	// Perform the request
@@ -182,9 +211,66 @@ func (dl *Downloader) streamFileToResponse(
 	}
 
 	buf := make([]byte, cfg.GetNetworkBufferSize())
-	io.CopyBuffer(resp, downloadResp.Body, buf)
+	n, _ := io.CopyBuffer(resp, downloadResp.Body, buf)
+
+	// Update the download statistics
+	reqBytes, _ := parseRangeHeader(req.Header.Get("Range"))
+	if reqBytes == 0 && unrestrict != nil {
+		reqBytes = uint64(unrestrict.Filesize)
+	}
+	dl.RequestedBytes.Add(reqBytes)
+	dl.TotalBytes.Add(uint64(n))
+	log.Debugf("Served %d MB of the requested %d MB of file %s (range=%s)", bToMb(uint64(n)), bToMb(reqBytes), unrestrict.Filename, req.Header.Get("Range"))
 }
 
 func redirect(resp http.ResponseWriter, req *http.Request, url string) {
 	http.Redirect(resp, req, url, http.StatusFound)
 }
+
+func parseRangeHeader(rangeHeader string) (uint64, error) {
+	if rangeHeader == "" { // Empty header means no range request
+		return 0, nil
+	}
+
+	if !strings.HasPrefix(rangeHeader, "bytes=") {
+		return 0, fmt.Errorf("invalid range header format")
+	}
+
+	parts := strings.SplitN(rangeHeader[6:], "-", 2) // [6:] removes "bytes="
+	if len(parts) != 2 {
+		return 0, fmt.Errorf("invalid range specification")
+	}
+
+	var start, end uint64
+	var err error
+
+	// Case 1: "bytes=100-" (from byte 100 to the end)
+	if parts[0] != "" {
+		start, err = strconv.ParseUint(parts[0], 10, 64)
+		if err != nil {
+			return 0, fmt.Errorf("invalid start value: %w", err)
+		}
+	}
+
+	// Case 2: "bytes=-200" (last 200 bytes)
+	if parts[1] != "" {
+		end, err = strconv.ParseUint(parts[1], 10, 64)
+		if err != nil {
+			return 0, fmt.Errorf("invalid end value: %w", err)
+		}
+	}
+
+	// Handle "bytes=500-100" (invalid range)
+	if start > end {
+		return 0, fmt.Errorf("invalid range: start cannot be greater than end")
+	}
+
+	// Calculate bytes to read
+	bytesToRead := end - start + 1 // +1 because ranges are inclusive
+
+	return bytesToRead, nil
+}
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}