hpz937/zurg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add bandwidth tracking

Browse Source
This commit is contained in:
Ben Adrian Sarmiento
2024-06-24 23:04:31 +02:00
parent ec25c1ef17
commit e9e31759c1
3 changed files with 131 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/app.go
View File

@@ -102,7 +102,11 @@ func MainApp(configPath string) {
log.Named("download_client"),
)
workerPool, err := ants.NewPool(config.GetNumberOfWorkers())
workerCount := config.GetNumberOfWorkers()
if workerCount < 10 {
workerCount = 10
}
workerPool, err := ants.NewPool(workerCount)
if err != nil {
zurglog.Errorf("Failed to create worker pool: %v", err)
os.Exit(1)
@@ -140,7 +144,7 @@ func MainApp(configPath string) {
log.Named("repair"),
)
downloader := universal.NewDownloader(downloadClient)
downloader := universal.NewDownloader(downloadClient, workerPool)
router := chi.NewRouter()
handlers.AttachHandlers(

46
internal/handlers/home.go
View File

@@ -29,7 +29,9 @@ type RootResponse struct {
Infuse string `json:"infuse"`
Logs string `json:"logs"`
UserInfo *realdebrid.User `json:"user_info"`
RDTrafficUsed uint64 `json:"rd_traffic_used"`
TrafficLogged uint64 `json:"traffic_logged"`
RequestedMB uint64 `json:"requested_mb"`
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
@@ -82,10 +84,10 @@ func (zr *Handlers) generateResponse(resp http.ResponseWriter, req *http.Request
sort.Strings(sortedIDs)
// check if real-debrid.com is in the traffic details
var rdTrafficUsed int64
rdTrafficUsed = 0
var trafficLogged int64
trafficLogged = 0
if _, ok := trafficDetails["real-debrid.com"]; ok {
rdTrafficUsed = trafficDetails["real-debrid.com"]
trafficLogged = trafficDetails["real-debrid.com"]
}
userInfo.Premium = userInfo.Premium / 86400
@@ -101,7 +103,9 @@ 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,
RDTrafficUsed: bToGb(uint64(rdTrafficUsed)),
TrafficLogged: bToMb(uint64(trafficLogged)),
RequestedMB: bToMb(zr.downloader.RequestedBytes.Load()),
ServedMB: bToMb(zr.downloader.TotalBytes.Load()),
LibrarySize: allTorrents.Count(),
TorrentsToRepair: repairQueueStr,
MemAlloc: bToMb(mem.Alloc),
@@ -190,6 +194,12 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
response.Logs,
)
denominator := response.RequestedMB
if denominator == 0 {
denominator = 1
}
efficiency := response.ServedMB / denominator
out += fmt.Sprintf(`
<tr>
<td>Library Size</td>
@@ -216,8 +226,20 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
<td colspan="2">%d</td>
</tr>
<tr>
<td>RD Traffic Used</td>
<td colspan="2">%d GB</td>
<td>Traffic Logged</td>
<td colspan="2">%d MB</td>
</tr>
<tr>
<td>Traffic Requested</td>
<td colspan="2">%d MB</td>
</tr>
<tr>
<td>Traffic Served</td>
<td colspan="2">%d MB</td>
</tr>
<tr>
<td>Traffic Efficiency</td>
<td colspan="2">%d%% (wasted %d MB)</td>
</tr>`,
response.LibrarySize,
response.MemAlloc,
@@ -225,7 +247,11 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
response.Sys,
response.NumGC,
response.PID,
response.RDTrafficUsed,
response.TrafficLogged,
response.RequestedMB,
response.ServedMB,
efficiency,
response.RequestedMB-response.ServedMB,
)
out += fmt.Sprintf(`
@@ -467,7 +493,3 @@ func (zr *Handlers) handleHome(resp http.ResponseWriter, req *http.Request) {
func bToMb(b uint64) uint64 {
return b / 1024 / 1024
}
func bToGb(b uint64) uint64 {
return b / 1024 / 1024 / 1024
}

96
internal/universal/downloader.go
View File

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